package asymm

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"

	"github.com/emmansun/gmsm/sm2"
	"github.com/emmansun/gmsm/smx509"
	"golang.org/x/crypto/ssh"
)

func EncodePrivateKey(private crypto.PrivateKey, secret string) ([]byte, error) {
	switch key := private.(type) {
	case *rsa.PrivateKey:
		return EncodeRsaPrivateKey(key, secret)
	case *ecdsa.PrivateKey:
		return EncodeEcdsaPrivateKey(key, secret)
	case *sm2.PrivateKey:
		return EncodeSM2PrivateKey(key, secret)
	default:
		b, err := x509.MarshalPKCS8PrivateKey(private)
		if err != nil {
			return nil, err
		}
		if secret == "" {
			return pem.EncodeToMemory(&pem.Block{
				Type:  "PRIVATE KEY",
				Bytes: b,
			}), nil
		}
		blk, err := x509.EncryptPEMBlock(rand.Reader, "PRIVATE KEY", b, []byte(secret), x509.PEMCipherAES256)
		if err != nil {
			return nil, err
		}
		return pem.EncodeToMemory(blk), nil
	}
}

func EncodePublicKey(public crypto.PublicKey) ([]byte, error) {
	switch key := public.(type) {
	case *rsa.PublicKey:
		return EncodeRsaPublicKey(key)
	case *ecdsa.PublicKey:
		if sm2.IsSM2PublicKey(key) {
			return EncodeSM2PublicKey(key)
		}
		return EncodeEcdsaPublicKey(key)
	default:
		publicBytes, err := x509.MarshalPKIXPublicKey(public)
		if err != nil {
			return nil, err
		}
		return pem.EncodeToMemory(&pem.Block{
			Type:  "PUBLIC KEY",
			Bytes: publicBytes,
		}), nil
	}
}

func DecodePrivateKey(private []byte, password string) (crypto.PrivateKey, error) {
	blk, _ := pem.Decode(private)
	if blk == nil {
		return nil, errors.New("private key error")
	}

	switch blk.Type {
	case "RSA PRIVATE KEY":
		return DecodeRsaPrivateKey(private, password)
	case "EC PRIVATE KEY":
		return DecodeEcdsaPrivateKey(private, password)
	case "SM2 PRIVATE KEY":
		return DecodeSM2PrivateKey(private, password)
	case "PRIVATE KEY":
		bytes, err := decodePEMBlockBytes(blk, password)
		if err != nil {
			return nil, err
		}
		key, err := x509.ParsePKCS8PrivateKey(bytes)
		if err == nil {
			return key, nil
		}
		return smx509.ParsePKCS8PrivateKey(bytes)
	case "OPENSSH PRIVATE KEY":
		if password == "" {
			return ssh.ParseRawPrivateKey(private)
		}
		return ssh.ParseRawPrivateKeyWithPassphrase(private, []byte(password))
	default:
		return nil, errors.New("private key type error")
	}
}

func EncodeOpenSSHPrivateKey(private crypto.PrivateKey, secret string) ([]byte, error) {
	key := interface{}(private)
	if k, ok := private.(*ed25519.PrivateKey); ok {
		key = *k
	}

	var (
		block *pem.Block
		err   error
	)
	if secret == "" {
		block, err = ssh.MarshalPrivateKey(key, "")
	} else {
		block, err = ssh.MarshalPrivateKeyWithPassphrase(key, "", []byte(secret))
	}
	if err != nil {
		return nil, err
	}
	return pem.EncodeToMemory(block), nil
}

func DecodePublicKey(pubStr []byte) (crypto.PublicKey, error) {
	blk, _ := pem.Decode(pubStr)
	if blk == nil {
		return nil, errors.New("public key error")
	}
	key, err := x509.ParsePKIXPublicKey(blk.Bytes)
	if err == nil {
		return key, nil
	}
	return smx509.ParsePKIXPublicKey(blk.Bytes)
}

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

func DecodeSSHPublicKey(pubStr []byte) (crypto.PublicKey, error) {
	return ssh.ParsePublicKey(pubStr)
}

func decodePEMBlockBytes(blk *pem.Block, password string) ([]byte, error) {
	if password == "" {
		if x509.IsEncryptedPEMBlock(blk) || smx509.IsEncryptedPEMBlock(blk) {
			return nil, errors.New("private key is encrypted but password is empty")
		}
		return blk.Bytes, nil
	}

	if x509.IsEncryptedPEMBlock(blk) {
		if b, err := x509.DecryptPEMBlock(blk, []byte(password)); err == nil {
			return b, nil
		}
	}
	if smx509.IsEncryptedPEMBlock(blk) {
		return smx509.DecryptPEMBlock(blk, []byte(password))
	}
	return blk.Bytes, nil
}