package crypto

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"os"
	"path/filepath"
	"runtime"
	"testing"
)

// vectorsJSON mirrors the JSON written by tools/export-kat (in Rust). Every field is hex.
type vectorsJSON struct {
	CAFingerprint       string `json:"ca_fingerprint"`
	ClientX25519Priv    string `json:"client_x25519_priv"`
	ClientX25519Pub     string `json:"client_x25519_pub"`
	ClientKyberPriv     string `json:"client_kyber_priv"`
	ClientKyberPub      string `json:"client_kyber_pub"`
	ServerX25519EphPriv string `json:"server_x25519_eph_priv"`
	ServerX25519EphPub  string `json:"server_x25519_eph_pub"`
	ServerKyberCt       string `json:"server_kyber_ct"`
	ClientNonce         string `json:"client_nonce"`
	ServerNonce         string `json:"server_nonce"`
	X25519SS            string `json:"x25519_ss"`
	KyberSS             string `json:"kyber_ss"`
	SessionKeys         struct {
		C2S string `json:"c2s"`
		S2C string `json:"s2c"`
	} `json:"session_keys"`
	TranscriptHash     string `json:"transcript_hash"`
	ClientFinishedHmac string `json:"client_finished_hmac"`
	ServerFinishedHmac string `json:"server_finished_hmac"`
	DatagramTest       struct {
		Seq          uint64 `json:"seq"`
		Frame        string `json:"frame"`
		Key          string `json:"key"`
		SealedRecord string `json:"sealed_record"`
	} `json:"datagram_test"`
	KnockTest struct {
		CAFingerprint string `json:"ca_fingerprint"`
		UnixMinute    uint64 `json:"unix_minute"`
		Knock         string `json:"knock"`
	} `json:"knock_test"`
}

// loadVectors finds the vectors file at <module>/kat/vectors.json. The file is created by
//
//	cargo run -p export-kat
//
// from the workspace root.
func loadVectors(t *testing.T) *vectorsJSON {
	t.Helper()
	// crypto_test.go is at singbox-aura/aura/crypto/. The KAT lives at singbox-aura/kat/.
	_, thisFile, _, ok := runtime.Caller(0)
	if !ok {
		t.Fatal("runtime.Caller failed")
	}
	path := filepath.Join(filepath.Dir(thisFile), "..", "..", "kat", "vectors.json")
	data, err := os.ReadFile(path)
	if err != nil {
		t.Skipf("KAT vectors.json not present at %s — run `cargo run -p export-kat` first: %v", path, err)
		return nil
	}
	var v vectorsJSON
	if err := json.Unmarshal(data, &v); err != nil {
		t.Fatalf("parse vectors.json: %v", err)
	}
	return &v
}

func mustHex(t *testing.T, s string) []byte {
	t.Helper()
	b, err := hex.DecodeString(s)
	if err != nil {
		t.Fatalf("hex decode %q: %v", s, err)
	}
	return b
}

func mustHex32(t *testing.T, s string) [32]byte {
	b := mustHex(t, s)
	if len(b) != 32 {
		t.Fatalf("want 32 bytes, got %d", len(b))
	}
	var out [32]byte
	copy(out[:], b)
	return out
}

// TestKAT_SessionKeys: HKDF-derive from the shared secrets in the vector reproduces the
// session_keys.{c2s,s2c} byte-for-byte.
func TestKAT_SessionKeys(t *testing.T) {
	v := loadVectors(t)
	if v == nil {
		return
	}
	xss := mustHex32(t, v.X25519SS)
	kss := mustHex32(t, v.KyberSS)
	cn := mustHex32(t, v.ClientNonce)
	sn := mustHex32(t, v.ServerNonce)
	wantC2S := mustHex(t, v.SessionKeys.C2S)
	wantS2C := mustHex(t, v.SessionKeys.S2C)

	shared := &HybridSharedSecret{X25519SS: xss, MLKEMSS: kss}
	keys := DeriveSessionKeys(shared, cn, sn)
	if !bytes.Equal(keys.ClientToServer[:], wantC2S) {
		t.Fatalf("c2s mismatch:\n got %x\nwant %x", keys.ClientToServer, wantC2S)
	}
	if !bytes.Equal(keys.ServerToClient[:], wantS2C) {
		t.Fatalf("s2c mismatch:\n got %x\nwant %x", keys.ServerToClient, wantS2C)
	}
}

// TestKAT_HybridDecapsulateRoundtrip: load the client's deterministic hybrid key from the
// vector, then run Decapsulate against the server's ciphertext. The derived shared secrets must
// match x25519_ss / kyber_ss in the vector.
func TestKAT_HybridDecapsulateRoundtrip(t *testing.T) {
	v := loadVectors(t)
	if v == nil {
		return
	}
	xPriv := mustHex32(t, v.ClientX25519Priv)
	// We don't ship the ml-kem seed in the JSON directly (the export tool uses a fixed seed and
	// stores only the expanded private key for diagnostics). Instead, reconstruct from the seed
	// the export tool documents — match the literal bytes in tools/export-kat/src/main.rs.
	var seed [64]byte
	copy(seed[:32], []byte("AURA-MLKEM-DSEED-CLIENT--FIXED32"))
	copy(seed[32:], []byte("AURA-MLKEM-ZSEED-CLIENT--FIXED32"))
	priv, pub, err := NewHybridPrivateFromBytes(xPriv, seed)
	if err != nil {
		t.Fatalf("rebuild hybrid: %v", err)
	}
	// Sanity: the recomputed encapsulation key must match what the Rust side emitted.
	if !bytes.Equal(pub.MLKEM, mustHex(t, v.ClientKyberPub)) {
		t.Fatalf("ml-kem ek mismatch: Go and Rust derive different bytes from the same seed")
	}
	if !bytes.Equal(pub.X25519[:], mustHex(t, v.ClientX25519Pub)) {
		t.Fatalf("x25519 pub mismatch")
	}
	// Decapsulate.
	ct := &HybridCiphertext{MLKEMCT: mustHex(t, v.ServerKyberCt)}
	copy(ct.X25519Eph[:], mustHex(t, v.ServerX25519EphPub))
	ss, err := priv.Decapsulate(ct)
	if err != nil {
		t.Fatalf("decapsulate: %v", err)
	}
	if !bytes.Equal(ss.X25519SS[:], mustHex(t, v.X25519SS)) {
		t.Fatalf("x25519_ss mismatch:\n got %x\nwant %s", ss.X25519SS, v.X25519SS)
	}
	if !bytes.Equal(ss.MLKEMSS[:], mustHex(t, v.KyberSS)) {
		t.Fatalf("kyber_ss mismatch:\n got %x\nwant %s", ss.MLKEMSS, v.KyberSS)
	}
}

// TestKAT_ClientFinishedHMAC: HMAC-SHA256(c2s, transcript_hash) reproduces the Rust value.
func TestKAT_ClientFinishedHMAC(t *testing.T) {
	v := loadVectors(t)
	if v == nil {
		return
	}
	key := mustHex(t, v.SessionKeys.C2S)
	transcript := mustHex(t, v.TranscriptHash)
	mac := hmac.New(sha256.New, key)
	mac.Write(transcript)
	got := mac.Sum(nil)
	want := mustHex(t, v.ClientFinishedHmac)
	if !bytes.Equal(got, want) {
		t.Fatalf("client finished mismatch:\n got %x\nwant %x", got, want)
	}
}

// TestKAT_ServerFinishedHMAC: HMAC-SHA256(s2c, transcript_hash) reproduces the Rust value.
func TestKAT_ServerFinishedHMAC(t *testing.T) {
	v := loadVectors(t)
	if v == nil {
		return
	}
	key := mustHex(t, v.SessionKeys.S2C)
	transcript := mustHex(t, v.TranscriptHash)
	mac := hmac.New(sha256.New, key)
	mac.Write(transcript)
	got := mac.Sum(nil)
	want := mustHex(t, v.ServerFinishedHmac)
	if !bytes.Equal(got, want) {
		t.Fatalf("server finished mismatch:\n got %x\nwant %x", got, want)
	}
}

// TestKAT_SealedDatagramRecord: ChaCha20-Poly1305.Seal under the c2s key at seq 2 with
// aad=seq_be reproduces the exact sealed_record bytes (seq_be || ciphertext).
func TestKAT_SealedDatagramRecord(t *testing.T) {
	v := loadVectors(t)
	if v == nil {
		return
	}
	key, err := NewAeadKey(mustHex(t, v.DatagramTest.Key))
	if err != nil {
		t.Fatal(err)
	}
	frameBytes := mustHex(t, v.DatagramTest.Frame)
	seq := v.DatagramTest.Seq
	var seqBE [8]byte
	binary.BigEndian.PutUint64(seqBE[:], seq)
	ct := key.Seal(seq, frameBytes, seqBE[:])
	got := append(append([]byte{}, seqBE[:]...), ct...)
	want := mustHex(t, v.DatagramTest.SealedRecord)
	if !bytes.Equal(got, want) {
		t.Fatalf("sealed datagram mismatch:\n got %x\nwant %x", got, want)
	}
	// Round-trip: opening at the same seq must return the original frame bytes.
	pt, err := key.Open(seq, ct, seqBE[:])
	if err != nil {
		t.Fatalf("open: %v", err)
	}
	if !bytes.Equal(pt, frameBytes) {
		t.Fatal("open returned different plaintext")
	}
}

// TestKAT_KnockToken: HMAC-SHA256(ca_fp, u64_be(minute))[:16] matches the Rust knock value.
func TestKAT_KnockToken(t *testing.T) {
	v := loadVectors(t)
	if v == nil {
		return
	}
	key := mustHex(t, v.KnockTest.CAFingerprint)
	var mb [8]byte
	binary.BigEndian.PutUint64(mb[:], v.KnockTest.UnixMinute)
	mac := hmac.New(sha256.New, key)
	mac.Write(mb[:])
	tag := mac.Sum(nil)
	if len(tag) < 16 {
		t.Fatalf("hmac too short: %d", len(tag))
	}
	got := tag[:16]
	want := mustHex(t, v.KnockTest.Knock)
	if !bytes.Equal(got, want) {
		t.Fatalf("knock mismatch:\n got %x\nwant %x", got, want)
	}
}

// TestNonceLayout: explicit sanity that NonceFor matches the documented LE(u64) || 0x00000000.
func TestNonceLayout(t *testing.T) {
	if got := NonceFor(0); got != ([NonceLen]byte{}) {
		t.Fatalf("counter 0: want zero, got %x", got)
	}
	n := NonceFor(0x0807060504030201)
	if !bytes.Equal(n[:8], []byte{1, 2, 3, 4, 5, 6, 7, 8}) {
		t.Fatalf("LE layout wrong: %x", n[:8])
	}
	if !bytes.Equal(n[8:], []byte{0, 0, 0, 0}) {
		t.Fatalf("upper 4 bytes not zero: %x", n[8:])
	}
}

// TestAeadSessionCounterMonotonic: Seal/Open lock-step advances the counter by exactly 1.
func TestAeadSessionCounterMonotonic(t *testing.T) {
	key := make([]byte, 32)
	for i := range key {
		key[i] = byte(i)
	}
	s, err := NewAeadSession(key)
	if err != nil {
		t.Fatal(err)
	}
	if s.Counter() != 0 {
		t.Fatalf("initial counter %d", s.Counter())
	}
	for want := uint64(1); want <= 5; want++ {
		_ = s.Seal([]byte("x"), nil)
		if s.Counter() != want {
			t.Fatalf("after %d seals: counter %d", want, s.Counter())
		}
	}
}