package license

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"encoding/asn1"
	"encoding/base64"
	"license-server/storage"
	"math/big"
)

var privateKey *ecdsa.PrivateKey

// 显式从 main 传入 DB，保证先建库再读写
func InitCrypto(db storage.Database) {
	// 1. 读取
	if pem, ok := db.GetPrivateKey(); ok {
		if der, err := base64.StdEncoding.DecodeString(pem); err == nil {
			if key, err := x509.ParseECPrivateKey(der); err == nil {
				privateKey = key
				return
			}
		}
	}
	// 2. 不存在或解析失败 → 重新生成并落库
	key, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	der, _ := x509.MarshalECPrivateKey(key)
	db.SavePrivateKey(base64.StdEncoding.EncodeToString(der))
	privateKey = key
}

// ecdsaSignature 是 ASN.1 编码中使用的结构体
type ecdsaSignature struct {
	R, S *big.Int
}

func SignPayload(message []byte) (string, error) {
	hash := sha256.Sum256(message)
	r, s, err := ecdsa.Sign(rand.Reader, privateKey, hash[:])
	if err != nil {
		return "", err
	}

	sig, err := asn1.Marshal(ecdsaSignature{r, s})
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(sig), nil
}

func VerifySignature(pub *ecdsa.PublicKey, message []byte, signatureBase64 string) bool {
	sigBytes, err := base64.StdEncoding.DecodeString(signatureBase64)
	if err != nil {
		return false
	}

	var sig ecdsaSignature
	_, err = asn1.Unmarshal(sigBytes, &sig)
	if err != nil {
		return false
	}

	hash := sha256.Sum256(message)
	return ecdsa.Verify(pub, hash[:], sig.R, sig.S)
}

func GetPublicKey() *ecdsa.PublicKey {
	return &privateKey.PublicKey
}