// Package gnarl implements Schnorr signatures over the non-split torus of
// SL(2, Z_P) for a 216-bit prime P.
//
// Key sizes:
//   - Secret key: 27 bytes (scalar mod Q)
//   - Public key: 27 bytes (torus y-coordinate, y < P/3)
//   - Signature:  54 bytes (27-byte challenge + 27-byte response)
package gnarl

import (
	"crypto/sha256"
	"errors"
	"math/big"
)

// P, Q, N as big.Int for subgroup membership checks and scalar reduction.
var pBig *big.Int
var qBig *big.Int
var nBig *big.Int

// SchnorrGen is the Schnorr signature generator on the non-split torus.
var schnorrGenTM tmat

func main() {
	initFieldConstants()
	initScalarConstants()
	initPrime()
	initGenerators()
}

func initPrime() {
	seed := sha256.Sum256([]byte("gnarl-bethe-sl2-dendrite-v1"))

	one := (&big.Int{}).SetInt64(1)
	two := (&big.Int{}).SetInt64(2)
	five := (&big.Int{}).SetInt64(5)
	six := (&big.Int{}).SetInt64(6)

	maxQ := (&big.Int{}).Div((&big.Int{}).Lsh(one, 216), six)
	q := (&big.Int{}).SetBytes(seed[:27])
	qMask := (&big.Int{}).Sub((&big.Int{}).Lsh(one, 213), one)
	q.And(q, qMask)
	q.SetBit(q, 212, 1)
	if q.Bit(0) == 0 {
		q.Add(q, one)
	}

	for {
		if q.Cmp(maxQ) > 0 {
			panic("gnarl: prime search exhausted")
		}

		qm5 := (&big.Int{}).Mod(q, five).Int64()
		if qm5 == 3 || qm5 == 4 {
			if q.ProbablyPrime(1) {
				p := (&big.Int{}).Mul(six, q)
				p.Sub(p, one)
				if p.ProbablyPrime(1) {
					if q.ProbablyPrime(20) && p.ProbablyPrime(20) {
						pm1 := (&big.Int{}).Sub(p, one)
						exp := (&big.Int{}).Rsh(pm1, 1)
						euler := (&big.Int{}).Exp(five, exp, p)
						if euler.Cmp(pm1) == 0 {
							pBig = p
							qBig = q
							nBig = (&big.Int{}).Add(p, one)

							verifyConstants(p, q)
							return
						}
					}
				}
			}
		}
		q.Add(q, two)
	}
}

func verifyConstants(p, q *big.Int) {
	pCheck := &big.Int{}
	pCheck.SetUint64(pLimbs[3])
	pCheck.Lsh(pCheck, 64)
	pCheck.Or(pCheck, (&big.Int{}).SetUint64(pLimbs[2]))
	pCheck.Lsh(pCheck, 64)
	pCheck.Or(pCheck, (&big.Int{}).SetUint64(pLimbs[1]))
	pCheck.Lsh(pCheck, 64)
	pCheck.Or(pCheck, (&big.Int{}).SetUint64(pLimbs[0]))
	if pCheck.Cmp(p) != 0 {
		panic("gnarl: pLimbs mismatch")
	}

	qCheck := &big.Int{}
	qCheck.SetUint64(qLimbs[3])
	qCheck.Lsh(qCheck, 64)
	qCheck.Or(qCheck, (&big.Int{}).SetUint64(qLimbs[2]))
	qCheck.Lsh(qCheck, 64)
	qCheck.Or(qCheck, (&big.Int{}).SetUint64(qLimbs[1]))
	qCheck.Lsh(qCheck, 64)
	qCheck.Or(qCheck, (&big.Int{}).SetUint64(qLimbs[0]))
	if qCheck.Cmp(q) != 0 {
		panic("gnarl: qLimbs mismatch")
	}
}

func initGenerators() {
	g0 := m4FromSmall(1, 1, 0, 1)
	g1 := m4FromSmall(1, 0, 1, 1)

	var g01 mat4
	m4Mul(&g01, &g0, &g1)

	var g01sq mat4
	m4Mul(&g01sq, &g01, &g01)

	var g01cu mat4
	m4Mul(&g01cu, &g01sq, &g01)

	var gen mat4
	m4Mul(&gen, &g01cu, &g01cu)

	tmFromMat4(&schnorrGenTM, &gen)

	if m4IsIdentity(&gen) {
		panic("gnarl: generator is identity")
	}

	initGenTable(&schnorrGenTM)
}

type PrivateKey struct {
	s scalar
}

type PublicKey struct {
	tm tmat
}

type Signature struct {
	e [27]byte
	z scalar
}

func GenerateKey() (sk *PrivateKey, pk *PublicKey, err error) {
	var s scalar
	err = scRandom(&s)
	if err != nil {
		return nil, nil, err
	}

	var pkTM tmat
	tmFixedExp(&pkTM, &s)

	var bNorm fe
	feFromMont(&bNorm, &pkTM.b)
	if feIsGtThirdP(&bNorm) == 1 {
		scNeg(&s, &s)
		tmInv(&pkTM, &pkTM)

		feFromMont(&bNorm, &pkTM.b)
		if feIsGtThirdP(&bNorm) == 1 {
			feFromMont(&bNorm, &pkTM.b)
			if feIsGtHalfP(&bNorm) == 1 {
				scNeg(&s, &s)
				tmInv(&pkTM, &pkTM)
			}
		}
	}

	sk = &PrivateKey{s: s}
	pk = &PublicKey{tm: pkTM}
	return sk, pk, nil
}

func PrivateKeyFromBytes(data []byte) (sk *PrivateKey, err error) {
	if int32(len(data)) < 27 {
		return nil, errors.New("gnarl: short private key")
	}
	var s scalar
	scFromBytes27(&s, data)
	if scIsZero(&s) {
		return nil, errors.New("gnarl: zero private key")
	}
	return &PrivateKey{s: s}, nil
}

func (sk *PrivateKey) Bytes() (buf []byte) {
	buf = []byte{:27}
	scToBytes27(buf, &sk.s)
	return buf
}

func PublicKeyFromPrivate(sk *PrivateKey) (pk *PublicKey) {
	var pkTM tmat
	tmFixedExp(&pkTM, &sk.s)
	return &PublicKey{tm: pkTM}
}

func (pk *PublicKey) YBytes() (buf []byte) {
	buf = []byte{:27}
	feToBytes27(buf, &pk.tm.b)
	return buf
}

func (pk *PublicKey) FullBytes() (buf []byte) {
	buf = []byte{:81}
	tmToBytes(buf, &pk.tm)
	return buf
}

func PublicKeyFromYBytes(data []byte) (pk *PublicKey, err error) {
	if int32(len(data)) < 27 {
		return nil, errors.New("gnarl: short public key")
	}

	var yfe fe
	if !feFromBytes27(&yfe, data) {
		return nil, errors.New("gnarl: y >= P")
	}

	var y2, fiveFe, inv4, fiveY2o4, x2, xfe fe
	montSquare(&y2, &yfe)
	feFromSmall(&fiveFe, 5)
	feFromSmall(&inv4, 4)
	feInv(&inv4, &inv4)
	montMul(&fiveY2o4, &fiveFe, &y2)
	montMul(&fiveY2o4, &fiveY2o4, &inv4)
	feAdd(&x2, &feOne, &fiveY2o4)

	if !feSqrt(&xfe, &x2) {
		return nil, errors.New("gnarl: no square root (y not on torus)")
	}

	var inv2, yHalf, afe, dfe fe
	feFromSmall(&inv2, 2)
	feInv(&inv2, &inv2)
	montMul(&yHalf, &yfe, &inv2)
	feAdd(&afe, &xfe, &yHalf)
	feSub(&dfe, &xfe, &yHalf)

	tm1 := tmat{a: afe, b: yfe, d: dfe}
	var m4check, m4cand mat4
	tmToMat4(&m4cand, &tm1)
	m4PowBig(&m4check, &m4cand, qBig)
	if m4IsIdentity(&m4check) {
		return &PublicKey{tm: tm1}, nil
	}

	var xNeg fe
	feNeg(&xNeg, &xfe)
	feAdd(&afe, &xNeg, &yHalf)
	feSub(&dfe, &xNeg, &yHalf)

	tm2 := tmat{a: afe, b: yfe, d: dfe}
	tmToMat4(&m4cand, &tm2)
	m4PowBig(&m4check, &m4cand, qBig)
	if m4IsIdentity(&m4check) {
		return &PublicKey{tm: tm2}, nil
	}

	return nil, errors.New("gnarl: neither root in subgroup")
}

func Sign(sk *PrivateKey, msg []byte, challengeFunc func([]byte) [27]byte) (sig *Signature, err error) {
	var k scalar
	err = scRandom(&k)
	if err != nil {
		return nil, err
	}

	var rTM tmat
	tmFixedExp(&rTM, &k)

	rBytes := []byte{:81}
	tmToBytes(rBytes, &rTM)

	hashInput := []byte{:int32(len(msg)) + 81}
	copy(hashInput, msg)
	copy(hashInput[int32(len(msg)):], rBytes)
	eBytes := challengeFunc(hashInput)

	var eSc, es, z scalar
	scFromBytes27(&eSc, eBytes[:])
	scMul(&es, &eSc, &sk.s)
	scSub(&z, &k, &es)

	return &Signature{e: eBytes, z: z}, nil
}

func Verify(pk *PublicKey, msg []byte, sig *Signature, challengeFunc func([]byte) [27]byte) (ok bool) {
	var eSc scalar
	scFromBytes27(&eSc, sig.e[:])

	var rPrime tmat
	tmShamirExp(&rPrime, &sig.z, &pk.tm, &eSc)

	rBytes := []byte{:81}
	tmToBytes(rBytes, &rPrime)

	hashInput := []byte{:int32(len(msg)) + 81}
	copy(hashInput, msg)
	copy(hashInput[int32(len(msg)):], rBytes)
	ePrime := challengeFunc(hashInput)

	return sig.e == ePrime
}

func (sig *Signature) Bytes() (buf []byte) {
	buf = []byte{:54}
	copy(buf[:27], sig.e[:])
	scToBytes27(buf[27:], &sig.z)
	return buf
}

func SignatureFromBytes(data []byte) (sig *Signature, err error) {
	if int32(len(data)) < 54 {
		return nil, errors.New("gnarl: short signature")
	}
	sig = &Signature{}
	copy(sig.e[:], data[:27])
	scFromBytes27(&sig.z, data[27:54])
	return sig, nil
}

func Params() (prime, subgroupOrder, torusOrder *big.Int) {
	return (&big.Int{}).Set(pBig), (&big.Int{}).Set(qBig), (&big.Int{}).Set(nBig)
}