package lafsession

import (
	"testing"
)

// setupPair runs a full handshake and returns established client and server sessions.
func setupPair(b *testing.B) (cli, srv *Session) {
	b.Helper()
	serverID, err := GenerateIdentity()
	if err != nil {
		b.Fatal(err)
	}
	clientID, err := GenerateIdentity()
	if err != nil {
		b.Fatal(err)
	}
	clientPub := clientID.PublicBytes()
	serverPub := serverID.PublicBytes()
	srv, err = New(Config{
		Identity: serverID,
		Role:     RoleServer,
		AcceptPeer: func(pub []byte) bool {
			return bytesEqual32(pub, clientPub)
		},
	})
	if err != nil {
		b.Fatal(err)
	}
	cli, err = New(Config{
		Identity: clientID,
		Role:     RoleClient,
		AcceptPeer: func(pub []byte) bool {
			return bytesEqual32(pub, serverPub)
		},
	})
	if err != nil {
		b.Fatal(err)
	}

	ch, err := cli.StartHandshake()
	if err != nil {
		b.Fatal(err)
	}
	sh, err := srv.HandleHandshake(ch)
	if err != nil {
		b.Fatal(err)
	}
	cf, err := cli.HandleHandshake(sh)
	if err != nil {
		b.Fatal(err)
	}
	if _, err = srv.HandleHandshake(cf); err != nil {
		b.Fatal(err)
	}
	sf, err := srv.SendFinished()
	if err != nil {
		b.Fatal(err)
	}
	if _, err = cli.HandleHandshake(sf); err != nil {
		b.Fatal(err)
	}
	if !cli.Established() || !srv.Established() {
		b.Fatal("not established")
	}
	return cli, srv
}

// BenchmarkHandshake measures the full mutual-auth handshake state machine
// (message build/parse, transcript, AcceptPeer, key install). Crypto ops are
// inherent to the flow but not isolated.
func BenchmarkHandshake(b *testing.B) {
	serverID, err := GenerateIdentity()
	if err != nil {
		b.Fatal(err)
	}
	clientID, err := GenerateIdentity()
	if err != nil {
		b.Fatal(err)
	}
	clientPub := clientID.PublicBytes()
	serverPub := serverID.PublicBytes()

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		srv, err := New(Config{
			Identity:   serverID,
			Role:       RoleServer,
			AcceptPeer: func(pub []byte) bool { return bytesEqual32(pub, clientPub) },
		})
		if err != nil {
			b.Fatal(err)
		}
		cli, err := New(Config{
			Identity:   clientID,
			Role:       RoleClient,
			AcceptPeer: func(pub []byte) bool { return bytesEqual32(pub, serverPub) },
		})
		if err != nil {
			b.Fatal(err)
		}
		ch, err := cli.StartHandshake()
		if err != nil {
			b.Fatal(err)
		}
		sh, err := srv.HandleHandshake(ch)
		if err != nil {
			b.Fatal(err)
		}
		cf, err := cli.HandleHandshake(sh)
		if err != nil {
			b.Fatal(err)
		}
		if _, err = srv.HandleHandshake(cf); err != nil {
			b.Fatal(err)
		}
		sf, err := srv.SendFinished()
		if err != nil {
			b.Fatal(err)
		}
		if _, err = cli.HandleHandshake(sf); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSeal measures session record framing + seal path for a fixed LAF-sized payload.
func BenchmarkSeal(b *testing.B) {
	cli, _ := setupPair(b)
	payload := make([]byte, 256) // typical small LAF envelope
	for i := range payload {
		payload[i] = byte(i)
	}

	b.SetBytes(int64(len(payload)))
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if _, err := cli.Seal(payload); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkOpen measures open + anti-replay accept for sequential records.
func BenchmarkOpen(b *testing.B) {
	cli, srv := setupPair(b)
	payload := make([]byte, 256)
	for i := range payload {
		payload[i] = byte(i)
	}

	frames := make([][]byte, b.N)
	for i := 0; i < b.N; i++ {
		f, err := cli.Seal(payload)
		if err != nil {
			b.Fatal(err)
		}
		frames[i] = append([]byte(nil), f...)
	}

	b.SetBytes(int64(len(payload)))
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if _, err := srv.Open(frames[i]); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSealOpenRoundTrip is client Seal → server Open for one payload size.
func BenchmarkSealOpenRoundTrip(b *testing.B) {
	cli, srv := setupPair(b)
	payload := make([]byte, 256)
	for i := range payload {
		payload[i] = byte(i)
	}

	b.SetBytes(int64(len(payload)))
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		frame, err := cli.Seal(payload)
		if err != nil {
			b.Fatal(err)
		}
		if _, err := srv.Open(frame); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSealOpenSizes exercises framing logic across payload sizes.
func BenchmarkSealOpenSizes(b *testing.B) {
	sizes := []int{64, 256, 1024, 4096}
	for _, sz := range sizes {
		sz := sz
		b.Run(itoa(sz), func(b *testing.B) {
			cli, srv := setupPair(b)
			payload := make([]byte, sz)
			b.SetBytes(int64(sz))
			b.ReportAllocs()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				frame, err := cli.Seal(payload)
				if err != nil {
					b.Fatal(err)
				}
				if _, err := srv.Open(frame); err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}

// BenchmarkReplayWindow measures only the anti-replay window accept path.
func BenchmarkReplayWindow(b *testing.B) {
	w := newReplayWindow(1024)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if !w.accept(uint64(i)) {
			b.Fatal("unexpected reject")
		}
	}
}

// BenchmarkReplayWindowWithGaps simulates out-of-order accepts within the window.
func BenchmarkReplayWindowWithGaps(b *testing.B) {
	w := newReplayWindow(1024)
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		seq := uint64(i)
		// deliver even first, then the preceding odd on next iteration pattern
		if i%2 == 0 {
			if !w.accept(seq + 1) {
				// may already be past; ignore
			}
		}
		_ = w.accept(seq)
	}
}

// BenchmarkFingerprint is package helper logic (hash + hex), not a crypto lib microbench.
func BenchmarkFingerprint(b *testing.B) {
	id, err := GenerateIdentity()
	if err != nil {
		b.Fatal(err)
	}
	pub := id.PublicBytes()
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = FingerprintOf(pub)
	}
}

func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	var buf [16]byte
	i := len(buf)
	for n > 0 {
		i--
		buf[i] = byte('0' + n%10)
		n /= 10
	}
	return string(buf[i:])
}

// BenchmarkSealToOpenTo measures steady-state 0-alloc path with caller-owned buffers.
func BenchmarkSealToOpenTo(b *testing.B) {
	cli, srv := setupPair(b)
	payload := make([]byte, 256)
	for i := range payload {
		payload[i] = byte(i)
	}
	// capacity: header + pt + tag
	sealDst := make([]byte, 0, recordHeaderLen+len(payload)+tagLen)
	openDst := make([]byte, 0, len(payload))

	b.SetBytes(int64(len(payload)))
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		frame, err := cli.SealTo(sealDst[:0], payload)
		if err != nil {
			b.Fatal(err)
		}
		// must copy frame before OpenTo if we reuse sealDst; use independent buffer
		pt, err := srv.OpenTo(openDst[:0], frame)
		if err != nil {
			b.Fatal(err)
		}
		if len(pt) != len(payload) {
			b.Fatal("len")
		}
		// frame still points into sealDst; next SealTo with [:0] reuses same cap
		_ = frame
	}
}

// BenchmarkSealSteady uses Seal after warm-up; internal scratch avoids allocs.
func BenchmarkSealSteady(b *testing.B) {
	cli, _ := setupPair(b)
	payload := make([]byte, 256)
	// warm-up grow
	if _, err := cli.Seal(payload); err != nil {
		b.Fatal(err)
	}
	b.SetBytes(int64(len(payload)))
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if _, err := cli.Seal(payload); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkOpenSteady uses Open after warm-up with pre-copied frames.
func BenchmarkOpenSteady(b *testing.B) {
	cli, srv := setupPair(b)
	payload := make([]byte, 256)
	frames := make([][]byte, b.N)
	for i := 0; i < b.N; i++ {
		f, err := cli.Seal(payload)
		if err != nil {
			b.Fatal(err)
		}
		frames[i] = append([]byte(nil), f...)
	}
	// warm openBuf
	if _, err := srv.Open(frames[0]); err != nil && err != ErrReplay {
		// first frame opens; remaining still valid if seq advances
	}
	// re-setup pair for clean window
	cli, srv = setupPair(b)
	frames = make([][]byte, b.N)
	for i := 0; i < b.N; i++ {
		f, err := cli.Seal(payload)
		if err != nil {
			b.Fatal(err)
		}
		frames[i] = append([]byte(nil), f...)
	}
	// warm
	if _, err := srv.Open(append([]byte(nil), frames[0]...)); err == nil {
		// consumed seq 0; rebuild
	}
	cli, srv = setupPair(b)
	frames = make([][]byte, b.N)
	for i := 0; i < b.N; i++ {
		f, _ := cli.Seal(payload)
		frames[i] = append([]byte(nil), f...)
	}
	if _, err := srv.Open(frames[0]); err != nil {
		b.Fatal(err)
	}
	b.SetBytes(int64(len(payload)))
	b.ReportAllocs()
	b.ResetTimer()
	for i := 1; i < b.N; i++ {
		if _, err := srv.Open(frames[i]); err != nil {
			b.Fatal(err)
		}
	}
}

func bytesEqual32(a, b []byte) bool {
	if len(a) != 32 || len(b) != 32 {
		return false
	}
	var v byte
	for i := 0; i < 32; i++ {
		v |= a[i] ^ b[i]
	}
	return v == 0
}