package elgamal

import (
	"crypto/rand"
	"errors"
	"io"
	"math/big"

	"xdx.jelly/xgcl/grand"
)

var Reader = grand.Reader

type Params struct {
	isMont bool // Is Montgomery represent
	p      *big.Int
	q      *big.Int
	g      *big.Int
	pMont  *big.Int
	qMont  *big.Int
	gMont  *big.Int
}

func (p *Params) P() *big.Int {
	return p.p
}
func (p *Params) Q() *big.Int {
	return p.q
}
func (p *Params) G() *big.Int {
	return p.g
}

type PublicKey struct {
	y      big.Int
	Params *Params
}

type PrivateKey struct {
	x big.Int
	PublicKey
}
type Cipher struct {
	u big.Int
	v big.Int
}

func GenerateKey(rnd io.Reader, params *Params) (*PrivateKey, *PublicKey, error) {
	x, err := rand.Int(rnd, params.q)
	if err != nil {
		return nil, nil, err
	}
	y := new(big.Int).Exp(params.g, x, params.p)
	sk := &PrivateKey{
		x: *x,
		PublicKey: PublicKey{
			y:      *y,
			Params: P1024,
		},
	}
	pk := &sk.PublicKey
	return sk, pk, nil
}

// m = g^M mod p for some M
// rnd can be nil, io.Reader or []byte
func (pk *PublicKey) Encryption(m *big.Int, rnd any) (*Cipher, error) {
	var r *big.Int
	var err error
	p := pk.Params.p
	q := pk.Params.q
	g := pk.Params.g

	switch rnd := rnd.(type) {
	case nil:
		rnd = Reader
	case io.Reader:
		r, err = rand.Int(rnd, q)
		if err != nil {
			return nil, err
		}
	case []byte:
		r = new(big.Int).SetBytes(rnd)
		if r.Cmp(q) >= 0 {
			return nil, errors.New("random bytes too large")
		}
	default:
		return nil, errors.New("rnd must be io.Reader or []bytes")
	}

	c := &Cipher{}
	c.u.Exp(g, r, p)
	c.v.Exp(&pk.y, r, p)
	c.v.Mul(&c.v, m)
	return c, nil
}

func (sk *PrivateKey) Decryption(c *Cipher) (*big.Int, error) {
	p := sk.PublicKey.Params.p
	u := &c.u
	v := &c.v
	z := new(big.Int)
	z.Exp(u, &sk.x, p)
	z.ModInverse(z, p)
	z.Mul(z, v)
	z.Mod(z, p)
	return z, nil
}

func (c *Cipher) HomoMap(c1 *Cipher, c2 *Cipher) (*Cipher, error) {
	c.u.Mul(&c1.u, &c2.u)
	c.v.Mul(&c1.v, &c2.v)
	return c, nil
}