package crypto

import (
	"crypto/rand"
	"crypto/sha256"
	"testing"

	"github.com/btcsuite/btcd/btcec/v2"
	"github.com/btcsuite/btcd/btcec/v2/schnorr"

	"git.smesh.lol/gnarl-hamadryad/crypto/gnarl"
	"git.smesh.lol/gnarl-hamadryad/crypto/ring"
)

// --- Hashing ---

func BenchmarkSHA256_32(b *testing.B) {
	msg := make([]byte, 32)
	for i := range msg {
		msg[i] = byte(i)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		sha256.Sum256(msg)
	}
}

func BenchmarkSHA256_128(b *testing.B) {
	msg := make([]byte, 128)
	for i := range msg {
		msg[i] = byte(i)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		sha256.Sum256(msg)
	}
}

func BenchmarkHamadryad_128(b *testing.B) {
	msg := make([]byte, 128)
	for i := range msg {
		msg[i] = byte(i)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		Hash(msg)
	}
}

func BenchmarkGHash_128(b *testing.B) {
	msg := make([]byte, 128)
	for i := range msg {
		msg[i] = byte(i)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		GHash(msg)
	}
}

func BenchmarkGMid_128(b *testing.B) {
	msg := make([]byte, 128)
	for i := range msg {
		msg[i] = byte(i)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		GMid(msg)
	}
}

// --- Key Generation ---

func BenchmarkBIP340_KeyGen(b *testing.B) {
	var seed [32]byte
	for i := 0; i < b.N; i++ {
		rand.Read(seed[:])
		privKey, _ := btcec.PrivKeyFromBytes(seed[:])
		_ = schnorr.SerializePubKey(privKey.PubKey())
	}
}

func BenchmarkGnarl_KeyGen(b *testing.B) {
	for i := 0; i < b.N; i++ {
		_, _, err := gnarl.GenerateKey()
		if err != nil {
			b.Fatal(err)
		}
	}
}

// --- Signing ---

func BenchmarkBIP340_Sign(b *testing.B) {
	var seed [32]byte
	rand.Read(seed[:])
	privKey, _ := btcec.PrivKeyFromBytes(seed[:])

	msg := sha256.Sum256([]byte("benchmark message for signing"))
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, err := schnorr.Sign(privKey, msg[:])
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkGnarl_Sign(b *testing.B) {
	sk, _, err := gnarl.GenerateKey()
	if err != nil {
		b.Fatal(err)
	}
	msg := []byte("benchmark message for signing")
	challengeFunc := func(data []byte) [27]byte {
		return GnarlSchnorrChallenge(data)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, err := gnarl.Sign(sk, msg, challengeFunc)
		if err != nil {
			b.Fatal(err)
		}
	}
}

// --- Verification ---

func BenchmarkBIP340_Verify(b *testing.B) {
	var seed [32]byte
	rand.Read(seed[:])
	privKey, _ := btcec.PrivKeyFromBytes(seed[:])
	pubKey := privKey.PubKey()

	msg := sha256.Sum256([]byte("benchmark message for verification"))
	sig, err := schnorr.Sign(privKey, msg[:])
	if err != nil {
		b.Fatal(err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if !sig.Verify(msg[:], pubKey) {
			b.Fatal("verification failed")
		}
	}
}

func BenchmarkGnarl_Verify(b *testing.B) {
	sk, pk, err := gnarl.GenerateKey()
	if err != nil {
		b.Fatal(err)
	}
	msg := []byte("benchmark message for verification")
	challengeFunc := func(data []byte) [27]byte {
		return GnarlSchnorrChallenge(data)
	}
	sig, err := gnarl.Sign(sk, msg, challengeFunc)
	if err != nil {
		b.Fatal(err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if !gnarl.Verify(pk, msg, sig, challengeFunc) {
			b.Fatal("verification failed")
		}
	}
}

// --- Wire Format ---

func BenchmarkGnarlSeal_128(b *testing.B) {
	var secret Hamadryad
	h := GHash([]byte("bench-secret"))
	copy(secret[:], h[:])
	identity := GMid([]byte("bench-id"))
	nonce := GnarlNonceFromCounter(0, 1)

	msg := make([]byte, 128)
	for i := range msg {
		msg[i] = byte(i)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		GnarlSeal(secret, identity, nonce, msg)
	}
}

func BenchmarkGnarlOpen_128(b *testing.B) {
	var secret Hamadryad
	h := GHash([]byte("bench-secret"))
	copy(secret[:], h[:])
	identity := GMid([]byte("bench-id"))
	nonce := GnarlNonceFromCounter(0, 1)

	msg := make([]byte, 128)
	for i := range msg {
		msg[i] = byte(i)
	}
	pkt := GnarlSeal(secret, identity, nonce, msg)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, err := GnarlOpen(secret, pkt)
		if err != nil {
			b.Fatal(err)
		}
	}
}

// --- Lattice Ring (post-quantum SVP-hard) ---

func BenchmarkRing_KeyGen(b *testing.B) {
	gp := ring.SmallGPVParams()
	for i := 0; i < b.N; i++ {
		ring.GPVKeyGen(gp)
	}
}

func BenchmarkRing_Sign(b *testing.B) {
	gp := ring.SmallGPVParams()
	_, sk := ring.GPVKeyGen(gp)
	msg := []byte("benchmark message for signing")
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ring.GPVSign(sk, msg)
	}
}

func BenchmarkRing_Verify(b *testing.B) {
	gp := ring.SmallGPVParams()
	pk, sk := ring.GPVKeyGen(gp)
	msg := []byte("benchmark message for verification")
	sig := ring.GPVSign(sk, msg)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if !ring.GPVVerify(pk, msg, sig) {
			b.Fatal("verification failed")
		}
	}
}

func BenchmarkRing_KEMKeyGen(b *testing.B) {
	kp := ring.DefaultKEMParams()
	for i := 0; i < b.N; i++ {
		ring.KEMKeyGen(kp)
	}
}

func BenchmarkRing_KEMEncapsulate(b *testing.B) {
	kp := ring.DefaultKEMParams()
	pk, _ := ring.KEMKeyGen(kp)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ring.Encapsulate(pk)
	}
}

func BenchmarkRing_KEMDecapsulate(b *testing.B) {
	kp := ring.DefaultKEMParams()
	pk, sk := ring.KEMKeyGen(kp)
	_, ct, _ := ring.Encapsulate(pk)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ring.Decapsulate(sk, ct)
	}
}

func BenchmarkRing_HEEncrypt(b *testing.B) {
	kp := ring.DefaultHEParams()
	pk, _, _ := ring.HEKeyGen(kp)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ring.HEEncrypt(pk, 1)
	}
}

func BenchmarkRing_HEAdd(b *testing.B) {
	kp := ring.DefaultHEParams()
	pk, _, _ := ring.HEKeyGen(kp)
	ct0 := ring.HEEncrypt(pk, 0)
	ct1 := ring.HEEncrypt(pk, 1)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ring.HEAdd(ct0, ct1)
	}
}

func BenchmarkRing_HEMul(b *testing.B) {
	kp := ring.DefaultHEParams()
	pk, _, rlk := ring.HEKeyGen(kp)
	ct0 := ring.HEEncrypt(pk, 1)
	ct1 := ring.HEEncrypt(pk, 1)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ring.HEMul(ct0, ct1, rlk)
	}
}

func BenchmarkRing_SISHash_128(b *testing.B) {
	sp := ring.HamadryadSISParams()
	h := ring.NewSISHasher(sp, "hamadryad-swifft-v1-dendrite-kismet")
	msg := make([]byte, 128)
	for i := range msg {
		msg[i] = byte(i)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		h.HashBytes(msg)
	}
}