package gmath

import (
	"crypto/rand"
	"encoding/hex"
	"math/big"

	"xdx.jelly/xgcl/grand"
)

// Zp is a simple wrap of big.Int, to comput in integers mod p
type Zp struct {
	d        *big.Int
	p        *big.Int
	byteSize int
}

func (z *Zp) init(p *big.Int) *Zp {
	z.d = new(big.Int)
	z.p = p
	z.byteSize = len(z.p.Bytes())
	return z
}

// a to bytes slice, big endian, padding 0
// 0x123456 =>
// 00 00 00 00 ... 12 34 56 78 -- total nBytes
// note:
// if a.Bytes() are longer than nBytes, then
// only the nBytes of lower are return. i.e
// the return value is a mod 2^{8*nBytes}
func intToByte(a *big.Int, nBytes int) []byte {
	abs := a.Bytes()
	l := len(abs)
	s := make([]byte, nBytes)
	if nBytes > l {
		copy(s[nBytes-l:], abs)
	} else {
		copy(s, abs[l-nBytes:])
	}
	return s
}

// NewZp ...
func NewZp(p *big.Int) *Zp {
	return new(Zp).init(p)
}

// Clear set memory to 0
func (z *Zp) Clear() {
	abs := z.d.Bytes()
	for i := range abs {
		abs[i] = 0
	}
	z.SetInt(0)
}

// GetInteger return the d (0 < d < p)
func (z *Zp) GetInteger() *big.Int {
	return z.Normalize().d
}

// GetPrime return the p
func (z *Zp) GetPrime() *big.Int {
	return z.p
}

// Set z to x
func (z *Zp) Set(x *Zp) *Zp {
	z.d.Set(x.d)
	z.p = x.p
	return z
}

// SetInt ...
func (z *Zp) SetInt(n int) *Zp {
	z.d.SetInt64(int64(n))
	return z
}

// SetBigInt ...
func (z *Zp) SetBigInt(n *big.Int) *Zp {
	z.d.Set(n)
	return z
}

// SetString ...
func (z *Zp) SetString(s string, base int) (*Zp, bool) {
	_, err := z.d.SetString(s, base)
	return z, err
}

// SetBytes ...
func (z *Zp) SetBytes(buf []byte) *Zp {
	z.d.SetBytes(buf)
	return z
}

// Is0 ...
func (z *Zp) Is0() bool {
	z.Normalize()
	return IsBigInt0(z.d)
}

// Is1 ...
func (z *Zp) Is1() bool {
	z.Normalize()
	return IsBigInt1(z.d)
}

// Sign return the odd of z.d
func (z *Zp) Sign() uint {
	return z.Normalize().d.Bit(0)
}

// Equal ...
func (z *Zp) Equal(p *Zp) bool {
	return z.Normalize().d.Cmp(p.Normalize().d) == 0
}

// GetItem return the d, this is for a extension field,
// the d is an array, and return the d[i]
func (z *Zp) GetItem(i int) *big.Int {
	if i == 0 {
		return z.d
	}
	return nil
}

// Normalize put d to [0,p-1]
func (z *Zp) Normalize() *Zp {
	z.d.Mod(z.d, z.p)
	return z
}

// Bytes return the byte slice of z.d, with a const slice size with p.Bytes()
func (z *Zp) Bytes() []byte {
	z.Normalize()
	return intToByte(z.d, z.byteSize)
}

func (z *Zp) String() string {
	return hex.EncodeToString(z.Bytes())
}

// Add z = p+q
// make sure that z,p,q are in the same field, or the result is a shit
func (z *Zp) Add(p, q *Zp) *Zp {
	pp := p.d
	qq := q.d
	z.d.Add(pp, qq)
	return z
}

// AddInt z = p+n
func (z *Zp) AddInt(p *Zp, n int) *Zp {
	z.d.Add(p.d, big.NewInt(int64(n)))
	return z
}

// AddBigInt ...
func (z *Zp) AddBigInt(p *Zp, n *big.Int) *Zp {
	z.d.Add(p.d, n)
	return z
}

// Neg ...
func (z *Zp) Neg(p *Zp) *Zp {
	p.Normalize() // p maybe changed, but not matter
	if !p.Is0() {
		z.d.Sub(z.d, p.d)
	} else {
		z.Set(p)
	}
	return z
}

// Sub ...
func (z *Zp) Sub(p, q *Zp) *Zp {
	z.d.Sub(p.d, q.d)
	return z
}

// SubInt ...
func (z *Zp) SubInt(p *Zp, n int) *Zp {
	z.d.Sub(p.d, big.NewInt(int64(n)))
	return z
}

// SubBigInt ...
func (z *Zp) SubBigInt(p *Zp, n *big.Int) *Zp {
	z.d.Sub(p.d, n)
	return z
}

// Mul ...
func (z *Zp) Mul(p, q *Zp) *Zp {
	p.Normalize()
	q.Normalize()
	z.d.Mul(p.d, q.d)
	z.Normalize()
	return z
}

// MulInt ...
func (z *Zp) MulInt(p *Zp, n int) *Zp {
	z.d.Mul(p.d, big.NewInt(int64(n)))
	z.Normalize()
	return z
}

// MulBigInt ...
func (z *Zp) MulBigInt(p *Zp, n *big.Int) *Zp {
	z.d.Mul(p.d, n)
	z.Normalize()
	return z
}

// Inv z=p^(-1), if p == 0, panic
func (z *Zp) Inv(p *Zp) *Zp {
	if p.Is0() {
		panic("zp:Inv try to invert 0")
	}
	z.d.ModInverse(p.d, z.p)
	return z
}

// Random return a random number from entropy
func (z *Zp) Random() error {
	n, err := rand.Int(grand.Reader, z.p)
	if err != nil {
		return err
	}
	z.d = n
	return nil
}