package testkit

import (
	"crypto/ed25519"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"math/big"
	"slices"

	"g.activething.com/go/DateKeys/internal/ed25519strict"
)

// Ed25519StrictFile is testdata/vectors/ed25519_strict.json: Ed25519
// signatures with the result of the strict profile of the author signature
// (spec v0.11, §29.9), built after the cases of «Taming the many EdDSAs»
// (Chalkias, Garillot, Nikolaenko, 2020): small-order and non-canonical
// keys, non-canonical R and S, keys of mixed order, and signatures that only
// the equation with the cofactor accepts. Stdlib is what crypto/ed25519 of
// Go says, for the record: where it says true and Valid is false, an
// implementation needs the checks of package ed25519strict.
type Ed25519StrictFile struct {
	Spec        string                `json:"spec"`
	Description string                `json:"description"`
	Vectors     []Ed25519StrictVector `json:"vectors"`
}

// Ed25519StrictVector is one signature, its message and its public key, in
// hexadecimal.
type Ed25519StrictVector struct {
	Name      string `json:"name"`
	Message   string `json:"message"`
	PublicKey string `json:"public_key"`
	Signature string `json:"signature"`
	Valid     bool   `json:"valid"`
	Stdlib    bool   `json:"stdlib"`
}

// Ed25519StrictVectors builds ed25519_strict.json, and fails if
// ed25519strict.Verify does not give the result each case is built for.
func Ed25519StrictVectors() (Ed25519StrictFile, error) {
	f := Ed25519StrictFile{
		Spec: SpecVersion,
		Description: "Ed25519 signatures and the result of the strict profile of the author signature (spec v0.11, §29.9), after the cases of " +
			"«Taming the many EdDSAs»; stdlib is the result of crypto/ed25519 of Go, for the record. Generated by the reference implementation. See testdata/README.md.",
	}
	var errs []error
	for _, c := range ed25519Cases() {
		if got := ed25519strict.Verify(c.pub, c.msg, c.sig); got != c.valid {
			errs = append(errs, fmt.Errorf("ed25519 %q: Verify = %v, want %v", c.name, got, c.valid))
		}
		f.Vectors = append(f.Vectors, Ed25519StrictVector{
			Name:      c.name,
			Message:   hex.EncodeToString(c.msg),
			PublicKey: hex.EncodeToString(c.pub),
			Signature: hex.EncodeToString(c.sig),
			Valid:     c.valid,
			Stdlib:    ed25519.Verify(ed25519.PublicKey(c.pub), c.msg, c.sig),
		})
	}
	return f, errors.Join(errs...)
}

type ed25519Case struct {
	name          string
	msg, pub, sig []byte
	valid         bool
}

func ed25519Cases() []ed25519Case {
	seed := sha256.Sum256([]byte("DateKeys ed25519_strict vectors"))
	priv := ed25519.NewKeyFromSeed(seed[:])
	pub := []byte(priv.Public().(ed25519.PublicKey))
	msg := []byte("DateKeys")
	sig := ed25519.Sign(priv, msg)

	cases := []ed25519Case{{"a valid signature", msg, pub, sig, true}}

	// S + ℓ is below 2^253, so its top bits are clear, but S is not canonical.
	s := leInt(sig[32:])
	s.Add(s, edL)
	cases = append(cases, ed25519Case{"S + ℓ", msg, pub, slices.Concat(sig[:32], leBytes(s)), false})
	high := slices.Clone(sig)
	high[63] |= 0x20
	cases = append(cases, ed25519Case{"S with bit 253 set", msg, pub, high, false})
	r := slices.Clone(sig)
	copy(r[:32], nonCanonicalZero(0))
	cases = append(cases, ed25519Case{"R not canonical", msg, pub, r, false})

	// A of small order, R the identity and S = 0: [S]B − [k]A = −[k]A is
	// the identity when the order of A divides k, so a message is searched.
	identity := edEncode(edPoint{big.NewInt(0), big.NewInt(1)})
	forged := slices.Concat(identity, make([]byte, 32))
	for i, t := range edTorsion() {
		a := edEncode(t)
		m := messageFor(identity, a, func(k *big.Int) bool { return new(big.Int).Mod(k, big.NewInt(8)).Sign() == 0 })
		cases = append(cases, ed25519Case{fmt.Sprintf("A of small order, the point %d of the torsion, R the identity and S = 0", i), m, a, forged, false})
	}
	// The same with A not canonical: y = p + 0, a point of order 4, with
	// either sign; and the identity with its sign bit set.
	for _, sign := range []byte{0, 0x80} {
		a := nonCanonicalZero(sign)
		m := messageFor(identity, a, func(k *big.Int) bool { return new(big.Int).Mod(k, big.NewInt(8)).Sign() == 0 })
		cases = append(cases, ed25519Case{fmt.Sprintf("A not canonical, y = p, sign %d, R the identity and S = 0", sign>>7), m, a, forged, false})
	}
	negZero := slices.Clone(identity)
	negZero[31] |= 0x80
	m := messageFor(identity, negZero, func(*big.Int) bool { return true })
	cases = append(cases, ed25519Case{"A the identity with the sign bit, R the identity and S = 0", m, negZero, forged, false})

	// A of mixed order, [a]B plus a point of order 8: the equation without
	// the cofactor holds only when [k]T is the identity, that is, when 8
	// divides k; the equation with the cofactor holds always.
	a := new(big.Int).Mod(leInt(seed[:]), edL)
	t8 := edTorsion()[edOrder8]
	mixed := edAdd(edMul(a, edBase()), t8)
	am := edEncode(mixed)
	rr := new(big.Int).Mod(leInt(slices.Concat(seed[:], seed[:])), edL)
	rp := edEncode(edMul(rr, edBase()))
	for _, holds := range []bool{true, false} {
		m := messageFor(rp, am, func(k *big.Int) bool { return (new(big.Int).Mod(k, big.NewInt(8)).Sign() == 0) == holds })
		k := hramScalar(rp, am, m)
		sv := new(big.Int).Mod(new(big.Int).Add(rr, new(big.Int).Mul(k, a)), edL)
		name := "A of mixed order, 8 divides k: the equation without the cofactor holds"
		if !holds {
			name = "A of mixed order, 8 does not divide k: only the equation with the cofactor holds"
		}
		cases = append(cases, ed25519Case{name, m, am, slices.Concat(rp, leBytes(sv)), holds})
	}

	// R the identity with A of prime order: S = k·a makes [S]B − [k]A the
	// identity, which is R; libsodium rejects an R of small order, and this
	// profile does not.
	ap := edEncode(edMul(a, edBase()))
	m = []byte("DateKeys: R the identity")
	k := hramScalar(identity, ap, m)
	sv := new(big.Int).Mod(new(big.Int).Mul(k, a), edL)
	cases = append(cases, ed25519Case{"R the identity, A of prime order", m, ap, slices.Concat(identity, leBytes(sv)), true})
	return cases
}

// messageFor returns the first message "DateKeys n", n = 0, 1, ..., whose k =
// SHA-512(R ‖ A ‖ M) mod ℓ satisfies ok.
func messageFor(r, a []byte, ok func(k *big.Int) bool) []byte {
	for n := uint64(0); ; n++ {
		m := binary.BigEndian.AppendUint64([]byte("DateKeys "), n)
		if ok(hramScalar(r, a, m)) {
			return m
		}
	}
}

func hramScalar(r, a, m []byte) *big.Int {
	h := sha512.Sum512(slices.Concat(r, a, m))
	return new(big.Int).Mod(leInt(h[:]), edL)
}

// nonCanonicalZero is y = p, the non-canonical encoding of y = 0, with the
// sign bit given.
func nonCanonicalZero(sign byte) []byte {
	b := make([]byte, 32)
	b[0] = 0xed
	for i := 1; i < 31; i++ {
		b[i] = 0xff
	}
	b[31] = 0x7f | sign
	return b
}

// Arithmetic on edwards25519 with math/big, slow and simple, only to build
// these vectors: the reference never computes on points itself.

var (
	edP = new(big.Int).Sub(new(big.Int).Lsh(big.NewInt(1), 255), big.NewInt(19))
	edL = func() *big.Int {
		l, _ := new(big.Int).SetString("7237005577332262213973186563042994240857116359379907606001950938285454250989", 10)
		return l
	}()
	edD = func() *big.Int {
		d := new(big.Int).Mul(big.NewInt(-121665), edInv(big.NewInt(121666)))
		return d.Mod(d, edP)
	}()
	edSqrtM1 = new(big.Int).Exp(big.NewInt(2), new(big.Int).Rsh(new(big.Int).Sub(edP, big.NewInt(1)), 2), edP)
)

// edOrder8 is the index in edTorsion of a point of order 8.
const edOrder8 = 1

type edPoint struct{ x, y *big.Int }

func edInv(x *big.Int) *big.Int {
	return new(big.Int).Exp(x, new(big.Int).Sub(edP, big.NewInt(2)), edP)
}

func edAdd(a, b edPoint) edPoint {
	t := new(big.Int).Mul(edD, a.x)
	t.Mul(t, b.x).Mul(t, a.y).Mul(t, b.y).Mod(t, edP)
	x := new(big.Int).Add(new(big.Int).Mul(a.x, b.y), new(big.Int).Mul(b.x, a.y))
	x.Mul(x, edInv(new(big.Int).Add(big.NewInt(1), t))).Mod(x, edP)
	y := new(big.Int).Add(new(big.Int).Mul(a.y, b.y), new(big.Int).Mul(a.x, b.x))
	y.Mul(y, edInv(new(big.Int).Mod(new(big.Int).Sub(big.NewInt(1), t), edP))).Mod(y, edP)
	return edPoint{x, y}
}

func edMul(k *big.Int, a edPoint) edPoint {
	r := edPoint{big.NewInt(0), big.NewInt(1)}
	for i := k.BitLen() - 1; i >= 0; i-- {
		r = edAdd(r, r)
		if k.Bit(i) == 1 {
			r = edAdd(r, a)
		}
	}
	return r
}

// edX recovers x from y and its sign bit, or nil when y is not on the curve.
func edX(y *big.Int, sign uint) *big.Int {
	yy := new(big.Int).Mul(y, y)
	num := new(big.Int).Sub(yy, big.NewInt(1))
	den := new(big.Int).Add(new(big.Int).Mul(edD, yy), big.NewInt(1))
	xx := new(big.Int).Mul(num, edInv(den.Mod(den, edP)))
	xx.Mod(xx, edP)
	if xx.Sign() == 0 {
		return big.NewInt(0)
	}
	x := new(big.Int).Exp(xx, new(big.Int).Rsh(new(big.Int).Add(edP, big.NewInt(3)), 3), edP)
	if new(big.Int).Mod(new(big.Int).Sub(new(big.Int).Mul(x, x), xx), edP).Sign() != 0 {
		x.Mul(x, edSqrtM1).Mod(x, edP)
	}
	if new(big.Int).Mod(new(big.Int).Sub(new(big.Int).Mul(x, x), xx), edP).Sign() != 0 {
		return nil
	}
	if x.Bit(0) != sign {
		x.Sub(edP, x)
	}
	return x
}

func edBase() edPoint {
	y := new(big.Int).Mul(big.NewInt(4), edInv(big.NewInt(5)))
	y.Mod(y, edP)
	return edPoint{edX(y, 0), y}
}

func edEncode(a edPoint) []byte {
	b := leBytes(a.y)
	b[31] |= byte(a.x.Bit(0)) << 7
	return b
}

// edTorsion returns the eight points of small order, [i]T for a point T of
// order 8 and i from 0 to 7: T is [ℓ]P for the first point P, by y, whose
// [ℓ]P is not of order 4 or less.
func edTorsion() []edPoint {
	for y := int64(2); ; y++ {
		x := edX(big.NewInt(y), 0)
		if x == nil {
			continue
		}
		t := edMul(edL, edPoint{x, big.NewInt(y)})
		if q := edMul(big.NewInt(4), t); q.x.Sign() == 0 && q.y.Cmp(big.NewInt(1)) == 0 {
			continue
		}
		out := make([]edPoint, 8)
		out[0] = edPoint{big.NewInt(0), big.NewInt(1)}
		for i := 1; i < 8; i++ {
			out[i] = edAdd(out[i-1], t)
		}
		return out
	}
}

func leInt(b []byte) *big.Int {
	be := slices.Clone(b)
	slices.Reverse(be)
	return new(big.Int).SetBytes(be)
}

func leBytes(x *big.Int) []byte {
	b := x.FillBytes(make([]byte, 32))
	slices.Reverse(b)
	return b
}

// Ed25519Decodes reports whether the encoding a, with y below p, decodes to a
// point, by computing its x with the square root of RFC 8032 5.1.3: an
// oracle for ed25519strict.OnCurve that does not use Euler's criterion.
func Ed25519Decodes(a []byte) bool {
	b := slices.Clone(a)
	b[31] &= 0x7f
	return edX(leInt(b), uint(a[31]>>7)) != nil
}

// Ed25519Torsion returns the canonical encodings of the eight points of small
// order, computed from the curve, to check the table of ed25519strict.
func Ed25519Torsion() [][32]byte {
	var out [][32]byte
	for _, p := range edTorsion() {
		out = append(out, [32]byte(edEncode(p)))
	}
	return out
}