package gnarlring

import (
	"crypto/rand"
	"io"
)

// LWEPublicKey is a Ring-LWE public key: (A, B = A·S + E).
type LWEPublicKey struct {
	A *Poly27 // uniform ring element (coefficient form)
	B *Poly27 // B = A·S + E (coefficient form)
}

// LWESecretKey is a Ring-LWE secret key.
type LWESecretKey struct {
	S  *Poly27 // secret (ternary or Gaussian)
	PK *LWEPublicKey
}

// LWECiphertext is a Ring-LWE encryption: (U, V).
type LWECiphertext struct {
	U *Poly27 // U = A·R + E1
	V *Poly27 // V = B·R + E2 + encode(m)
}

// LWEKeyGen generates a Ring-LWE key pair with ternary secret.
func LWEKeyGen() (*LWEPublicKey, *LWESecretKey) {
	return LWEKeyGenFrom(rand.Reader)
}

// LWEKeyGenFrom generates a key pair from the given randomness source.
func LWEKeyGenFrom(rng io.Reader) (*LWEPublicKey, *LWESecretKey) {
	if rng == nil {
		rng = rand.Reader
	}
	gs := NewGaussSamplerFrom(DefaultSigma(), rng)

	// A ← uniform in R_q.
	a := NewPoly27()
	for i := 0; i < N; i++ {
		var buf [2]byte
		io.ReadFull(rng, buf[:])
		a.Coeffs[i] = uint16(uint32(buf[0])<<8|uint32(buf[1])) % Q
	}

	// S ← ternary (coefficients in {0, 1, Q-1}).
	s := ternaryPoly(rng)

	// E ← small Gaussian.
	e := gs.SamplePoly()
	for i := range e.Coeffs {
		// Scale down to ternary-like noise.
		if e.Coeffs[i] > Q/2 {
			e.Coeffs[i] = Q - 1 // -1 mod Q
		} else if e.Coeffs[i] != 0 {
			e.Coeffs[i] = 1
		}
	}

	// B = A·S + E.
	as := Mul(a, s)
	b := Add(as, e)

	pk := &LWEPublicKey{A: a, B: b}
	sk := &LWESecretKey{S: s, PK: pk}
	return pk, sk
}

// LWEEncrypt encrypts a single bit using the public key.
// Returns (U, V) ciphertext.
func LWEEncrypt(pk *LWEPublicKey, bit int) *LWECiphertext {
	return LWEEncryptFrom(pk, bit, rand.Reader)
}

// LWEEncryptFrom encrypts with a given randomness source.
func LWEEncryptFrom(pk *LWEPublicKey, bit int, rng io.Reader) *LWECiphertext {
	if rng == nil {
		rng = rand.Reader
	}

	// R ← ternary.
	r := ternaryPoly(rng)

	// E1, E2 ← ternary noise.
	e1 := ternaryPoly(rng)
	e2 := ternaryPoly(rng)

	// U = A·R + E1.
	ar := Mul(pk.A, r)
	u := Add(ar, e1)

	// V = B·R + E2 + encode(bit).
	br := Mul(pk.B, r)
	enc := encodeBit(bit)
	v := Add(br, e2)
	v = Add(v, enc)

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

// LWEDecrypt decrypts a ciphertext using the secret key.
// Returns 0 or 1.
func LWEDecrypt(sk *LWESecretKey, ct *LWECiphertext) int {
	// V - S·U = B·R + E2 + encode(m) - S·(A·R + E1)
	//         = (A·S + E)·R + E2 + encode(m) - S·A·R - S·E1
	//         = A·S·R + E·R + E2 + encode(m) - S·A·R - S·E1
	//         = encode(m) + (E·R + E2 - S·E1)
	// Noise = E·R + E2 - S·E1

	su := Mul(sk.S, ct.U)
	noisy := Sub(ct.V, su)

	return decodeBit(noisy)
}

// LWEAdd homomorphically adds two ciphertexts.
func LWEAdd(ct1, ct2 *LWECiphertext) *LWECiphertext {
	return &LWECiphertext{
		U: Add(ct1.U, ct2.U),
		V: Add(ct1.V, ct2.V),
	}
}

func encodeBit(bit int) *Poly27 {
	p := NewPoly27()
	if bit == 1 {
		p.Coeffs[0] = uint16(Q / 2) // 135
	}
	return p
}

func decodeBit(p *Poly27) int {
	c := p.Coeffs[0]
	half := uint16(Q / 2)   // 135
	quarter := uint16(Q / 4) // 67
	if c > half {
		c = Q - c
	}
	if c > quarter {
		return 1
	}
	return 0
}

func ternaryPoly(rng io.Reader) *Poly27 {
	p := NewPoly27()
	var buf [1]byte
	for i := 0; i < N; i++ {
		io.ReadFull(rng, buf[:])
		switch buf[0] % 3 {
		case 0:
			p.Coeffs[i] = 0
		case 1:
			p.Coeffs[i] = 1
		case 2:
			p.Coeffs[i] = Q - 1
		}
	}
	return p
}

// Serialize for public key: 2 × PolyBytes = 62 bytes.
func (pk *LWEPublicKey) MarshalBinary() []byte {
	buf := make([]byte, 2*PolyBytes)
	copy(buf[:PolyBytes], pk.A.MarshalBinary())
	copy(buf[PolyBytes:], pk.B.MarshalBinary())
	return buf
}

func UnmarshalLWEPK(data []byte) (*LWEPublicKey, error) {
	if len(data) < 2*PolyBytes {
		return nil, errShortData
	}
	a, err := UnmarshalBinary(data[:PolyBytes])
	if err != nil {
		return nil, err
	}
	b, err := UnmarshalBinary(data[PolyBytes : 2*PolyBytes])
	if err != nil {
		return nil, err
	}
	return &LWEPublicKey{A: a, B: b}, nil
}

// Ciphertext: 2 × PolyBytes = 62 bytes.
func (ct *LWECiphertext) MarshalBinary() []byte {
	buf := make([]byte, 2*PolyBytes)
	copy(buf[:PolyBytes], ct.U.MarshalBinary())
	copy(buf[PolyBytes:], ct.V.MarshalBinary())
	return buf
}

func UnmarshalLWECT(data []byte) (*LWECiphertext, error) {
	if len(data) < 2*PolyBytes {
		return nil, errShortData
	}
	u, err := UnmarshalBinary(data[:PolyBytes])
	if err != nil {
		return nil, err
	}
	v, err := UnmarshalBinary(data[PolyBytes : 2*PolyBytes])
	if err != nil {
		return nil, err
	}
	return &LWECiphertext{U: u, V: v}, nil
}