package mls

// MLS cipher suites (RFC 9420 §5.1, §17.1 + RFC 9180).
// Supported:
//   0x0003: DHKEM(X25519, HKDF-SHA256) + ChaCha20-Poly1305    + SHA-256 + Ed25519
//
// Suite 0x0001 (AES-128-GCM) is defined in the wire format but not implemented here.
// Smesh hardcodes suite 0x0003 in marmot; receiving a 0x0001 group is an error.

import (
	"crypto/ed25519"
	"errors"
	"smesh.lol/web/common/crypto/chacha20poly1305"
	"smesh.lol/web/common/crypto/hkdf"
	"crypto/hmac"
	"crypto/sha256"
	"smesh.lol/web/common/jsbridge/subtle"
	"smesh.lol/web/common/jsbridge/x25519"
)

var (
	errInvalidKeySize    = errors.New("mls: invalid key size")
	errAEADOpenFailed    = errors.New("mls: AEAD open failed")
	errHPKEDecryptFailed = errors.New("mls: HPKE decryption failed")
)

// --- Suite parameters (shared) ---

const (
	hashSize   = 32 // SHA-256 (both suites)
	aeadNonce  = 12 // AEAD nonce size (both suites)
	aeadTag    = 16 // AEAD tag size (both suites)
	kemKeySize = 32 // X25519 (both suites)
)

func (cs CipherSuite) assertSupported() {
	switch cs {
	case CipherSuite0x0003:
	default:
		panic("mls: unsupported cipher suite")
	}
}

// --- Public metadata ---

func (cs CipherSuite) HashSize() int {
	cs.assertSupported()
	return hashSize
}

func (cs CipherSuite) AEADKeySize() int {
	cs.assertSupported()
	return 32 // ChaCha20-Poly1305
}

func (cs CipherSuite) AEADNonceSize() int {
	cs.assertSupported()
	return aeadNonce
}

func (cs CipherSuite) ExtractSize() int {
	cs.assertSupported()
	return hashSize
}

// --- KEM/HPKE suite IDs (RFC 9180 §4, §5) ---

// kemSuiteID = "KEM" || I2OSP(kem_id=0x0020, 2) — DHKEM(X25519, HKDF-SHA256) for both suites.
var kemSuiteID = []byte("KEM\x00\x20")

// hpkeSuiteID = "HPKE" || I2OSP(kem=0x0020,2) || I2OSP(kdf=0x0001,2) || I2OSP(aead=0x0003,2)
// Suite 0x0003: ChaCha20-Poly1305 only.
func (cs CipherSuite) hpkeSuiteID() []byte {
	cs.assertSupported()
	return []byte("HPKE\x00\x20\x00\x01\x00\x03")
}

// --- Hash ---

func (cs CipherSuite) hash(data []byte) []byte {
	cs.assertSupported()
	h := sha256.Sum(data)
	return h[:]
}

// --- MAC ---

func (cs CipherSuite) signMAC(key, message []byte) []byte {
	cs.assertSupported()
	mac := hmac.Sum(key, message)
	return mac[:]
}

func (cs CipherSuite) verifyMAC(key, message, tag []byte) bool {
	cs.assertSupported()
	expected := hmac.Sum(key, message)
	return hmacEqual(tag, expected[:])
}

func hmacEqual(a, b []byte) bool {
	if len(a) != len(b) {
		return false
	}
	var v byte
	for i := range a {
		v |= a[i] ^ b[i]
	}
	return v == 0
}

// --- HKDF ---

func (cs CipherSuite) hkdfExtract(salt, ikm []byte) []byte {
	cs.assertSupported()
	h := hkdf.Extract(salt, ikm)
	return h[:]
}

func (cs CipherSuite) hkdfExpand(prk, info []byte, length int) []byte {
	cs.assertSupported()
	return hkdf.Expand(prk, info, length)
}

// --- MLS Labeled Expand/Extract (RFC 9420 §8) ---

func (cs CipherSuite) expandWithLabel(secret, label, context []byte, length uint16) ([]byte, error) {
	cs.assertSupported()
	mlsLabel := append([]byte("MLS 1.0 "), label...)

	// KDFLabel = length (2) ‖ opaqueVec(label) ‖ opaqueVec(context)
	var w Writer
	w.addUint16(length)
	w.writeOpaqueVec(mlsLabel)
	w.writeOpaqueVec(context)
	kdfLabel, err := w.bytes()
	if err != nil {
		return nil, err
	}

	return hkdf.Expand(secret, kdfLabel, int(length)), nil
}

func (cs CipherSuite) deriveSecret(secret, label []byte) ([]byte, error) {
	return cs.expandWithLabel(secret, label, nil, uint16(hashSize))
}

// --- Ref hash ---

func (cs CipherSuite) refHash(label, value []byte) ([]byte, error) {
	cs.assertSupported()
	var w Writer
	w.writeOpaqueVec(label)
	w.writeOpaqueVec(value)
	input, err := w.bytes()
	if err != nil {
		return nil, err
	}
	h := sha256.Sum(input)
	return h[:], nil
}

// --- Signatures (Ed25519 for both suites) ---

func (cs CipherSuite) signWithLabel(signKey signaturePrivateKey, label, content []byte) ([]byte, error) {
	cs.assertSupported()
	signContent, err := marshalSignContent(label, content)
	if err != nil {
		return nil, err
	}
	privkey := ed25519.NewKeyFromSeed([]byte(signKey))
	return ed25519.Sign(privkey, signContent), nil
}

func (cs CipherSuite) verifyWithLabel(verifKey signaturePublicKey, label, content, signValue []byte) bool {
	cs.assertSupported()
	signContent, err := marshalSignContent(label, content)
	if err != nil {
		return false
	}
	return ed25519.Verify(ed25519.PublicKey(verifKey), signContent, signValue)
}

func marshalSignContent(label, content []byte) ([]byte, error) {
	mlsLabel := append([]byte("MLS 1.0 "), label...)
	var w Writer
	w.writeOpaqueVec(mlsLabel)
	w.writeOpaqueVec(content)
	return w.bytes()
}

func (cs CipherSuite) generateSignatureKeyPair() (signaturePublicKey, signaturePrivateKey, error) {
	cs.assertSupported()
	seed := []byte{:32}
	subtle.RandomBytes(seed)
	privkey := ed25519.NewKeyFromSeed(seed)
	pub := []byte{:32}
	copy(pub, []byte(privkey)[32:])
	return signaturePublicKey(pub), signaturePrivateKey(seed), nil
}

// --- AEAD dispatch ---

func (cs CipherSuite) aeadSeal(key, nonce, plaintext, aad []byte) ([]byte, error) {
	cs.assertSupported()
	if len(key) != cs.AEADKeySize() {
		return nil, errInvalidKeySize
	}
	var k [32]byte
	var n [12]byte
	copy(k[:], key)
	copy(n[:], nonce)
	return chacha20poly1305.Seal(k, n, plaintext, aad), nil
}

func (cs CipherSuite) aeadOpen(key, nonce, ciphertext, aad []byte) ([]byte, error) {
	cs.assertSupported()
	if len(key) != cs.AEADKeySize() {
		return nil, errInvalidKeySize
	}
	var k [32]byte
	var n [12]byte
	copy(k[:], key)
	copy(n[:], nonce)
	pt, ok := chacha20poly1305.Open(k, n, ciphertext, aad)
	if !ok {
		return nil, errAEADOpenFailed
	}
	return pt, nil
}

// --- DHKEM(X25519, HKDF-SHA256) — RFC 9180 §4.1 (both suites) ---

// labeledExtractKEM: HKDF-Extract(salt, "HPKE-v1" || kemSuiteID || label || ikm)
func labeledExtractKEM(salt, label, ikm []byte) []byte {
	labeled := []byte("HPKE-v1")
	labeled = append(labeled, kemSuiteID...)
	labeled = append(labeled, label...)
	labeled = append(labeled, ikm...)
	h := hkdf.Extract(salt, labeled)
	return h[:]
}

// labeledExpandKEM: HKDF-Expand(prk, I2OSP(L,2) || "HPKE-v1" || kemSuiteID || label || info, L)
func labeledExpandKEM(prk, label, info []byte, length int) []byte {
	labeled := []byte{byte(length >> 8), byte(length)}
	labeled = append(labeled, []byte("HPKE-v1")...)
	labeled = append(labeled, kemSuiteID...)
	labeled = append(labeled, label...)
	labeled = append(labeled, info...)
	return hkdf.Expand(prk, labeled, length)
}

// labeledExtractHPKE: HKDF-Extract(salt, "HPKE-v1" || hpkeSuiteID || label || ikm)
func (cs CipherSuite) labeledExtractHPKE(salt, label, ikm []byte) []byte {
	labeled := []byte("HPKE-v1")
	labeled = append(labeled, cs.hpkeSuiteID()...)
	labeled = append(labeled, label...)
	labeled = append(labeled, ikm...)
	h := hkdf.Extract(salt, labeled)
	return h[:]
}

// labeledExpandHPKE: HKDF-Expand(prk, I2OSP(L,2) || "HPKE-v1" || hpkeSuiteID || label || info, L)
func (cs CipherSuite) labeledExpandHPKE(prk, label, info []byte, length int) []byte {
	labeled := []byte{byte(length >> 8), byte(length)}
	labeled = append(labeled, []byte("HPKE-v1")...)
	labeled = append(labeled, cs.hpkeSuiteID()...)
	labeled = append(labeled, label...)
	labeled = append(labeled, info...)
	return hkdf.Expand(prk, labeled, length)
}

// extractAndExpand derives the shared secret from DH output per RFC 9180 §4.1.
//   eae_prk = LabeledExtract("", "eae_prk", dh)
//   shared_secret = LabeledExpand(eae_prk, "shared_secret", kem_context, Nsecret)
func extractAndExpand(dh, kemContext []byte) []byte {
	prk := labeledExtractKEM(nil, []byte("eae_prk"), dh)
	return labeledExpandKEM(prk, []byte("shared_secret"), kemContext, hashSize)
}

// --- KEM operations (X25519 for both suites) ---

func (cs CipherSuite) kemEncap(pkR hpkePublicKey) (sharedSecret, enc []byte, err error) {
	cs.assertSupported()
	// Generate ephemeral keypair
	skE := []byte{:kemKeySize}
	subtle.RandomBytes(skE)
	pkE := x25519.ScalarBaseMult(skE)

	// DH(skE, pkR)
	dh := x25519.ScalarMult(skE, []byte(pkR))

	enc = pkE
	kemContext := append(enc, []byte(pkR)...)
	sharedSecret = extractAndExpand(dh, kemContext)
	return sharedSecret, enc, nil
}

func (cs CipherSuite) kemDecap(enc []byte, skR hpkePrivateKey) ([]byte, error) {
	cs.assertSupported()
	pkR := x25519.ScalarBaseMult([]byte(skR))
	dh := x25519.ScalarMult([]byte(skR), enc)
	kemContext := append([]byte{:0:len(enc)+len(pkR)}, enc...)
	kemContext = append(kemContext, pkR...)
	return extractAndExpand(dh, kemContext), nil
}

// --- HPKE Base mode setup (RFC 9180 §5.1) ---

type hpkeSealer struct {
	cs    CipherSuite
	key   []byte
	nonce []byte
}

func (s *hpkeSealer) seal(plaintext, aad []byte) []byte {
	ct, err := s.cs.aeadSeal(s.key, s.nonce, plaintext, aad)
	if err != nil {
		return nil
	}
	return ct
}

type hpkeOpener struct {
	cs    CipherSuite
	key   []byte
	nonce []byte
}

func (o *hpkeOpener) open(ciphertext, aad []byte) ([]byte, bool) {
	pt, err := o.cs.aeadOpen(o.key, o.nonce, ciphertext, aad)
	if err != nil {
		return nil, false
	}
	return pt, true
}

func (cs CipherSuite) hpkeKeySchedule(sharedSecret, info []byte) (key, baseNonce []byte) {
	// mode = 0x00 (Base mode)
	mode := byte(0x00)

	// psk_id_hash = LabeledExtract("", "psk_id_hash", "")
	pskIDHash := cs.labeledExtractHPKE(nil, []byte("psk_id_hash"), nil)
	// info_hash = LabeledExtract("", "info_hash", info)
	infoHash := cs.labeledExtractHPKE(nil, []byte("info_hash"), info)

	// ks_context = mode || psk_id_hash || info_hash
	ksContext := []byte{mode}
	ksContext = append(ksContext, pskIDHash...)
	ksContext = append(ksContext, infoHash...)

	// secret = LabeledExtract(shared_secret, "secret", psk="")
	secret := cs.labeledExtractHPKE(sharedSecret, []byte("secret"), nil)

	// key = LabeledExpand(secret, "key", ks_context, Nk)
	key = cs.labeledExpandHPKE(secret, []byte("key"), ksContext, cs.AEADKeySize())
	// base_nonce = LabeledExpand(secret, "base_nonce", ks_context, Nn)
	baseNonce = cs.labeledExpandHPKE(secret, []byte("base_nonce"), ksContext, aeadNonce)
	return key, baseNonce
}

// encryptWithLabel implements HPKE single-shot encrypt (RFC 9180 §6.1).
func (cs CipherSuite) encryptWithLabel(publicKey hpkePublicKey, label, context, plaintext []byte) (kemOutput, ciphertext []byte, err error) {
	encryptContext, err := marshalEncryptContext(label, context)
	if err != nil {
		return nil, nil, err
	}

	sharedSecret, enc, err := cs.kemEncap(publicKey)
	if err != nil {
		return nil, nil, err
	}

	key, baseNonce := cs.hpkeKeySchedule(sharedSecret, encryptContext)
	sealer := &hpkeSealer{cs: cs, key: key, nonce: baseNonce}
	ciphertext = sealer.seal(plaintext, nil)
	return enc, ciphertext, nil
}

// decryptWithLabel implements HPKE single-shot decrypt (RFC 9180 §6.1).
func (cs CipherSuite) decryptWithLabel(privateKey hpkePrivateKey, label, context, kemOutput, ciphertext []byte) ([]byte, error) {
	encryptContext, err := marshalEncryptContext(label, context)
	if err != nil {
		return nil, err
	}
	sharedSecret, err := cs.kemDecap(kemOutput, privateKey)
	if err != nil {
		return nil, err
	}
	key, baseNonce := cs.hpkeKeySchedule(sharedSecret, encryptContext)
	opener := &hpkeOpener{cs: cs, key: key, nonce: baseNonce}
	plaintext, ok := opener.open(ciphertext, nil)
	if !ok {
		return nil, errHPKEDecryptFailed
	}
	return plaintext, nil
}

func marshalEncryptContext(label, context []byte) ([]byte, error) {
	mlsLabel := append([]byte("MLS 1.0 "), label...)
	var w Writer
	w.writeOpaqueVec(mlsLabel)
	w.writeOpaqueVec(context)
	return w.bytes()
}

// --- Key generation ---

func (cs CipherSuite) generateEncryptionKeyPair() (hpkePublicKey, hpkePrivateKey, error) {
	cs.assertSupported()
	sk := []byte{:kemKeySize}
	subtle.RandomBytes(sk)
	pk := x25519.ScalarBaseMult(sk)
	return hpkePublicKey(pk), hpkePrivateKey(sk), nil
}

func (cs CipherSuite) deriveEncryptionKeyPair(seed []byte) (hpkePublicKey, hpkePrivateKey, error) {
	cs.assertSupported()
	// RFC 9180 §7.1.3: DeriveKeyPair for X25519
	// dkp_prk = LabeledExtract("", "dkp_prk", seed)
	dkpPrk := labeledExtractKEM(nil, []byte("dkp_prk"), seed)
	// sk = LabeledExpand(dkp_prk, "sk", "", Nsk=32)
	sk := labeledExpandKEM(dkpPrk, []byte("sk"), nil, kemKeySize)
	pk := x25519.ScalarBaseMult(sk)
	return hpkePublicKey(pk), hpkePrivateKey(sk), nil
}

func (cs CipherSuite) randomBytes(n int) []byte {
	buf := []byte{:n}
	subtle.RandomBytes(buf)
	return buf
}