package group

import (
	"crypto"
	"crypto/ecdh"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	_ "crypto/sha256"
	_ "crypto/sha512"
	"crypto/subtle"
	"encoding/binary"
	"fmt"
	"io"
	"math/big"
	"slices"

	optP384 "github.com/cloudflare/circl/ecc/p384"
	"github.com/cloudflare/circl/expander"
	"github.com/cloudflare/circl/internal/conv"
	"golang.org/x/crypto/cryptobyte"
)

var (
	// P256 is the group generated by P-256 elliptic curve.
	P256 Group = wG{ellC: elliptic.P256, ecdhC: ecdh.P256, c: elliptic.P256(), order: orderP256[:]}
	// P384 is the group generated by P-384 elliptic curve.
	P384 Group = wG{ellC: elliptic.P384, ecdhC: ecdh.P384, c: optP384.P384(), order: orderP384[:]}
	// P521 is the group generated by P-521 elliptic curve.
	P521 Group = wG{ellC: elliptic.P521, ecdhC: ecdh.P521, c: elliptic.P521(), order: orderP521[:]}
)

type wG struct {
	c     elliptic.Curve
	ellC  func() elliptic.Curve
	ecdhC func() ecdh.Curve
	order []byte
}

func (g wG) String() string      { return g.c.Params().Name }
func (g wG) NewElement() Element { return g.zeroElement() }
func (g wG) NewScalar() Scalar   { return g.zeroScalar() }
func (g wG) Identity() Element   { return g.zeroElement() }
func (g wG) byteSize() int       { return (g.c.Params().BitSize + 7) / 8 }
func (g wG) zeroScalar() *wScl   { return &wScl{g, make([]byte, g.byteSize())} }
func (g wG) zeroElement() *wElt  { return &wElt{new(big.Int), new(big.Int), g} }
func (g wG) Generator() Element  { return &wElt{g.c.Params().Gx, g.c.Params().Gy, g} }
func (g wG) RandomElement(rd io.Reader) Element {
	b := make([]byte, g.byteSize())
	if n, err := io.ReadFull(rd, b); err != nil || n != len(b) {
		panic(err)
	}
	return g.HashToElement(b, nil)
}

func (g wG) RandomScalar(rd io.Reader) Scalar {
	r, err := rand.Int(rd, g.c.Params().N)
	if err != nil {
		panic(err)
	}

	return &wScl{g, r.FillBytes(make([]byte, g.byteSize()))}
}

func (g wG) RandomNonZeroScalar(rd io.Reader) Scalar {
	zero := g.zeroScalar()
	for {
		s := g.RandomScalar(rd)
		if !s.IsEqual(zero) {
			return s
		}
	}
}

func (g wG) cvtElt(e Element) *wElt {
	if e == nil {
		return g.zeroElement()
	}
	ee, ok := e.(*wElt)
	if !ok || g.c.Params().BitSize != ee.c.Params().BitSize {
		panic(ErrType)
	}
	return ee
}

func (g wG) cvtScl(s Scalar) *wScl {
	if s == nil {
		return g.zeroScalar()
	}
	ss, ok := s.(*wScl)
	if !ok || g.c.Params().BitSize != ss.c.Params().BitSize {
		panic(ErrType)
	}
	return ss
}

func (g wG) Params() *Params {
	fieldLen := uint(g.byteSize())
	return &Params{
		ElementLength:           1 + 2*fieldLen,
		CompressedElementLength: 1 + fieldLen,
		ScalarLength:            fieldLen,
	}
}

func (g wG) HashToElementNonUniform(b, dst []byte) Element {
	var u [1]big.Int
	mapping, h, L := g.mapToCurveParams()
	xmd := expander.NewExpanderMD(h, dst)
	HashToField(u[:], b, xmd, g.c.Params().P, L)
	return mapping(&u[0])
}

func (g wG) HashToElement(b, dst []byte) Element {
	var u [2]big.Int
	mapping, h, L := g.mapToCurveParams()
	xmd := expander.NewExpanderMD(h, dst)
	HashToField(u[:], b, xmd, g.c.Params().P, L)
	Q0 := mapping(&u[0])
	Q1 := mapping(&u[1])
	return Q0.Add(Q0, Q1)
}

func (g wG) HashToScalar(b, dst []byte) Scalar {
	var u [1]big.Int
	_, h, L := g.mapToCurveParams()
	xmd := expander.NewExpanderMD(h, dst)
	HashToField(u[:], b, xmd, g.c.Params().N, L)
	return g.zeroScalar().SetBigInt(&u[0])
}

type wElt struct {
	x, y *big.Int
	wG
}

func (e *wElt) Group() Group     { return e.wG }
func (e *wElt) String() string   { return fmt.Sprintf("x: 0x%v\ny: 0x%v", e.x.Text(16), e.y.Text(16)) }
func (e *wElt) IsIdentity() bool { return e.x.Sign() == 0 && e.y.Sign() == 0 }
func (e *wElt) IsEqual(o Element) bool {
	oo := e.cvtElt(o)
	return e.x.Cmp(oo.x) == 0 && e.y.Cmp(oo.y) == 0
}

func (e *wElt) Set(a Element) Element {
	aa := e.cvtElt(a)
	e.x.Set(aa.x)
	e.y.Set(aa.y)
	return e
}

func (e *wElt) Copy() Element { return e.wG.zeroElement().Set(e) }

func (e *wElt) CMov(v int, a Element) Element {
	if !(v == 0 || v == 1) {
		panic(ErrSelector)
	}
	aa := e.cvtElt(a)
	l := e.wG.byteSize()
	bufE := make([]byte, l)
	bufA := make([]byte, l)
	e.x.FillBytes(bufE)
	aa.x.FillBytes(bufA)
	subtle.ConstantTimeCopy(v, bufE, bufA)
	e.x.SetBytes(bufE)

	e.y.FillBytes(bufE)
	aa.y.FillBytes(bufA)
	subtle.ConstantTimeCopy(v, bufE, bufA)
	e.y.SetBytes(bufE)

	return e
}

func (e *wElt) CSelect(v int, a Element, b Element) Element {
	if !(v == 0 || v == 1) {
		panic(ErrSelector)
	}
	aa, bb := e.cvtElt(a), e.cvtElt(b)
	l := e.wG.byteSize()
	bufE := make([]byte, l)
	bufA := make([]byte, l)
	bufB := make([]byte, l)

	aa.x.FillBytes(bufA)
	bb.x.FillBytes(bufB)
	subtle.ConstantTimeCopy(v, bufE, bufA)
	subtle.ConstantTimeCopy(1-v, bufE, bufB)
	e.x.SetBytes(bufE)

	aa.y.FillBytes(bufA)
	bb.y.FillBytes(bufB)
	subtle.ConstantTimeCopy(v, bufE, bufA)
	subtle.ConstantTimeCopy(1-v, bufE, bufB)
	e.y.SetBytes(bufE)

	return e
}

func (e *wElt) Add(a, b Element) Element {
	aa, bb := e.cvtElt(a), e.cvtElt(b)
	e.x, e.y = e.c.Add(aa.x, aa.y, bb.x, bb.y)
	return e
}

func (e *wElt) Dbl(a Element) Element {
	aa := e.cvtElt(a)
	e.x, e.y = e.c.Double(aa.x, aa.y)
	return e
}

func (e *wElt) Neg(a Element) Element {
	aa := e.cvtElt(a)
	e.x.Set(aa.x)
	e.y.Neg(aa.y).Mod(e.y, e.c.Params().P)
	return e
}

func (e *wElt) Mul(a Element, s Scalar) Element {
	aa, ss := e.cvtElt(a), e.cvtScl(s)
	e.x, e.y = e.c.ScalarMult(aa.x, aa.y, ss.k)
	return e
}

func (e *wElt) MulGen(s Scalar) Element {
	ss := e.cvtScl(s)
	e.x, e.y = e.c.ScalarBaseMult(ss.k)
	return e
}

func (e *wElt) MarshalBinary() ([]byte, error) {
	if e.IsIdentity() {
		return []byte{0x0}, nil
	}

	e.x.Mod(e.x, e.c.Params().P)
	e.y.Mod(e.y, e.c.Params().P)
	pk, err := (&ecdsa.PublicKey{Curve: e.wG.ellC(), X: e.x, Y: e.y}).ECDH()
	if err != nil {
		return nil, err
	}

	return pk.Bytes(), nil
}

func (e *wElt) MarshalBinaryCompress() ([]byte, error) {
	if e.IsIdentity() {
		return []byte{0x0}, nil
	}
	e.x.Mod(e.x, e.c.Params().P)
	e.y.Mod(e.y, e.c.Params().P)
	return elliptic.MarshalCompressed(e.wG.c, e.x, e.y), nil
}

func (e *wElt) UnmarshalBinary(b []byte) error {
	byteLen := e.wG.byteSize()
	l := len(b)
	switch {
	case l == 1 && b[0] == 0x00: // point at infinity
		e.x.SetInt64(0)
		e.y.SetInt64(0)
	case l == 1+byteLen && (b[0] == 0x02 || b[0] == 0x03): // compressed
		x, y := elliptic.UnmarshalCompressed(e.wG.c, b)
		if x == nil {
			return ErrUnmarshal
		}
		e.x, e.y = x, y
	case l == 1+2*byteLen && b[0] == 0x04: // uncompressed
		_, err := e.wG.ecdhC().NewPublicKey(b)
		if err != nil {
			return ErrUnmarshal
		}

		e.x.SetBytes(b[1 : 1+byteLen])
		e.y.SetBytes(b[1+byteLen : 1+2*byteLen])
	default:
		return ErrUnmarshal
	}
	return nil
}

type wScl struct {
	wG
	k []byte
}

func (s *wScl) Group() Group   { return s.wG }
func (s *wScl) String() string { return fmt.Sprintf("0x%x", s.k) }
func (s *wScl) SetUint64(n uint64) Scalar {
	l := s.wG.byteSize()
	s.k = slices.Grow(s.k, l)
	clear(s.k)
	binary.BigEndian.PutUint64(s.k[l-8:], n)
	return s
}

func (s *wScl) SetBigInt(x *big.Int) Scalar {
	s.k = slices.Grow(s.k, s.wG.byteSize())
	x.Mod(x, s.c.Params().N).FillBytes(s.k)
	return s
}

func (s *wScl) IsZero() bool {
	return subtle.ConstantTimeCompare(s.k, make([]byte, s.wG.byteSize())) == 1
}

func (s *wScl) IsEqual(a Scalar) bool {
	aa := s.cvtScl(a)
	return subtle.ConstantTimeCompare(s.k, aa.k) == 1
}

func (s *wScl) Set(a Scalar) Scalar {
	aa := s.cvtScl(a)
	s.k = slices.Grow(s.k, s.wG.byteSize())
	copy(s.k, aa.k)
	return s
}

func (s *wScl) Copy() Scalar { return s.wG.zeroScalar().Set(s) }

func (s *wScl) CMov(v int, a Scalar) Scalar {
	if !(v == 0 || v == 1) {
		panic(ErrSelector)
	}
	aa := s.cvtScl(a)
	subtle.ConstantTimeCopy(v, s.k, aa.k)
	return s
}

func (s *wScl) CSelect(v int, a Scalar, b Scalar) Scalar {
	if !(v == 0 || v == 1) {
		panic(ErrSelector)
	}
	aa, bb := s.cvtScl(a), s.cvtScl(b)
	subtle.ConstantTimeCopy(v, s.k, aa.k)
	subtle.ConstantTimeCopy(1-v, s.k, bb.k)
	return s
}

func (s *wScl) Add(a, b Scalar) Scalar {
	aa, bb := s.cvtScl(a), s.cvtScl(b)
	r := new(big.Int)
	r.SetBytes(aa.k).Add(r, new(big.Int).SetBytes(bb.k))
	return s.SetBigInt(r)
}

func (s *wScl) Sub(a, b Scalar) Scalar {
	aa, bb := s.cvtScl(a), s.cvtScl(b)
	r := new(big.Int)
	r.SetBytes(aa.k).Sub(r, new(big.Int).SetBytes(bb.k))
	return s.SetBigInt(r)
}

func (s *wScl) Mul(a, b Scalar) Scalar {
	aa, bb := s.cvtScl(a), s.cvtScl(b)
	r := new(big.Int)
	r.SetBytes(aa.k).Mul(r, new(big.Int).SetBytes(bb.k))
	return s.SetBigInt(r)
}

func (s *wScl) Neg(a Scalar) Scalar {
	aa := s.cvtScl(a)
	r := new(big.Int)
	r.SetBytes(aa.k).Neg(r)
	return s.SetBigInt(r)
}

func (s *wScl) Inv(a Scalar) Scalar {
	aa := s.cvtScl(a)
	r := new(big.Int)
	r.SetBytes(aa.k).ModInverse(r, s.c.Params().N)
	return s.SetBigInt(r)
}

func (s *wScl) MarshalBinary() (data []byte, err error) {
	return slices.Clone(s.k), nil
}

func (s *wScl) UnmarshalBinary(b []byte) error {
	return conv.UnmarshalBinary(s, b)
}

func (s *wScl) Marshal(b *cryptobyte.Builder) error {
	b.AddBytes(s.k)
	return nil
}

func (s *wScl) Unmarshal(str *cryptobyte.String) bool {
	*s = *s.wG.zeroScalar()
	return str.CopyBytes(s.k) && isLessThanBE(s.k, s.wG.order) == 1
}

// isLessThan returns 1 if 0 <= x < y, otherwise 0. Assumes that slices have the same length.
func isLessThanBE(x, y []byte) int {
	j := len(x) - 1
	for i := len(x) - 1; i >= 0; i-- {
		eq := subtle.ConstantTimeByteEq(x[i], y[i])
		j = subtle.ConstantTimeSelect(eq, j, i)
	}

	return 1 - subtle.ConstantTimeLessOrEq(int(y[j]), int(x[j]))
}

func (g wG) mapToCurveParams() (mapping func(u *big.Int) *wElt, h crypto.Hash, L uint) {
	var Z, C2 big.Int
	switch g.c.Params().BitSize {
	case 256:
		Z.SetInt64(-10)
		C2.SetString("0x78bc71a02d89ec07214623f6d0f955072c7cc05604a5a6e23ffbf67115fa5301", 0)
		h = crypto.SHA256
		L = 48
	case 384:
		Z.SetInt64(-12)
		C2.SetString("0x19877cc1041b7555743c0ae2e3a3e61fb2aaa2e0e87ea557a563d8b598a0940d0a697a9e0b9e92cfaa314f583c9d066", 0)
		h = crypto.SHA384
		L = 72
	case 521:
		Z.SetInt64(-4)
		C2.SetInt64(8)
		h = crypto.SHA512
		L = 98
	default:
		panic("curve not supported")
	}
	return func(u *big.Int) *wElt { return g.sswu3mod4Map(u, &Z, &C2) }, h, L
}

func (g wG) sswu3mod4Map(u *big.Int, Z, C2 *big.Int) *wElt {
	tv1 := new(big.Int)
	tv2 := new(big.Int)
	tv3 := new(big.Int)
	tv4 := new(big.Int)
	xn := new(big.Int)
	xd := new(big.Int)
	x1n := new(big.Int)
	x2n := new(big.Int)
	gx1 := new(big.Int)
	gxd := new(big.Int)
	y1 := new(big.Int)
	y2 := new(big.Int)
	x := new(big.Int)
	y := new(big.Int)

	A := big.NewInt(-3)
	B := g.c.Params().B
	p := g.c.Params().P
	c1 := new(big.Int)
	c1.Sub(p, big.NewInt(3)).Rsh(c1, 2) // 1.  c1 = (q - 3) / 4

	add := func(c, a, b *big.Int) { c.Add(a, b).Mod(c, p) }
	mul := func(c, a, b *big.Int) { c.Mul(a, b).Mod(c, p) }
	sqr := func(c, a *big.Int) { c.Mul(a, a).Mod(c, p) }
	exp := func(c, a, b *big.Int) { c.Exp(a, b, p) }
	sgn := func(a *big.Int) uint { a.Mod(a, p); return a.Bit(0) }
	cmv := func(c, a, b *big.Int, k bool) {
		if k {
			c.Set(b)
		} else {
			c.Set(a)
		}
	}

	sqr(tv1, u)                 // 1.  tv1 = u^2
	mul(tv3, Z, tv1)            // 2.  tv3 = Z * tv1
	sqr(tv2, tv3)               // 3.  tv2 = tv3^2
	add(xd, tv2, tv3)           // 4.   xd = tv2 + tv3
	add(x1n, xd, big.NewInt(1)) // 5.  x1n = xd + 1
	mul(x1n, x1n, B)            // 6.  x1n = x1n * B
	tv4.Neg(A)                  //
	mul(xd, tv4, xd)            // 7.   xd = -A * xd
	e1 := xd.Sign() == 0        // 8.   e1 = xd == 0
	mul(tv4, Z, A)              //
	cmv(xd, xd, tv4, e1)        // 9.   xd = CMOV(xd, Z * A, e1)
	sqr(tv2, xd)                // 10. tv2 = xd^2
	mul(gxd, tv2, xd)           // 11. gxd = tv2 * xd
	mul(tv2, A, tv2)            // 12. tv2 = A * tv2
	sqr(gx1, x1n)               // 13. gx1 = x1n^2
	add(gx1, gx1, tv2)          // 14. gx1 = gx1 + tv2
	mul(gx1, gx1, x1n)          // 15. gx1 = gx1 * x1n
	mul(tv2, B, gxd)            // 16. tv2 = B * gxd
	add(gx1, gx1, tv2)          // 17. gx1 = gx1 + tv2
	sqr(tv4, gxd)               // 18. tv4 = gxd^2
	mul(tv2, gx1, gxd)          // 19. tv2 = gx1 * gxd
	mul(tv4, tv4, tv2)          // 20. tv4 = tv4 * tv2
	exp(y1, tv4, c1)            // 21.  y1 = tv4^c1
	mul(y1, y1, tv2)            // 22.  y1 = y1 * tv2
	mul(x2n, tv3, x1n)          // 23. x2n = tv3 * x1n
	mul(y2, y1, C2)             // 24.  y2 = y1 * c2
	mul(y2, y2, tv1)            // 25.  y2 = y2 * tv1
	mul(y2, y2, u)              // 26.  y2 = y2 * u
	sqr(tv2, y1)                // 27. tv2 = y1^2
	mul(tv2, tv2, gxd)          // 28. tv2 = tv2 * gxd
	e2 := tv2.Cmp(gx1) == 0     // 29.  e2 = tv2 == gx1
	cmv(xn, x2n, x1n, e2)       // 30.  xn = CMOV(x2n, x1n, e2)
	cmv(y, y2, y1, e2)          // 31.   y = CMOV(y2, y1, e2)
	e3 := sgn(u) == sgn(y)      // 32.  e3 = sgn0(u) == sgn0(y)
	tv1.Neg(y)                  //
	cmv(y, tv1, y, e3)          // 33.   y = CMOV(-y, y, e3)
	tv1.ModInverse(xd, p)       //
	mul(x, xn, tv1)             // 34. return (xn, xd, y, 1)
	y.Mod(y, p)
	return &wElt{x, y, g}
}

var (
	// Order of the P256 group.
	orderP256 = [...]byte{
		0xff, 0xff, 0xff, 0xff, 0x00, 0x00, 0x00, 0x00,
		0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
		0xbc, 0xe6, 0xfa, 0xad, 0xa7, 0x17, 0x9e, 0x84,
		0xf3, 0xb9, 0xca, 0xc2, 0xfc, 0x63, 0x25, 0x51,
	}
	// Order of the P384 group.
	orderP384 = [...]byte{
		0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
		0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
		0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
		0xc7, 0x63, 0x4d, 0x81, 0xf4, 0x37, 0x2d, 0xdf,
		0x58, 0x1a, 0x0d, 0xb2, 0x48, 0xb0, 0xa7, 0x7a,
		0xec, 0xec, 0x19, 0x6a, 0xcc, 0xc5, 0x29, 0x73,
	}
	// Order of the P521 group.
	orderP521 = [...]byte{
		0x01, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
		0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
		0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
		0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
		0xff, 0xfa, 0x51, 0x86, 0x87, 0x83, 0xbf, 0x2f,
		0x96, 0x6b, 0x7f, 0xcc, 0x01, 0x48, 0xf7, 0x09,
		0xa5, 0xd0, 0x3b, 0xb5, 0xc9, 0xb8, 0x89, 0x9c,
		0x47, 0xae, 0xbb, 0x6f, 0xb7, 0x1e, 0x91, 0x38,
		0x64, 0x09,
	}
)
