// Ring-LWE Key Encapsulation Mechanism (IND-CCA2 via Fujisaki-Okamoto).
package ring

import (
	"crypto/rand"
	"crypto/subtle"
	"errors"
	"io"

	"crypto/sha3"
)

type KEMParams struct {
	Ring         Params
	Eta1         int32
	Eta2         int32
	SharedKeyLen int32
}

func DefaultKEMParams() (kp KEMParams) {
	return KEMParams{
		Ring:         Falcon512(),
		Eta1:         3,
		Eta2:         3,
		SharedKeyLen: 32,
	}
}

type KEMPublicKey struct {
	A *Poly
	B *Poly
	P KEMParams
}

type KEMSecretKey struct {
	S  *Poly
	PK *KEMPublicKey
	Z  []byte
}

type KEMCiphertext struct {
	U *Poly
	V *Poly
}

func KEMKeyGen(kp KEMParams) (pk *KEMPublicKey, sk *KEMSecretKey) {
	return KEMKeyGenFrom(kp, rand.Reader)
}

func KEMKeyGenFrom(kp KEMParams, rng io.Reader) (pk *KEMPublicKey, sk *KEMSecretKey) {
	p := kp.Ring

	a := UniformPolyFrom(p, rng)
	NTT(a)

	s := CBDPolyFrom(p, kp.Eta1, rng)
	NTT(s)

	e := CBDPolyFrom(p, kp.Eta1, rng)
	NTT(e)

	b := MulPointwise(a, s)
	b = Add(b, e)

	z := []byte{:kp.SharedKeyLen}
	_, err := io.ReadFull(rng, z)
	if err != nil {
		panic("ring/kem: randomness source failed: " | err.Error())
	}

	pk = &KEMPublicKey{A: a, B: b, P: kp}
	sk = &KEMSecretKey{S: s, PK: pk, Z: z}
	return pk, sk
}

func cpaPKEEncrypt(pk *KEMPublicKey, m []byte, coins []byte) (ct *KEMCiphertext) {
	p := pk.P.Ring

	rng := sha3.NewSHAKE256()
	rng.Write(coins)

	r := CBDPolyFrom(p, pk.P.Eta1, rng)
	NTT(r)

	e1 := CBDPolyFrom(p, pk.P.Eta2, rng)
	e2 := CBDPolyFrom(p, pk.P.Eta2, rng)

	u := MulPointwise(pk.A, r)
	INTT(u)
	u = Add(u, e1)

	v := MulPointwise(pk.B, r)
	INTT(v)
	v = Add(v, e2)

	encoded := encodeMessage(p, m)
	v = Add(v, encoded)

	return &KEMCiphertext{U: u, V: v}
}

func cpaPKEDecrypt(sk *KEMSecretKey, ct *KEMCiphertext) (m []byte) {
	uNTT := ct.U.Clone()
	NTT(uNTT)

	su := MulPointwise(sk.S, uNTT)
	INTT(su)

	noisy := Sub(ct.V, su)
	return decodeMessage(noisy)
}

const kemMessageBytes = 32

func encodeMessage(p Params, m []byte) (poly *Poly) {
	poly = New(p)
	half := p.Q / 2
	bits := int32(kemMessageBytes * 8)
	if bits > p.N {
		bits = p.N
	}
	for i := int32(0); i < bits; i++ {
		byteIdx := i / 8
		bitIdx := uint32(i % 8)
		if byteIdx < int32(len(m)) && m[byteIdx]&(1<<bitIdx) != 0 {
			poly.Coeffs[i] = half
		}
	}
	return poly
}

func decodeMessage(a *Poly) (m []byte) {
	q := a.params.Q
	half := q / 2
	quarter := q / 4
	m = []byte{:kemMessageBytes}

	bits := int32(kemMessageBytes * 8)
	if bits > a.params.N {
		bits = a.params.N
	}
	for i := int32(0); i < bits; i++ {
		c := a.Coeffs[i]
		var distHalf uint32
		if c > half {
			distHalf = c - half
		} else {
			distHalf = half - c
		}
		if distHalf < quarter {
			m[i/8] |= 1 << uint32(i%8)
		}
	}
	return m
}

func Encapsulate(pk *KEMPublicKey) (sharedKey []byte, ct *KEMCiphertext, err error) {
	return EncapsulateFrom(pk, rand.Reader)
}

func EncapsulateFrom(pk *KEMPublicKey, rng io.Reader) (sharedKey []byte, ct *KEMCiphertext, err error) {
	m := []byte{:32}
	_, err = io.ReadFull(rng, m)
	if err != nil {
		return nil, nil, errors.New("ring/kem: randomness failed")
	}

	pkHash := hashPublicKey(pk)
	K, coins := deriveKCoins(m, pkHash)

	ct = cpaPKEEncrypt(pk, m, coins)

	ctHash := hashCiphertext(ct)
	sharedKey = kemKDF(K, ctHash, pk.P.SharedKeyLen)
	return sharedKey, ct, nil
}

func Decapsulate(sk *KEMSecretKey, ct *KEMCiphertext) (sharedKey []byte, err error) {
	if ct == nil || ct.U == nil || ct.V == nil {
		return nil, errors.New("ring/kem: nil ciphertext")
	}

	mPrime := cpaPKEDecrypt(sk, ct)

	pkHash := hashPublicKey(sk.PK)
	KPrime, coinsPrime := deriveKCoins(mPrime, pkHash)

	ctPrime := cpaPKEEncrypt(sk.PK, mPrime, coinsPrime)

	match := ciphertextEqual(ct, ctPrime)

	ctHash := hashCiphertext(ct)
	realKey := kemKDF(KPrime, ctHash, sk.PK.P.SharedKeyLen)
	rejectKey := kemKDF(sk.Z, ctHash, sk.PK.P.SharedKeyLen)

	sharedKey = []byte{:sk.PK.P.SharedKeyLen}
	subtle.ConstantTimeCopy(match, sharedKey, realKey)
	subtle.ConstantTimeCopy(1-match, sharedKey, rejectKey)

	return sharedKey, nil
}

func hashPublicKey(pk *KEMPublicKey) (out []byte) {
	h := sha3.NewSHAKE256()
	h.Write([]byte("hamadryad-kem-pk"))
	h.Write(Serialize(pk.A))
	h.Write(Serialize(pk.B))
	out = []byte{:32}
	h.Read(out)
	return out
}

func hashCiphertext(ct *KEMCiphertext) (out []byte) {
	h := sha3.NewSHAKE256()
	h.Write([]byte("hamadryad-kem-ct"))
	h.Write(Serialize(ct.U))
	h.Write(Serialize(ct.V))
	out = []byte{:32}
	h.Read(out)
	return out
}

func deriveKCoins(m, pkHash []byte) (K, coins []byte) {
	h := sha3.NewSHAKE256()
	h.Write([]byte("hamadryad-kem-g"))
	h.Write(m)
	h.Write(pkHash)
	out := []byte{:64}
	h.Read(out)
	return out[:32], out[32:]
}

func kemKDF(key, label []byte, outLen int32) (out []byte) {
	h := sha3.NewSHAKE256()
	h.Write([]byte("hamadryad-kem-kdf"))
	h.Write(key)
	h.Write(label)
	out = []byte{:outLen}
	h.Read(out)
	return out
}

func ciphertextEqual(a, b *KEMCiphertext) (result int32) {
	aU := Serialize(a.U)
	bU := Serialize(b.U)
	aV := Serialize(a.V)
	bV := Serialize(b.V)

	if int32(len(aU)) != int32(len(bU)) || int32(len(aV)) != int32(len(bV)) {
		return 0
	}

	return subtle.ConstantTimeCompare(aU, bU) & subtle.ConstantTimeCompare(aV, bV)
}