// Anti-malleability defenses for homomorphic ciphertexts.
package ring

import (
	"crypto/rand"
	"io"

	"crypto/sha3"
)

func NoiseFlood(pk *KEMPublicKey, ct *HECiphertext, floodBits int32) (result *HECiphertext) {
	return NoiseFloodFrom(pk, ct, floodBits, rand.Reader)
}

func NoiseFloodFrom(pk *KEMPublicKey, ct *HECiphertext, floodBits int32, rng io.Reader) (result *HECiphertext) {
	p := pk.P.Ring
	q := p.Q

	floodU := New(p)
	floodV := New(p)

	bound := uint32(1) << uint32(floodBits)

	var buf [4]byte
	for i := int32(0); i < p.N; i++ {
		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 & ^int64(1)
		if noise >= 0 {
			floodU.Coeffs[i] = uint32(noise) % q
		} else {
			floodU.Coeffs[i] = q - (uint32(-noise) % q)
		}

		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 & ^int64(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,
	}
}

type SecureHEResult struct {
	Ciphertext *HECiphertext
	Tag        []byte
	Signature  *GPVSignature
}

func WrapResult(pk *KEMPublicKey, ct *HECiphertext, sessionID []byte, gpvSK *GPVSecretKey) (result *SecureHEResult) {
	return WrapResultFrom(pk, ct, sessionID, gpvSK, rand.Reader)
}

func WrapResultFrom(pk *KEMPublicKey, ct *HECiphertext, sessionID []byte, gpvSK *GPVSecretKey, rng io.Reader) (result *SecureHEResult) {
	ct = RerandomizeFrom(pk, ct, rng)

	floodBits := int32(8)
	ct = NoiseFloodFrom(pk, ct, floodBits, rng)

	tag := computeTag(ct, sessionID)

	var sig *GPVSignature
	if gpvSK != nil {
		sig = GPVSignFrom(gpvSK, tag, rng)
	}

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

func VerifyResult(result *SecureHEResult, sessionID []byte, gpvPK *GPVPublicKey) (ok bool) {
	expectedTag := computeTag(result.Ciphertext, sessionID)
	if int32(len(result.Tag)) != int32(len(expectedTag)) {
		return false
	}
	for i := int32(0); i < int32(len(result.Tag)); i++ {
		if result.Tag[i] != expectedTag[i] {
			return false
		}
	}

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

	return true
}

func computeTag(ct *HECiphertext, sessionID []byte) (tag []byte) {
	h := sha3.NewSHAKE256()
	h.Write([]byte("he-result-tag-v1"))
	h.Write(sessionID)

	uBytes := Serialize(ct.U)
	vBytes := Serialize(ct.V)
	h.Write(uBytes)
	h.Write(vBytes)

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

type SessionWrapper struct {
	SessionID []byte
	HEPK      *KEMPublicKey
	GPVPK     *GPVPublicKey
	GPVSK     *GPVSecretKey
}

func NewSession(hePK *KEMPublicKey, gpvPK *GPVPublicKey, gpvSK *GPVSecretKey) (sw *SessionWrapper) {
	sessionID := []byte{:32}
	rand.Read(sessionID)

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

func (sw *SessionWrapper) Wrap(ct *HECiphertext) (result *SecureHEResult) {
	return WrapResult(sw.HEPK, ct, sw.SessionID, sw.GPVSK)
}

func (sw *SessionWrapper) Verify(result *SecureHEResult) (ok bool) {
	return VerifyResult(result, sw.SessionID, sw.GPVPK)
}