package mlkem768

import (
	"bytes"
	"crypto/mlkem"
	"crypto/sha256"
	"crypto/sha3"
	"errors"
	"fmt"
	"io"

	"git.tswf.io/infra/go-synapse-backupper/pkg/domain/crypto"
)

// suiteID is the scheme identifier for ML-KEM-768.
const suiteID uint16 = 0x0006

// ErrDecapsulationFailed is returned when ciphertext decapsulation fails,
// typically due to a tampered or invalid ciphertext.
var ErrDecapsulationFailed = errors.New("decapsulation failed")

// DefaultRegistry is the package-level registry for ML-KEM-768.
var DefaultRegistry = crypto.NewRegistry()

// kemAdapter wraps the Go stdlib crypto/mlkem implementation to satisfy
// the pkg/domain/crypto.KEM interface.
type kemAdapter struct{}

// New creates a new KEM adapter instance.
func New() crypto.KEM {
	return &kemAdapter{}
}

// SchemeID returns the ML-KEM-768 scheme identifier (0x0006).
func (k *kemAdapter) SchemeID() uint16 {
	return suiteID
}

// GenerateKeyPair generates a new ML-KEM-768 key pair.
//
// Note: the rand parameter is part of the KEM interface contract but is
// ignored here because the stdlib crypto/mlkem.GenerateKey768 uses
// crypto/rand internally.
func (k *kemAdapter) GenerateKeyPair(
	rand io.Reader,
) (
	crypto.RecipientPub,
	crypto.RecipientPriv,
	error,
) {
	dk, err := mlkem.GenerateKey768()
	if err != nil {
		return nil, nil, err
	}

	ek := dk.EncapsulationKey()
	rawPub := ek.Bytes()
	keyID := computeKeyID(rawPub)

	pub := &pubKey{
		key:   ek,
		keyID: keyID,
	}
	priv := &privKey{
		key:   dk,
		keyID: keyID,
	}

	return pub, priv, nil
}

// Encapsulate generates a shared secret and ciphertext for the given public key.
//
// Note: the rand parameter is part of the KEM interface contract but is
// ignored here because the stdlib (*EncapsulationKey768).Encapsulate uses
// crypto/rand internally.
//
// CRITICAL: the stdlib returns (sharedKey, ciphertext) but this adapter
// swaps the order to (ciphertext, sharedSecret) to match the domain KEM
// interface contract.
func (k *kemAdapter) Encapsulate(
	pub crypto.RecipientPub,
	rand io.Reader,
) (
	ciphertext []byte,
	sharedSecret []byte,
	err error,
) {
	p, ok := pub.(*pubKey)
	if !ok {
		raw := pub.Raw()
		ek, parseErr := mlkem.NewEncapsulationKey768(raw)
		if parseErr != nil {
			return nil, nil, fmt.Errorf("invalid public key for ML-KEM-768: %w", parseErr)
		}
		p = &pubKey{key: ek, keyID: computeKeyID(raw)}
	}

	// stdlib returns (sharedKey, ciphertext); we swap to (ciphertext, sharedSecret).
	ss, ct := p.key.Encapsulate()

	return ct, ss, nil
}

// LoadPriv loads an ML-KEM-768 private key from raw bytes.
func (k *kemAdapter) LoadPriv(
	raw []byte,
) (
	crypto.RecipientPriv,
	error,
) {
	dk, err := mlkem.NewDecapsulationKey768(raw)
	if err != nil {
		return nil, fmt.Errorf("invalid ML-KEM-768 private key: %w", err)
	}
	pubRaw := dk.EncapsulationKey().Bytes()
	return &privKey{key: dk, keyID: computeKeyID(pubRaw)}, nil
}

// Decapsulate recovers the shared secret from a ciphertext using the private key.
// For tampered ciphertexts, it returns ErrDecapsulationFailed.
func (k *kemAdapter) Decapsulate(
	priv crypto.RecipientPriv,
	ciphertext []byte,
) (
	sharedSecret []byte,
	err error,
) {
	p, ok := priv.(*privKey)
	if !ok {
		return nil, fmt.Errorf("invalid private key type for ML-KEM-768")
	}

	ss, err := p.key.Decapsulate(ciphertext)
	if err != nil {
		return nil, errors.Join(ErrDecapsulationFailed, err)
	}

	// crypto/mlkem.Decapsulate implements implicit rejection: it returns a
	// pseudorandom shared secret instead of an error for invalid ciphertexts.
	// Perform explicit rejection by recomputing the implicit-rejection value
	// Kout = SHAKE256(z || ciphertext) and comparing it with the result.
	// If they match, the ciphertext was invalid.
	seed := p.key.Bytes()
	z := seed[32:]
	shake := sha3.NewSHAKE256()
	// sha3.ShakeHash.Write/Read never return an error; explicitly ignore to satisfy errcheck.
	_, _ = shake.Write(z)
	_, _ = shake.Write(ciphertext)
	computedKout := make([]byte, mlkem.SharedKeySize)
	_, _ = shake.Read(computedKout)

	if bytes.Equal(ss, computedKout) {
		return nil, ErrDecapsulationFailed
	}

	return ss, nil
}

// pubKey wraps *mlkem.EncapsulationKey768 to satisfy crypto.RecipientPub.
type pubKey struct {
	key   *mlkem.EncapsulationKey768
	keyID []byte
}

func (p *pubKey) SchemeID() uint16 { return suiteID }
func (p *pubKey) KeyID() []byte    { return p.keyID }
func (p *pubKey) Raw() []byte      { return p.key.Bytes() }

// privKey wraps *mlkem.DecapsulationKey768 to satisfy crypto.RecipientPriv.
type privKey struct {
	key   *mlkem.DecapsulationKey768
	keyID []byte
}

func (p *privKey) SchemeID() uint16 { return suiteID }
func (p *privKey) KeyID() []byte    { return p.keyID }
func (p *privKey) Raw() []byte      { return p.key.Bytes() }

// computeKeyID derives the first 8 bytes of SHA-256 over the first 8 bytes of raw key material.
func computeKeyID(raw []byte) []byte {
	h := sha256.Sum256(raw[:8])
	return h[:8]
}

// init registers the ML-KEM-768 factory under suiteID 0x0006.
func init() {
	_ = DefaultRegistry.Register(
		suiteID,
		func() crypto.KEM {
			return New()
		},
	)
}