package encrypt

import (
	"crypto/ecdsa"
	"encoding/base64"
	"encoding/hex"
	"encoding/pem"
	"fmt"
	"github.com/emmansun/gmsm/sm2"
	"github.com/emmansun/gmsm/smx509"
)

type KeyType string

const (
	Raw    KeyType = "raw"
	Hex            = "hex"
	Base64         = "base64"
	Pem            = "pem"
)

const publicPemFormat = `-----BEGIN PUBLIC KEY-----
%s
-----END PUBLIC KEY-----`

const privatePemFormat = `-----BEGIN PRIVATE KEY-----
%s
-----END PRIVATE KEY-----`

func ParsePublicKey(key string, kt KeyType) (*ecdsa.PublicKey, error) {
	if len(key) == 0 {
		return nil, nil
	}

	switch kt {
	case Raw:
		return sm2.NewPublicKey([]byte(key))
	case Hex:
		pubBytes, err := hex.DecodeString(key)
		if err != nil {
			return nil, err
		}
		return sm2.NewPublicKey(pubBytes)
	case Base64:
		pubBytes, err := base64.StdEncoding.DecodeString(key)
		if err != nil {
			return nil, err
		}
		return sm2.NewPublicKey(pubBytes)
	case Pem:
		pubPEM := fmt.Sprintf(publicPemFormat, key)
		block, _ := pem.Decode([]byte(pubPEM))
		if block == nil {
			return nil, fmt.Errorf("failed to parse PEM block containing the public key")
		}

		pub, err := smx509.ParsePKIXPublicKey(block.Bytes)
		if err != nil {
			return nil, err
		}

		if !sm2.IsSM2PublicKey(pub) {
			return nil, fmt.Errorf("not a SM2 public key")
		}

		if key, ok := pub.(*ecdsa.PublicKey); ok {
			return key, nil
		}

		return nil, fmt.Errorf("not a valid SM2 public key")
	default:
		return nil, nil
	}
}

func ParsePrivateKey(key string, kt KeyType) (*sm2.PrivateKey, error) {
	if len(key) == 0 {
		return nil, nil
	}

	switch kt {
	case Raw:
		return sm2.NewPrivateKey([]byte(key))
	case Hex:
		priBytes, err := hex.DecodeString(key)
		if err != nil {
			return nil, err
		}
		return sm2.NewPrivateKey(priBytes)
	case Base64:
		priBytes, err := base64.StdEncoding.DecodeString(key)
		if err != nil {
			return nil, err
		}
		return sm2.NewPrivateKey(priBytes)
	case Pem:
		priPEM := fmt.Sprintf(privatePemFormat, key)
		block, _ := pem.Decode([]byte(priPEM))
		if block == nil {
			return nil, fmt.Errorf("failed to parse PEM block containing the private key")
		}

		pri, err := smx509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return nil, err
		}

		if pk, ok := pri.(*sm2.PrivateKey); ok {
			return pk, nil
		}
		return nil, fmt.Errorf("not a valid SM2 private key")
	default:
		return nil, fmt.Errorf("unsupported key type")
	}
}

func PriToPub(priKey string) (string, error) {
	priBytes, err := hex.DecodeString(priKey)
	if err != nil {
		return "", err
	}
	priObj, err := sm2.NewPrivateKey(priBytes)
	if err != nil {
		return "", err
	}

	pubBytes, err := smx509.MarshalPKIXPublicKey(&priObj.PublicKey)
	if err != nil {
		return "", err
	}
	return hex.EncodeToString(pubBytes), nil
}