package dleq_test

import (
	"bytes"
	"crypto"
	"crypto/rand"
	"errors"
	"fmt"
	"testing"

	"github.com/cloudflare/circl/group"
	"github.com/cloudflare/circl/internal/test"
	"github.com/cloudflare/circl/zk/dleq"
)

func TestDLEQ(t *testing.T) {
	for _, g := range []group.Group{
		group.P256,
		group.P384,
		group.P521,
		group.Ristretto255,
	} {
		t.Run(g.(fmt.Stringer).String(), func(t *testing.T) {
			params := dleq.Params{g, crypto.SHA256, []byte("domain_sep_string")}
			Peggy := dleq.Prover{params}
			Victor := dleq.Verifier{params}

			k := g.RandomScalar(rand.Reader)
			A := g.RandomElement(rand.Reader)
			kA := g.NewElement().Mul(A, k)

			B := g.RandomElement(rand.Reader)
			kB := g.NewElement().Mul(B, k)

			proof, err := Peggy.Prove(k, A, kA, B, kB, rand.Reader)
			test.CheckNoErr(t, err, "wrong proof generation")
			test.CheckOk(Victor.Verify(A, kA, B, kB, proof), "proof must verify", t)

			rr := g.RandomScalar(rand.Reader)
			proof, err = Peggy.ProveWithRandomness(k, A, kA, B, kB, rr)
			test.CheckNoErr(t, err, "wrong proof generation")
			test.CheckOk(Victor.Verify(A, kA, B, kB, proof), "proof must verify", t)

			const N = 4
			C := make([]group.Element, N)
			kC := make([]group.Element, N)
			for i := 0; i < N; i++ {
				C[i] = g.RandomElement(rand.Reader)
				kC[i] = g.NewElement().Mul(C[i], k)
			}
			proof, err = Peggy.ProveBatch(k, A, kA, C, kC, rand.Reader)
			test.CheckNoErr(t, err, "wrong proof generation")
			test.CheckOk(Victor.VerifyBatch(A, kA, C, kC, proof), "proof must verify", t)

			testMarshal(t, g, proof)
			testErrors(t, &Peggy, &Victor, g, k, A, kA, B, kB)
		})
	}
}

func TestChallengeBindsBase(t *testing.T) {
	g := group.P256
	params := dleq.Params{G: g, H: crypto.SHA256, DST: []byte("domain_sep_string")}
	prover := dleq.Prover{Params: params}
	verifier := dleq.Verifier{Params: params}

	k := g.RandomScalar(rand.Reader)
	b := g.RandomElement(rand.Reader)
	kb := g.NewElement().Mul(b, k)
	ka := g.RandomElement(rand.Reader)

	r := g.RandomScalar(rand.Reader)
	a := g.Generator()
	proof, err := prover.ProveWithRandomness(k, a, ka, b, kb, r)
	test.CheckNoErr(t, err, "wrong proof generation")

	rawProof, err := proof.MarshalBinary()
	test.CheckNoErr(t, err, "error on marshaling proof")
	scalarLen := int(g.Params().ScalarLength)
	c := g.NewScalar()
	err = c.UnmarshalBinary(rawProof[:scalarLen])
	test.CheckNoErr(t, err, "error on unmarshaling challenge")
	s := g.NewScalar()
	err = s.UnmarshalBinary(rawProof[scalarLen:])
	test.CheckNoErr(t, err, "error on unmarshaling response")

	rA := g.NewElement().Mul(a, r)
	cKA := g.NewElement().Mul(ka, c)
	diff := g.NewElement().Add(rA, g.NewElement().Neg(cKA))
	forgedA := g.NewElement().Mul(diff, g.NewScalar().Inv(s))

	if g.NewElement().Mul(forgedA, k).IsEqual(ka) {
		t.Fatal("unexpected true statement")
	}
	if verifier.Verify(forgedA, ka, b, kb, proof) {
		t.Fatal("Verify accepted a proof for a false statement with an adaptively chosen base")
	}
}

func TestMalformedBatch(t *testing.T) {
	g := group.P256
	params := dleq.Params{G: g, H: crypto.SHA256, DST: []byte("domain_sep_string")}
	prover := dleq.Prover{Params: params}
	verifier := dleq.Verifier{Params: params}
	k := g.RandomScalar(rand.Reader)
	a := g.RandomElement(rand.Reader)
	ka := g.NewElement().Mul(a, k)
	bi := []group.Element{g.RandomElement(rand.Reader), g.RandomElement(rand.Reader)}
	kbi := []group.Element{
		g.NewElement().Mul(bi[0], k),
		g.NewElement().Mul(bi[1], k),
	}
	rnd := g.RandomScalar(rand.Reader)
	proof, err := prover.ProveBatchWithRandomness(k, a, ka, bi, kbi, rnd)
	test.CheckNoErr(t, err, "wrong proof generation")

	tests := []struct {
		name string
		bi   []group.Element
		kbi  []group.Element
	}{
		{"empty", nil, nil},
		{"short evaluated vector", bi, kbi[:1]},
		{"long evaluated vector", bi, append(kbi, g.RandomElement(rand.Reader))},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			if _, err := prover.ProveBatchWithRandomness(
				k, a, ka, tc.bi, tc.kbi, rnd,
			); !errors.Is(err, dleq.ErrInvalidBatch) {
				t.Fatalf("got proving error %v, want %v", err, dleq.ErrInvalidBatch)
			}
			if verifier.VerifyBatch(a, ka, tc.bi, tc.kbi, proof) {
				t.Fatal("malformed batch verified")
			}
		})
	}

	if verifier.VerifyBatch(a, ka, bi, kbi, nil) {
		t.Fatal("nil proof verified")
	}
}

func testMarshal(t *testing.T, g group.Group, proof *dleq.Proof) {
	t.Helper()

	wantProofBytes, err := proof.MarshalBinary()
	test.CheckNoErr(t, err, "error on marshaling proof")

	gotProof := new(dleq.Proof)
	err = gotProof.UnmarshalBinary(g, wantProofBytes)
	test.CheckNoErr(t, err, "error on unmarshaling proof")

	gotProofBytes, err := gotProof.MarshalBinary()
	test.CheckNoErr(t, err, "error on marshaling proof")

	if !bytes.Equal(gotProofBytes, wantProofBytes) {
		test.ReportError(t, gotProofBytes, wantProofBytes)
	}
}

func testErrors(
	t *testing.T,
	Peggy *dleq.Prover,
	Victor *dleq.Verifier,
	g group.Group,
	k group.Scalar, a, ka, b, kb group.Element,
) {
	goodProof, err := Peggy.Prove(k, a, ka, b, kb, rand.Reader)
	test.CheckNoErr(t, err, "wrong proof generation")

	proofBytes, err := goodProof.MarshalBinary()
	test.CheckNoErr(t, err, "error on marshaling proof")

	// Tamper proof (in transit)
	proofBytes[7] ^= 0xFF

	tamperedProof := new(dleq.Proof)
	err = tamperedProof.UnmarshalBinary(g, proofBytes[:5])
	test.CheckIsErr(t, err, "unmarshal must fail")

	// Trailing bytes must be rejected to keep the encoding canonical.
	err = tamperedProof.UnmarshalBinary(g, append(append([]byte{}, proofBytes...), 0x00))
	test.CheckIsErr(t, err, "unmarshal with trailing bytes must fail")

	err = tamperedProof.UnmarshalBinary(g, proofBytes)
	test.CheckNoErr(t, err, "proof must be unmarshaled")
	test.CheckOk(false == Victor.Verify(a, ka, b, kb, tamperedProof), "proof must not verify", t)

	// Tamper elements
	bada := g.NewElement().Neg(a)
	test.CheckOk(false == Victor.Verify(bada, ka, b, kb, goodProof), "proof must not verify", t)
	badka := g.NewElement().Neg(ka)
	test.CheckOk(false == Victor.Verify(a, badka, b, kb, goodProof), "proof must not verify", t)
	badb := g.NewElement().Neg(b)
	test.CheckOk(false == Victor.Verify(a, ka, badb, kb, goodProof), "proof must not verify", t)
	badkb := g.NewElement().Neg(kb)
	test.CheckOk(false == Victor.Verify(a, ka, b, badkb, goodProof), "proof must not verify", t)
}

func BenchmarkDLEQ(b *testing.B) {
	g := group.P256
	params := dleq.Params{g, crypto.SHA256, []byte("domain_sep_string")}
	Peggy := dleq.Prover{params}
	Victor := dleq.Verifier{params}

	k := g.RandomScalar(rand.Reader)
	A := g.Generator()
	kA := g.NewElement().MulGen(k)

	B := g.RandomElement(rand.Reader)
	kB := g.NewElement().Mul(B, k)
	rr := g.RandomScalar(rand.Reader)

	proof, _ := Peggy.ProveWithRandomness(k, A, kA, B, kB, rr)

	const N = 4
	C := make([]group.Element, N)
	kC := make([]group.Element, N)
	for i := 0; i < N; i++ {
		C[i] = g.RandomElement(rand.Reader)
		kC[i] = g.NewElement().Mul(C[i], k)
	}

	proofBatched, _ := Peggy.ProveBatchWithRandomness(k, A, kA, C, kC, rr)

	b.Run("Prove", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			_, _ = Peggy.ProveWithRandomness(k, A, kA, B, kB, rr)
		}
	})
	b.Run("Verify", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			_ = Victor.Verify(A, kA, B, kB, proof)
		}
	})
	b.Run(fmt.Sprint("ProveBatch=", N), func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			_, _ = Peggy.ProveBatchWithRandomness(k, A, kA, C, kC, rr)
		}
	})
	b.Run(fmt.Sprint("VerifyBatch=", N), func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			_ = Victor.VerifyBatch(A, kA, C, kC, proofBatched)
		}
	})
}
