package cert

import (
	"crypto/mldsa"
	"crypto/mlkem"
	"fmt"
	"slices"
	"strings"
)

// PQSuite names the post-quantum algorithms of a version 3 certificate, the way Curve names its classical ones: the
// key encapsulation mechanism of a host's key, and the signature algorithm of a CA's key and of the certificates it
// signs. A CA and every certificate it signs share a suite. Every version 3 certificate names its suite.
type PQSuite uint8

const (
	// PQSuiteUnknown is never valid in a certificate.
	PQSuiteUnknown PQSuite = 0

	// PQSuite_MLKEM1024_MLDSA87 gives hosts ML-KEM-1024 keys and CAs ML-DSA-87 keys and signatures.
	PQSuite_MLKEM1024_MLDSA87 PQSuite = 1
)

// DefaultPQSuite is the suite tools use when none is asked for.
const DefaultPQSuite = PQSuite_MLKEM1024_MLDSA87

// pqKEM is the key encapsulation mechanism of a suite, what host keys use. Private keys are seeds.
type pqKEM struct {
	name           string // as it appears in key banners
	publicKeySize  int
	privateKeySize int
	generate       func() (pub, priv []byte, err error)
	publicKey      func(priv []byte) ([]byte, error)
	checkPublicKey func(pub []byte) bool
}

// pqDSA is the signature algorithm of a suite, what CA keys and certificate signatures use. Private keys are seeds.
type pqDSA struct {
	name           string // as it appears in key banners
	publicKeySize  int
	privateKeySize int
	signatureSize  int
	generate       func() (pub, priv []byte, err error)
	publicKey      func(priv []byte) ([]byte, error)
	checkPublicKey func(pub []byte) bool
	signer         func(priv []byte) (SignerLambda, error) // must sign deterministically
	verify         func(pub, msg, sig []byte) bool
}

type pqSuite struct {
	kem pqKEM
	dsa pqDSA
}

// pqSuites holds every known suite, see registerPQSuite.
var pqSuites = map[PQSuite]*pqSuite{}

func init() {
	registerPQSuite(PQSuite_MLKEM1024_MLDSA87, &pqSuite{kem: mlkem1024, dsa: mldsa87})
}

// registerPQSuite makes a suite known, along with the banners of its composite keys.
func registerPQSuite(id PQSuite, s *pqSuite) {
	pqSuites[id] = s
	for _, curve := range []Curve{Curve_CURVE25519, Curve_P256} {
		for _, kind := range pqKeyKinds {
			pqKeyBanners[pqBannerName(kind, curve, s)] = pqKeyBanner{kind: kind, curve: curve, size: pqKeySize(kind, curve, s)}
		}
	}
}

var mlkem1024 = pqKEM{
	name:           "MLKEM1024",
	publicKeySize:  mlkem.EncapsulationKeySize1024,
	privateKeySize: mlkem.SeedSize,
	generate: func() ([]byte, []byte, error) {
		dk, err := mlkem.GenerateKey1024()
		if err != nil {
			return nil, nil, err
		}
		return dk.EncapsulationKey().Bytes(), dk.Bytes(), nil
	},
	publicKey: func(priv []byte) ([]byte, error) {
		dk, err := mlkem.NewDecapsulationKey1024(priv)
		if err != nil {
			return nil, err
		}
		return dk.EncapsulationKey().Bytes(), nil
	},
	checkPublicKey: func(pub []byte) bool {
		_, err := mlkem.NewEncapsulationKey1024(pub)
		return err == nil
	},
}

var mldsa87 = pqDSA{
	name:           "MLDSA87",
	publicKeySize:  mldsa.MLDSA87PublicKeySize,
	privateKeySize: mldsa.PrivateKeySize,
	signatureSize:  mldsa.MLDSA87SignatureSize,
	generate: func() ([]byte, []byte, error) {
		sk, err := mldsa.GenerateKey(mldsa.MLDSA87())
		if err != nil {
			return nil, nil, err
		}
		return sk.PublicKey().Bytes(), sk.Bytes(), nil
	},
	publicKey: func(priv []byte) ([]byte, error) {
		sk, err := mldsa.NewPrivateKey(mldsa.MLDSA87(), priv)
		if err != nil {
			return nil, err
		}
		return sk.PublicKey().Bytes(), nil
	},
	checkPublicKey: func(pub []byte) bool {
		_, err := mldsa.NewPublicKey(mldsa.MLDSA87(), pub)
		return err == nil
	},
	// Deterministic signing (FIPS 204 section 3.4): the hedged variant guards signers an attacker can fault
	// physically, which a CA key in a file is not, and the Ed25519 signature beside it is deterministic too.
	signer: func(priv []byte) (SignerLambda, error) {
		sk, err := mldsa.NewPrivateKey(mldsa.MLDSA87(), priv)
		if err != nil {
			return nil, err
		}
		return func(b []byte) ([]byte, error) {
			return sk.SignDeterministic(b, &mldsa.Options{Context: pqSignatureContext})
		}, nil
	},
	verify: func(pub, msg, sig []byte) bool {
		pk, err := mldsa.NewPublicKey(mldsa.MLDSA87(), pub)
		if err != nil {
			return false
		}
		return mldsa.Verify(pk, msg, sig, &mldsa.Options{Context: pqSignatureContext}) == nil
	},
}

// PQSuites returns every known suite.
func PQSuites() []PQSuite {
	s := make([]PQSuite, 0, len(pqSuites))
	for k := range pqSuites {
		s = append(s, k)
	}
	slices.Sort(s)
	return s
}

// ParsePQSuite returns the suite named by s, as String names it, ignoring case.
func ParsePQSuite(s string) (PQSuite, error) {
	for _, k := range PQSuites() {
		if strings.EqualFold(s, k.String()) {
			return k, nil
		}
	}
	return PQSuiteUnknown, fmt.Errorf("%w: %q", ErrUnknownPQSuite, s)
}

// PQSuiteOf returns the suite of a version 3 certificate, or PQSuiteUnknown for any other version.
func PQSuiteOf(c Certificate) PQSuite {
	if c3, ok := c.(*certificateV3); ok {
		return c3.pqSuite
	}
	return PQSuiteUnknown
}

func (s PQSuite) String() string {
	if p := pqSuites[s]; p != nil {
		return p.kem.name + "_" + p.dsa.name
	}
	return fmt.Sprintf("PQSuite(%d)", uint8(s))
}

// Valid reports whether s is a known suite.
func (s PQSuite) Valid() bool {
	return pqSuites[s] != nil
}

// PublicKeySize returns the size of the post-quantum half of a composite public key, a CA's or a host's.
func (s PQSuite) PublicKeySize(isCA bool) int {
	p := pqSuites[s]
	switch {
	case p == nil:
		return 0
	case isCA:
		return p.dsa.publicKeySize
	default:
		return p.kem.publicKeySize
	}
}

// PrivateKeySize returns the size of the post-quantum half of a composite private key, a CA's or a host's.
func (s PQSuite) PrivateKeySize(isCA bool) int {
	p := pqSuites[s]
	switch {
	case p == nil:
		return 0
	case isCA:
		return p.dsa.privateKeySize
	default:
		return p.kem.privateKeySize
	}
}

// GenerateKey makes a post-quantum keypair, a CA's signing key or a host's KEM key, to append to a classical one.
func (s PQSuite) GenerateKey(isCA bool) (pub, priv []byte, err error) {
	p := pqSuites[s]
	switch {
	case p == nil:
		return nil, nil, fmt.Errorf("%w: %v", ErrUnknownPQSuite, s)
	case isCA:
		return p.dsa.generate()
	default:
		return p.kem.generate()
	}
}

// publicKey derives the post-quantum public key from the post-quantum half of a private key.
func (s PQSuite) publicKey(isCA bool, priv []byte) ([]byte, error) {
	p := pqSuites[s]
	switch {
	case p == nil:
		return nil, fmt.Errorf("%w: %v", ErrUnknownPQSuite, s)
	case isCA:
		return p.dsa.publicKey(priv)
	default:
		return p.kem.publicKey(priv)
	}
}

// checkPublicKey reports whether pub is a well formed post-quantum public key, a CA's or a host's.
func (s PQSuite) checkPublicKey(isCA bool, pub []byte) bool {
	p := pqSuites[s]
	switch {
	case p == nil, len(pub) != s.PublicKeySize(isCA):
		return false
	case isCA:
		return p.dsa.checkPublicKey(pub)
	default:
		return p.kem.checkPublicKey(pub)
	}
}

// signatureSize returns the size of the post-quantum signature of a certificate in this suite.
func (s PQSuite) signatureSize() int {
	if p := pqSuites[s]; p != nil {
		return p.dsa.signatureSize
	}
	return 0
}

// SplitPublicKey splits a composite public key into its classical and post-quantum halves.
func (s PQSuite) SplitPublicKey(isCA bool, key []byte) (classical, pq []byte, err error) {
	n := s.PublicKeySize(isCA)
	if n == 0 || len(key) <= n {
		return nil, nil, ErrInvalidPublicKey
	}
	return key[:len(key)-n], key[len(key)-n:], nil
}

// SplitPrivateKey splits a composite private key into its classical and post-quantum halves.
func (s PQSuite) SplitPrivateKey(isCA bool, key []byte) (classical, pq []byte, err error) {
	n := s.PrivateKeySize(isCA)
	if n == 0 || len(key) <= n {
		return nil, nil, ErrInvalidPrivateKey
	}
	return key[:len(key)-n], key[len(key)-n:], nil
}