// Anti-malleability defenses for homomorphic ciphertexts.
//
// Homomorphic encryption is inherently malleable: anyone can transform
// ciphertexts without knowing the plaintext. This is a feature (enables
// computation on encrypted data) but also a vulnerability (enables
// ciphertext manipulation attacks).
//
// Defense layers:
//
// 1. Rerandomization: add Enc(0) to break ciphertext linkability.
//    Already implemented in he.go. This prevents tracking a ciphertext
//    through a computation graph.
//
// 2. Noise flooding: inject large noise into computation results.
//    The "circuit noise" from homomorphic operations carries information
//    about intermediate values. Flooding with fresh noise drowns this out.
//    Cost: reduces remaining noise budget.
//
// 3. CCA2 session wrapping: wrap HE ciphertexts in a CCA2-secure
//    session using the FO-transformed KEM from kem.go. An adversary
//    who modifies the outer layer gets caught by the FO check.
//
// 4. Signed computation: attach a GPV signature to computation results,
//    proving the computation was performed by an authorized party.
//    Without the signing key, an adversary can't forge valid results.
//
// The combination: rerandomize after each operation, noise-flood before
// returning results, wrap in CCA2 for transport, sign to prove provenance.

package ring

import (
	"crypto/rand"
	"io"

	"golang.org/x/crypto/sha3"
)

// NoiseFlood adds a large amount of fresh noise to a ciphertext.
// This drowns out any information leaked through the noise pattern.
//
// The flood noise must be much larger than the circuit noise but still
// smaller than the decryption threshold (q/4 for BGV).
//
// floodBits controls the magnitude: flood noise ∈ [-2^floodBits, 2^floodBits].
func NoiseFlood(pk *KEMPublicKey, ct *HECiphertext, floodBits int) *HECiphertext {
	return NoiseFloodFrom(pk, ct, floodBits, rand.Reader)
}

// NoiseFloodFrom adds flood noise using the given randomness source.
func NoiseFloodFrom(pk *KEMPublicKey, ct *HECiphertext, floodBits int, rng io.Reader) *HECiphertext {
	p := pk.P.Ring
	q := p.Q

	// Generate flood noise polynomial with coefficients in [-2^floodBits, 2^floodBits].
	// Must be even (divisible by 2) to preserve the BGV parity structure.
	floodU := New(p)
	floodV := New(p)

	bound := uint32(1) << floodBits

	var buf [4]byte
	for i := range p.N {
		// u noise
		io.ReadFull(rng, buf[:])
		val := uint32(buf[0])<<16 | uint32(buf[1])<<8 | uint32(buf[2])
		val = val % (2 * bound)
		noise := int64(val) - int64(bound)
		// Make even (BGV structure).
		noise = noise & ^1
		if noise >= 0 {
			floodU.Coeffs[i] = uint32(noise) % q
		} else {
			floodU.Coeffs[i] = q - (uint32(-noise) % q)
		}

		// v noise
		io.ReadFull(rng, buf[:])
		val = uint32(buf[0])<<16 | uint32(buf[1])<<8 | uint32(buf[2])
		val = val % (2 * bound)
		noise = int64(val) - int64(bound)
		noise = noise & ^1
		if noise >= 0 {
			floodV.Coeffs[i] = uint32(noise) % q
		} else {
			floodV.Coeffs[i] = q - (uint32(-noise) % q)
		}
	}

	return &HECiphertext{
		U:             Add(ct.U, floodU),
		V:             Add(ct.V, floodV),
		NoiseEstimate: ct.NoiseEstimate + float64(bound),
		params:        ct.params,
	}
}

// SecureHEResult wraps a homomorphic computation result with all defenses.
type SecureHEResult struct {
	// Ciphertext is the rerandomized + noise-flooded result.
	Ciphertext *HECiphertext

	// Tag is a MAC-like binding: H(ciphertext || session_id).
	// Detects any modification during transport.
	Tag []byte

	// Signature is an optional GPV signature on the tag,
	// proving the computation was performed by a specific party.
	Signature *GPVSignature
}

// WrapResult applies the full anti-malleability stack to a computation result.
//
// Steps:
// 1. Rerandomize (fresh Enc(0))
// 2. Noise flood (drown circuit noise)
// 3. Compute binding tag
// 4. Optionally sign
func WrapResult(pk *KEMPublicKey, ct *HECiphertext, sessionID []byte, gpvSK *GPVSecretKey) *SecureHEResult {
	return WrapResultFrom(pk, ct, sessionID, gpvSK, rand.Reader)
}

// WrapResultFrom wraps with the given randomness source.
func WrapResultFrom(pk *KEMPublicKey, ct *HECiphertext, sessionID []byte, gpvSK *GPVSecretKey, rng io.Reader) *SecureHEResult {
	// Step 1: Rerandomize.
	ct = RerandomizeFrom(pk, ct, rng)

	// Step 2: Noise flood with moderate noise.
	// For HE64 (q ≈ 10M), we can afford ~10 bits of flood noise
	// after one multiplication (noise budget ≈ 22 bits, mul uses ~20).
	floodBits := 8
	ct = NoiseFloodFrom(pk, ct, floodBits, rng)

	// Step 3: Compute binding tag.
	tag := computeTag(ct, sessionID)

	// Step 4: Sign (optional).
	var sig *GPVSignature
	if gpvSK != nil {
		sig = GPVSignFrom(gpvSK, tag, rng)
	}

	return &SecureHEResult{
		Ciphertext: ct,
		Tag:        tag,
		Signature:  sig,
	}
}

// VerifyResult checks the anti-malleability protections.
//
// Checks:
// 1. Tag matches the ciphertext (integrity)
// 2. Signature verifies (authenticity, if present)
func VerifyResult(result *SecureHEResult, sessionID []byte, gpvPK *GPVPublicKey) bool {
	// Check tag.
	expectedTag := computeTag(result.Ciphertext, sessionID)
	if len(result.Tag) != len(expectedTag) {
		return false
	}
	for i := range result.Tag {
		if result.Tag[i] != expectedTag[i] {
			return false
		}
	}

	// Check signature (if present).
	if result.Signature != nil {
		if gpvPK == nil {
			return false
		}
		if !GPVVerify(gpvPK, result.Tag, result.Signature) {
			return false
		}
	}

	return true
}

// computeTag creates a binding tag from a ciphertext and session ID.
func computeTag(ct *HECiphertext, sessionID []byte) []byte {
	h := sha3.NewShake256()
	h.Write([]byte("he-result-tag-v1"))
	h.Write(sessionID)

	// Encode ciphertext components.
	uBytes := Serialize(ct.U)
	vBytes := Serialize(ct.V)
	h.Write(uBytes)
	h.Write(vBytes)

	tag := make([]byte, 32)
	h.Read(tag)
	return tag
}

// SessionWrapper provides CCA2-secure transport for HE ciphertexts.
// Uses the KEM to establish a session key, then wraps each result.
type SessionWrapper struct {
	// SessionID is the unique session identifier.
	SessionID []byte

	// HEPK is the homomorphic encryption public key.
	HEPK *KEMPublicKey

	// GPVPK is the optional signing public key.
	GPVPK *GPVPublicKey

	// GPVSK is the optional signing secret key (only for the sender).
	GPVSK *GPVSecretKey
}

// NewSession creates a new session for wrapping HE results.
func NewSession(hePK *KEMPublicKey, gpvPK *GPVPublicKey, gpvSK *GPVSecretKey) *SessionWrapper {
	sessionID := make([]byte, 32)
	rand.Read(sessionID)

	return &SessionWrapper{
		SessionID: sessionID,
		HEPK:      hePK,
		GPVPK:     gpvPK,
		GPVSK:     gpvSK,
	}
}

// Wrap applies anti-malleability defenses to a computation result.
func (sw *SessionWrapper) Wrap(ct *HECiphertext) *SecureHEResult {
	return WrapResult(sw.HEPK, ct, sw.SessionID, sw.GPVSK)
}

// Verify checks the anti-malleability protections on a received result.
func (sw *SessionWrapper) Verify(result *SecureHEResult) bool {
	return VerifyResult(result, sw.SessionID, sw.GPVPK)
}