package lafsession

import (
	"testing"
)

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
}

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)
		}
	}
}

func BenchmarkSeal(b *testing.B) {
	cli, _ := 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++ {
		if _, err := cli.Seal(payload); err != nil {
			b.Fatal(err)
		}
	}
}

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)
		}
	}
}

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)
		}
	}
}

func BenchmarkSealOpenSizes(b *testing.B) {
	sizes := []int{64, 512, 1024}
	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)
				}
			}
		})
	}
}

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")
		}
	}
}

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:])
}

func BenchmarkSealToOpenTo(b *testing.B) {
	cli, srv := setupPair(b)
	payload := make([]byte, 256)
	for i := range payload {
		payload[i] = byte(i)
	}
	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)
		}
		pt, err := srv.OpenTo(openDst[:0], frame)
		if err != nil {
			b.Fatal(err)
		}
		if len(pt) != len(payload) {
			b.Fatal("len")
		}
		_ = frame
	}
}

func BenchmarkSealSteady(b *testing.B) {
	cli, _ := setupPair(b)
	payload := make([]byte, 256)
	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)
		}
	}
}

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...)
	}
	if _, err := srv.Open(frames[0]); err != nil && err != ErrReplay {
	}
	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...)
	}
	if _, err := srv.Open(append([]byte(nil), frames[0]...)); err == nil {
	}
	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
}