package asymm

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"encoding/pem"
	"errors"

	"golang.org/x/crypto/ssh"
)

func GenerateEcdsaKey(pubkeyCurve elliptic.Curve) (*ecdsa.PrivateKey, *ecdsa.PublicKey, error) {
	priv, err := ecdsa.GenerateKey(pubkeyCurve, rand.Reader)
	if err != nil {
		return nil, nil, err
	}
	return priv, &priv.PublicKey, nil
}

func EncodeEcdsaPrivateKey(private *ecdsa.PrivateKey, secret string) ([]byte, error) {
	b, err := x509.MarshalECPrivateKey(private)
	if err != nil {
		return nil, err
	}
	if secret == "" {
		return pem.EncodeToMemory(&pem.Block{
			Type:  "EC PRIVATE KEY",
			Bytes: b,
		}), nil
	}
	blk, err := x509.EncryptPEMBlock(rand.Reader, "EC PRIVATE KEY", b, []byte(secret), x509.PEMCipherAES256)
	if err != nil {
		return nil, err
	}
	return pem.EncodeToMemory(blk), nil
}

func EncodeEcdsaPublicKey(public *ecdsa.PublicKey) ([]byte, error) {
	publicBytes, err := x509.MarshalPKIXPublicKey(public)
	if err != nil {
		return nil, err
	}
	return pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: publicBytes,
	}), nil
}

func DecodeEcdsaPrivateKey(private []byte, password string) (*ecdsa.PrivateKey, error) {
	blk, _ := pem.Decode(private)
	if blk == nil {
		return nil, errors.New("private key error")
	}
	bytes, err := decodePEMBlockBytes(blk, password)
	if err != nil {
		return nil, err
	}

	if key, err := x509.ParseECPrivateKey(bytes); err == nil {
		return key, nil
	}
	pkcs8, err := x509.ParsePKCS8PrivateKey(bytes)
	if err != nil {
		return nil, err
	}
	key, ok := pkcs8.(*ecdsa.PrivateKey)
	if !ok {
		return nil, errors.New("private key is not ECDSA")
	}
	return key, nil
}

func DecodeEcdsaPublicKey(pubStr []byte) (*ecdsa.PublicKey, error) {
	blk, _ := pem.Decode(pubStr)
	if blk == nil {
		return nil, errors.New("public key error")
	}
	pub, err := x509.ParsePKIXPublicKey(blk.Bytes)
	if err != nil {
		return nil, err
	}
	key, ok := pub.(*ecdsa.PublicKey)
	if !ok {
		return nil, errors.New("public key is not ECDSA")
	}
	return key, nil
}

func EncodeEcdsaSSHPublicKey(public *ecdsa.PublicKey) ([]byte, error) {
	publicKey, err := ssh.NewPublicKey(public)
	if err != nil {
		return nil, err
	}
	return ssh.MarshalAuthorizedKey(publicKey), nil
}

func GenerateEcdsaSSHKeyPair(pubkeyCurve elliptic.Curve, secret string) (string, string, error) {
	pkey, pubkey, err := GenerateEcdsaKey(pubkeyCurve)
	if err != nil {
		return "", "", err
	}
	pub, err := EncodeEcdsaSSHPublicKey(pubkey)
	if err != nil {
		return "", "", err
	}
	priv, err := EncodeEcdsaPrivateKey(pkey, secret)
	if err != nil {
		return "", "", err
	}
	return string(priv), string(pub), nil
}