// Package composite implements the v2 .pqenc artifact format: a dual-KEM
// (post-quantum + classical) hybrid AEAD with inherent downgrade resistance.
//
// Layout of an artifact (all multi-byte fields are big-endian):
//
//	[magic u32 = 0x47535051 "GSPQ"]              [0:4]
//	[version u16 = 0x0002]                       [4:6]
//	[flags u32 = 0x00000000 (reserved, must==0)]  [6:10]
//	[nRecipients u8]                             [10]
//
//	for each recipient slot i (positional: slot 0 = PQ, slot 1 = classical):
//	  [schemeID u16] [keyID 8B] [ctLen u32] [ciphertext ctLen bytes]
//
//	[wrapNonce 12B] [wrappedCEK 48B] [firstPayloadNonce 12B]
//
//	chunk records (until a chunk with flags&0x01==1 is seen):
//	  [len u32 = ciphertext length incl. 16B tag] [flags u8] [ciphertext]
//
// The content key (CEK, 32B) is wrapped with kekFinal — the HKDF combiner
// of the two KEM shared secrets — using AES-256-GCM (AAD = version u16 BE).
// Payload is encrypted with AES-256-GCM under CEK in 64 KiB chunks; a
// zero-length final marker chunk (flags&0x01==1) is ALWAYS emitted on EOF
// regardless of the previous chunk's fullness, providing explicit AEAD
// integrity for the logical end-of-stream.
package composite

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hkdf"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"io"

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

// Sentinel errors surfaced by the composite Encryptor/Decryptor.
var (
	// ErrMalformedHeader indicates the artifact header is structurally
	// invalid (bad magic, reserved flags, out-of-range nRecipients, oversized
	// recipient ciphertext length, or premature EOF while reading fixed
	// metadata).
	ErrMalformedHeader = errors.New("composite: malformed header")

	// ErrUnsupportedVersion indicates the artifact's version field is not
	// 0x0002. Decryption stops at the version check — no AES/GCM operations
	// are attempted and no recipient state is allocated.
	ErrUnsupportedVersion = errors.New("composite: unsupported artifact version")

	// ErrWrongKeys indicates one of the supplied private keys does not match
	// the recipient slot it was routed to (positional). Signal: the priv's
	// KeyID() does not equal the slot's keyID, OR the KEM Decapsulate failed
	// for the slot's ciphertext.
	ErrWrongKeys = errors.New("composite: wrong recipient keys")

	// ErrTamperingDetected indicates the wrappedCEK or a payload chunk failed
	// AES-GCM authentication: the cancellation or modification of ciphertext
	// bytes is cryptographically rejected.
	ErrTamperingDetected = errors.New("composite: tampering detected")

	// ErrNonceCounterWrapped indicates the per-chunk 64-bit counter
	// (chunkNonce[4:12]) wrapped around to zero while encrypting or
	// decrypting an additional chunk — the nonce sequence is exhausted.
	ErrNonceCounterWrapped = errors.New("composite: nonce counter wrapped")

	// ErrMalformedChunk indicates a chunk record failed structural
	// validation: zero-length non-final chunk (infinite-loop DoS) or
	// oversized ciphertext (over the 64 KiB+16 maximum).
	ErrMalformedChunk = errors.New("composite: malformed chunk")

	// ErrUnexpectedEOF indicates the chunk stream ended before any chunk
	// with flags&0x01==1 (logical end-of-stream marker) was observed.
	ErrUnexpectedEOF = errors.New("composite: unexpected end of stream")

	// ErrPlaintextTooLarge indicates the decrypted payload would exceed the
	// configured maximum plaintext size; decryption is aborted before the
	// limit is crossed to prevent unbounded disk consumption.
	ErrPlaintextTooLarge = errors.New("composite: plaintext size limit exceeded")
)

// Format constants.
const (
	magic                     uint32 = 0x47535051 // "GSPQ"
	version                   uint16 = 0x0002
	flags                     uint32 = 0x00000000
	maxRecipients             int    = 2
	chunkSize                 int    = 64 * 1024
	gcmTagLen                 int    = 16
	maxRecipientCiphertextLen int    = 1 << 20 // MiB cap on a single recipient ciphertext
	wrapNonceLen              int    = 12
	wrappedCekLen             int    = 48 // 32-byte CEK + 16-byte GCM tag
	firstPayloadNonceLen      int    = 12
	kekLen                    int    = 32
	maxPlaintextSize          int64  = 1 << 40 // 1 TiB cap on decrypted output
	infoPq                    string = "git.tswf.io/infra/go-synapse-backupper/v2/kek/pq"
	infoComposite             string = "git.tswf.io/infra/go-synapse-backupper/v2/kek/composite"
)

// Chunk flag bits.
const (
	flagFinal byte = 0x01
)

// encryptor is the composite Encryptor implementation backed by a Registry
// of KEM factories. Recipients are routed positionally: slot 0 = PQ, slot 1
// = classical.
type encryptor struct {
	registry crypto.Registry
}

// NewEncryptor returns a composite Encryptor that resolves KEM schemes via
// the provided Registry (Registry.Lookup(schemeID)).
func NewEncryptor(
	registry crypto.Registry,
) crypto.Encryptor {
	return &encryptor{
		registry: registry,
	}
}

// Encrypt encrypts plaintext to multiple recipients under the v2 artifact
// format and streams the result to sink. Exactly two recipients must be
// supplied — slot 0 (PQ) and slot 1 (classical).
func (e *encryptor) Encrypt(
	plaintext io.Reader,
	recipients []crypto.RecipientPub,
	sink io.Writer,
	rand io.Reader,
) error {
	if len(recipients) != maxRecipients {
		return ErrMalformedHeader
	}
	pqPub, classicalPub := recipients[0], recipients[1]

	// Generate the per-message content key (32B for AES-256).
	cek := make([]byte, kekLen)
	if _, err := io.ReadFull(rand, cek); err != nil {
		return err
	}

	// Bind KEM adapters through the registry — no direct adapter imports.
	pqFactory, err := e.registry.Lookup(pqPub.SchemeID())
	if err != nil {
		return err
	}
	classicalFactory, err := e.registry.Lookup(classicalPub.SchemeID())
	if err != nil {
		return err
	}
	pqKem := pqFactory()
	classicalKem := classicalFactory()

	// Encapsulate to each recipient.
	// ADAPTER CONTRACT (pinned verbatim): adapters return (ct, ss) —
	// composite unpacks in that order at each call site.
	pqCt, ssPq, err := pqKem.Encapsulate(pqPub, rand)
	if err != nil {
		return err
	}
	classicalCt, ssClassical, err := classicalKem.Encapsulate(classicalPub, rand)
	if err != nil {
		return err
	}

	// HKDF combiner (verbatim per plan Metis B1):
	//   PRK1 = HKDF-Extract(ssPq, salt=nil)
	//   kek1 = HKDF-Expand(prk1, infoPq, 32)
	//   IKM  = kek1 || ssClassical     (with defensive copy of kek1)
	//   PRK2 = HKDF-Extract(IKM, salt=nil)
	//   kekFinal = HKDF-Expand(prk2, infoComposite, 32)
	kekFinal, err := deriveCompositeKEK(ssPq, ssClassical)
	if err != nil {
		return err
	}

	// Wrap the CEK via AES-256-GCM with AAD = version u16 BE = {0x00, 0x02}.
	wrapNonce := make([]byte, wrapNonceLen)
	if _, err := io.ReadFull(rand, wrapNonce); err != nil {
		return err
	}
	kekBlock, err := aes.NewCipher(kekFinal)
	if err != nil {
		return err
	}
	wrapGcm, err := cipher.NewGCM(kekBlock)
	if err != nil {
		return err
	}
	wrappedCek := wrapGcm.Seal(nil, wrapNonce, cek, []byte{0x00, 0x02})

	// firstPayloadNonce seeds the per-chunk nonce stream.
	firstPayloadNonce := make([]byte, firstPayloadNonceLen)
	if _, err := io.ReadFull(rand, firstPayloadNonce); err != nil {
		return err
	}

	// Emit the artifact header.
	if err := writeHeader(
		sink,
		pqPub, classicalPub,
		pqCt, classicalCt,
		wrapNonce, wrappedCek, firstPayloadNonce,
	); err != nil {
		return err
	}

	// Encrypt the payload into chunks.
	payloadBlock, err := aes.NewCipher(cek)
	if err != nil {
		return err
	}
	payloadGcm, err := cipher.NewGCM(payloadBlock)
	if err != nil {
		return err
	}
	return encryptChunks(plaintext, sink, payloadGcm, firstPayloadNonce)
}

// writeHeader serializes the v2 artifact header.
//
// Layout (see package doc):
//
//	magic(4) + version(2) + flags(4) + nRecipients(1)
//	  + per-recipient: schemeID(2) + keyID(8) + ctLen(4) + ciphertext
//	  + wrapNonce(12) + wrappedCEK(48) + firstPayloadNonce(12)
func writeHeader(
	w io.Writer,
	pqPub, classicalPub crypto.RecipientPub,
	pqCt, classicalCt,
	wrapNonce, wrappedCek, firstPayloadNonce []byte,
) error {
	var buf bytes.Buffer
	var scratchBytes [4]byte

	binary.BigEndian.PutUint32(scratchBytes[:], magic)
	buf.Write(scratchBytes[:]) // [0:4] magic

	binary.BigEndian.PutUint16(scratchBytes[:2], version)
	buf.Write(scratchBytes[:2]) // [4:6] version

	binary.BigEndian.PutUint32(scratchBytes[:], flags)
	buf.Write(scratchBytes[:]) // [6:10] flags

	// [10] nRecipients — composite v2 always carries exactly two slots.
	buf.WriteByte(byte(maxRecipients))

	// Slot 0 (PQ).
	binary.BigEndian.PutUint16(scratchBytes[:2], pqPub.SchemeID())
	buf.Write(scratchBytes[:2])
	if len(pqPub.KeyID()) != 8 {
		return ErrMalformedHeader
	}
	buf.Write(pqPub.KeyID())
	binary.BigEndian.PutUint32(scratchBytes[:], uint32(len(pqCt)))
	buf.Write(scratchBytes[:])
	buf.Write(pqCt)

	// Slot 1 (classical).
	binary.BigEndian.PutUint16(scratchBytes[:2], classicalPub.SchemeID())
	buf.Write(scratchBytes[:2])
	if len(classicalPub.KeyID()) != 8 {
		return ErrMalformedHeader
	}
	buf.Write(classicalPub.KeyID())
	binary.BigEndian.PutUint32(scratchBytes[:], uint32(len(classicalCt)))
	buf.Write(scratchBytes[:])
	buf.Write(classicalCt)

	// wrapNonce + wrappedCEK + firstPayloadNonce.
	buf.Write(wrapNonce)
	buf.Write(wrappedCek)
	buf.Write(firstPayloadNonce)

	_, err := w.Write(buf.Bytes())
	return err
}

// encryptChunks encrypts plaintext into 64 KiB AES-256-GCM chunks under CEK
// and streams them to w. A zero-length final marker chunk (flags&0x01==1)
// is ALWAYS emitted on EOF regardless of the previous chunk's fullness.
func encryptChunks(
	plaintext io.Reader,
	w io.Writer,
	gcm cipher.AEAD,
	firstPayloadNonce []byte,
) error {
	chunkNonce := make([]byte, gcm.NonceSize())
	copy(chunkNonce, firstPayloadNonce)

	buf := make([]byte, chunkSize)
	var lenB [4]byte

	for {
		readN, readErr := io.ReadFull(plaintext, buf)
		hasData := readN > 0
		eof := readErr == io.EOF || readErr == io.ErrUnexpectedEOF

		if hasData {
			// Body chunk, flags = 0x00.
			ciphertext := gcm.Seal(nil, chunkNonce, buf[:readN], nil)
			binary.BigEndian.PutUint32(lenB[:], uint32(len(ciphertext)))
			if _, err := w.Write(lenB[:]); err != nil {
				return err
			}
			if _, err := w.Write([]byte{0x00}); err != nil {
				return err
			}
			if _, err := w.Write(ciphertext); err != nil {
				return err
			}
			// Increment counter for the next chunk; the top 4 bytes of the
			// nonce ([0:4]) are untouched.
			if err := incrementCounter(chunkNonce); err != nil {
				return err
			}
		}

		if eof {
			// Final marker chunk — zero-length plaintext, flags = 0x01,
			// ciphertext is just the 16-byte GCM tag. ALWAYS emitted.
			ciphertext := gcm.Seal(nil, chunkNonce, nil, nil)
			binary.BigEndian.PutUint32(lenB[:], uint32(len(ciphertext)))
			if _, err := w.Write(lenB[:]); err != nil {
				return err
			}
			if _, err := w.Write([]byte{flagFinal}); err != nil {
				return err
			}
			if _, err := w.Write(ciphertext); err != nil {
				return err
			}
			return nil
		}
		if readErr != nil {
			return readErr
		}
	}
}

// incrementCounter mutates chunkNonce in place: reads the 64-bit big-endian
// counter at chunkNonce[4:12], adds one, rejects wrap, writes back.
// chunkNonce[0:4] (the random base) is preserved.
func incrementCounter(chunkNonce []byte) error {
	counter := binary.BigEndian.Uint64(chunkNonce[4:12])
	newCounter := counter + 1
	if newCounter <= counter {
		return ErrNonceCounterWrapped
	}
	binary.BigEndian.PutUint64(chunkNonce[4:12], newCounter)
	return nil
}

// decryptor is the composite Decryptor implementation. Slot routing is
// positional; keyID equality is enforced as a fast wrong-key reject before
// any AEAD operation.
type decryptor struct {
	registry crypto.Registry
}

// NewDecryptor returns a composite Decryptor that resolves KEM schemes via
// the provided Registry.
func NewDecryptor(
	registry crypto.Registry,
) crypto.Decryptor {
	return &decryptor{
		registry: registry,
	}
}

// Decrypt parses the v2 artifact from src, decapsulates per slot using the
// supplied private keys (positional: slot 0 ← privs[0], slot 1 ← privs[1]),
// re-derives the composite KEK, unwraps the CEK, and streams decrypted
// plaintext chunks to plaintext.
func (d *decryptor) Decrypt(
	src io.Reader,
	privs []crypto.RecipientPriv,
	plaintext io.Writer,
) error {
	// Prefix: magic(4) + version(2) + flags(4) + nRecipients(1) = 11 bytes.
	const prefixLen = 11
	var prefix [prefixLen]byte
	if _, err := io.ReadFull(src, prefix[:]); err != nil {
		return ErrMalformedHeader
	}

	if binary.BigEndian.Uint32(prefix[0:4]) != magic {
		return ErrMalformedHeader
	}

	fileVersion := binary.BigEndian.Uint16(prefix[4:6])
	if fileVersion != version {
		// Stops BEFORE any flags/nRecipients validation, before any GCM
		// work, before any recipient allocation.
		return ErrUnsupportedVersion
	}

	if binary.BigEndian.Uint32(prefix[6:10]) != flags {
		return ErrMalformedHeader
	}

	nRecipients := int(prefix[10])
	if nRecipients < 1 || nRecipients > maxRecipients {
		return ErrMalformedHeader
	}
	if nRecipients != maxRecipients {
		// Composite v2 mandates exactly two slots.
		return ErrMalformedHeader
	}

	// Per-recipient: each slot is `(meta 14B) (ciphertext ctLen B)` INLINE —
	// slot0's ciphertext lives BETWEEN slot0's metadata and slot1's metadata
	// (the inline-ct layout; see package doc). So parse strictly per-slot:
	// read metadata → validate ctLen ≤ max → read ct → advance to next slot.
	// The ctLen ≤ maxRecipientCiphertextLen check must fire BEFORE allocating/reading
	// the per-slot ciphertext (test l: no OOM on malicious oversized value).
	const perSlotMetaLen = 14
	type slotMeta struct {
		schemeID uint16
		keyID    []byte
		ctLen    uint32
		ct       []byte
	}
	slots := make([]slotMeta, nRecipients)
	for slotIndex := 0; slotIndex < nRecipients; slotIndex++ {
		var meta [perSlotMetaLen]byte
		if _, err := io.ReadFull(src, meta[:]); err != nil {
			return ErrMalformedHeader
		}
		slot := &slots[slotIndex]
		slot.schemeID = binary.BigEndian.Uint16(meta[0:2])
		slot.keyID = append([]byte(nil), meta[2:10]...)
		slot.ctLen = binary.BigEndian.Uint32(meta[10:14])
		if int(slot.ctLen) > maxRecipientCiphertextLen {
			return ErrMalformedHeader
		}
		slot.ct = make([]byte, slot.ctLen)
		if slot.ctLen > 0 {
			if _, err := io.ReadFull(src, slot.ct); err != nil {
				return ErrMalformedHeader
			}
		}
	}

	// Trailing fixed region: wrapNonce(12) + wrappedCEK(48) + firstPayloadNonce(12) = 72B.
	var tail [wrapNonceLen + wrappedCekLen + firstPayloadNonceLen]byte
	if _, err := io.ReadFull(src, tail[:]); err != nil {
		return ErrMalformedHeader
	}
	wrapNonce := tail[:wrapNonceLen]
	wrappedCek := tail[wrapNonceLen : wrapNonceLen+wrappedCekLen]
	firstPayloadNonce := tail[wrapNonceLen+wrappedCekLen:]

	// Decapsulate per slot, routing privs positionally.
	if len(privs) < nRecipients {
		return ErrWrongKeys
	}
	sharedSecrets := make([][]byte, nRecipients)
	for slotIndex := 0; slotIndex < nRecipients; slotIndex++ {
		slot := slots[slotIndex]
		priv := privs[slotIndex]
		if priv == nil {
			return ErrWrongKeys
		}
		if priv.SchemeID() != slot.schemeID {
			return ErrWrongKeys
		}
		if !bytes.Equal(priv.KeyID(), slot.keyID) {
			return ErrWrongKeys
		}
		factory, err := d.registry.Lookup(slot.schemeID)
		if err != nil {
			return ErrWrongKeys
		}
		kem := factory()
		ss, err := kem.Decapsulate(priv, slot.ct)
		if err != nil {
			return ErrWrongKeys
		}
		sharedSecrets[slotIndex] = ss
	}

	kekFinal, err := deriveCompositeKEK(sharedSecrets[0], sharedSecrets[1])
	if err != nil {
		return err
	}

	// Unwrap CEK via AES-256-GCM. AAD = version u16 BE.
	kekBlock, err := aes.NewCipher(kekFinal)
	if err != nil {
		return err
	}
	wrapGcm, err := cipher.NewGCM(kekBlock)
	if err != nil {
		return err
	}
	cek, err := wrapGcm.Open(nil, wrapNonce, wrappedCek, []byte{0x00, 0x02})
	if err != nil {
		return ErrTamperingDetected
	}

	// Setup payload AEAD under CEK.
	payloadBlock, err := aes.NewCipher(cek)
	if err != nil {
		return err
	}
	payloadGcm, err := cipher.NewGCM(payloadBlock)
	if err != nil {
		return err
	}

	plaintextLimiter := &limitedWriter{writer: plaintext, remaining: maxPlaintextSize}
	return decryptChunks(src, plaintextLimiter, payloadGcm, firstPayloadNonce)
}

// limitedWriter wraps an io.Writer and rejects writes that would exceed a
// maximum byte budget.
type limitedWriter struct {
	writer    io.Writer
	remaining int64
}

func (lw *limitedWriter) Write(p []byte) (int, error) {
	if int64(len(p)) > lw.remaining {
		return 0, ErrPlaintextTooLarge
	}
	n, err := lw.writer.Write(p)
	lw.remaining -= int64(n)
	return n, err
}

// decryptChunks reads and decrypts chunk records until a final chunk
// (flags & flagFinal != 0) is observed.
func decryptChunks(
	src io.Reader,
	plaintext io.Writer,
	gcm cipher.AEAD,
	firstPayloadNonce []byte,
) error {
	chunkNonce := make([]byte, gcm.NonceSize())
	copy(chunkNonce, firstPayloadNonce)

	var lenB [4]byte
	var flagB [1]byte

	maxChunkCtLen := uint32(chunkSize + gcmTagLen)

	for {
		// Read chunk length u32 BE.
		_, err := io.ReadFull(src, lenB[:])
		if err == io.EOF || err == io.ErrUnexpectedEOF {
			// No final marker chunk observed — premature end of stream.
			return ErrUnexpectedEOF
		}
		if err != nil {
			return err
		}
		length := binary.BigEndian.Uint32(lenB[:])

		// Read flags u8.
		_, err = io.ReadFull(src, flagB[:])
		if err == io.EOF || err == io.ErrUnexpectedEOF {
			return ErrUnexpectedEOF
		}
		if err != nil {
			return err
		}
		flagsByte := flagB[0]
		isFinal := flagsByte&flagFinal != 0

		// Structural validation.
		if length == 0 && !isFinal {
			// Zero-length non-final chunk — infinite-loop DoS.
			return ErrMalformedChunk
		}
		if length > maxChunkCtLen {
			return ErrMalformedChunk
		}

		// Read ciphertext.
		ciphertext := make([]byte, length)
		if length > 0 {
			if _, err := io.ReadFull(src, ciphertext); err != nil {
				return ErrUnexpectedEOF
			}
		}

		// AEAD open.
		plaintextChunk, err := gcm.Open(nil, chunkNonce, ciphertext, nil)
		if err != nil {
			return ErrTamperingDetected
		}

		if len(plaintextChunk) > 0 {
			if _, err := plaintext.Write(plaintextChunk); err != nil {
				return err
			}
		}

		if isFinal {
			return nil
		}

		if err := incrementCounter(chunkNonce); err != nil {
			return err
		}
	}
}

// deriveCompositeKEK applies the verbatim HKDF combiner from the plan:
//
//	PRK1      = HKDF-Extract(IKM=ss_pq, salt=nil)
//	kek1      = HKDF-Expand(prk1, infoPq, 32)
//	IKM       = kek1 || ss_classical     // defensive copy of kek1
//	PRK2      = HKDF-Extract(IKM, salt=nil)
//	kekFinal  = HKDF-Expand(prk2, infoComposite, 32)
//
// Go 1.26 stdlib crypto/hkdf returns ([]byte, error) directly from Extract
// and Expand — no infinite io.Reader is involved, so neither io.ReadAll nor
// io.ReadFull is needed; the spirit of the plan's "do not use io.ReadAll on
// hkdf.Expand" guidance is preserved trivially.
func deriveCompositeKEK(
	ssPq, ssClassical []byte,
) ([]byte, error) {
	prk1, err := hkdf.Extract(sha256.New, ssPq, nil)
	if err != nil {
		return nil, err
	}
	kek1, err := hkdf.Expand(sha256.New, prk1, infoPq, kekLen)
	if err != nil {
		return nil, err
	}
	// Defensive copy of kek1 — append([]byte(nil), ...) avoids aliasing
	// kek1's backing array when concatenating ssClassical (Metis B1).
	ikmComposite := append(append([]byte(nil), kek1...), ssClassical...)
	prk2, err := hkdf.Extract(sha256.New, ikmComposite, nil)
	if err != nil {
		return nil, err
	}
	kekFinal, err := hkdf.Expand(sha256.New, prk2, infoComposite, kekLen)
	if err != nil {
		return nil, err
	}
	return kekFinal, nil
}