package e2ee

import (
	"bytes"
	"encoding/json"
	"os"
	"reflect"
	"testing"
)

// vectors.json is generated from web/static/crypto.js by
// test/gen-e2ee-vectors.mjs — every blob here was produced by the code the
// browser actually runs. These tests are the cross-language contract.
type vectorFile struct {
	AESJSON []struct {
		Key       string          `json:"key"`
		Plaintext json.RawMessage `json:"plaintext"`
		Blob      string          `json:"blob"`
	} `json:"aesJSON"`
	AESBytes []struct {
		Key          string `json:"key"`
		PlaintextB64 string `json:"plaintextB64"`
		BlobB64      string `json:"blobB64"`
	} `json:"aesBytes"`
	Invites []struct {
		T               string `json:"t"`
		ID              string `json:"id"`
		ClaimSecret     string `json:"claimSecret"`
		ClaimHash       string `json:"claimHash"`
		GroupKey        string `json:"groupKey"`
		WrappedGroupKey string `json:"wrappedGroupKey"`
	} `json:"invites"`
	ThreadWraps []struct {
		PublicKey  string `json:"publicKey"`
		PrivateKey string `json:"privateKey"`
		ThreadKey  string `json:"threadKey"`
		Wrapped    string `json:"wrapped"`
	} `json:"threadWraps"`
	Seals []struct {
		PublicKey  string          `json:"publicKey"`
		PrivateKey string          `json:"privateKey"`
		Plaintext  json.RawMessage `json:"plaintext"`
		Sealed     string          `json:"sealed"`
	} `json:"seals"`
}

func loadVectors(t *testing.T) vectorFile {
	t.Helper()
	raw, err := os.ReadFile("testdata/vectors.json")
	if err != nil {
		t.Fatalf("read vectors (regenerate with `node test/gen-e2ee-vectors.mjs`): %v", err)
	}
	var v vectorFile
	if err := json.Unmarshal(raw, &v); err != nil {
		t.Fatalf("parse vectors: %v", err)
	}
	if len(v.AESJSON) == 0 || len(v.AESBytes) == 0 || len(v.Invites) == 0 ||
		len(v.ThreadWraps) == 0 || len(v.Seals) == 0 {
		t.Fatal("vectors.json is missing sections — regenerate it")
	}
	return v
}

// jsonEqual compares two JSON documents structurally, so JS/Go key order and
// number formatting differences don't matter.
func jsonEqual(t *testing.T, a, b []byte) bool {
	t.Helper()
	var av, bv any
	if err := json.Unmarshal(a, &av); err != nil {
		t.Fatalf("bad JSON: %v", err)
	}
	if err := json.Unmarshal(b, &bv); err != nil {
		t.Fatalf("bad JSON: %v", err)
	}
	return reflect.DeepEqual(av, bv)
}

func TestVectorsAESJSON(t *testing.T) {
	for i, c := range loadVectors(t).AESJSON {
		key, err := ImportKey(c.Key)
		if err != nil {
			t.Fatalf("case %d: %v", i, err)
		}
		var got json.RawMessage
		if err := DecryptJSON(key, c.Blob, &got); err != nil {
			t.Fatalf("case %d: decrypt JS blob: %v", i, err)
		}
		if !jsonEqual(t, got, c.Plaintext) {
			t.Errorf("case %d: got %s, want %s", i, got, c.Plaintext)
		}
	}
}

func TestVectorsAESBytes(t *testing.T) {
	for i, c := range loadVectors(t).AESBytes {
		key, err := ImportKey(c.Key)
		if err != nil {
			t.Fatalf("case %d: %v", i, err)
		}
		blob, err := B64uDecode(c.BlobB64)
		if err != nil {
			t.Fatalf("case %d: %v", i, err)
		}
		want, err := B64uDecode(c.PlaintextB64)
		if err != nil {
			t.Fatalf("case %d: %v", i, err)
		}
		got, err := DecryptBytes(key, blob)
		if err != nil {
			t.Fatalf("case %d: decrypt JS bytes: %v", i, err)
		}
		if !bytes.Equal(got, want) {
			t.Errorf("case %d: plaintext mismatch (%d bytes)", i, len(want))
		}
	}
}

// TestVectorsInvites is the strongest cross-check: deriveInvite is fully
// deterministic, so Go must reproduce the JS outputs exactly.
func TestVectorsInvites(t *testing.T) {
	for i, c := range loadVectors(t).Invites {
		inv, err := DeriveInvite(c.T)
		if err != nil {
			t.Fatalf("case %d: %v", i, err)
		}
		if inv.ID != c.ID {
			t.Errorf("case %d: id: got %s, want %s", i, inv.ID, c.ID)
		}
		if inv.ClaimSecret != c.ClaimSecret {
			t.Errorf("case %d: claimSecret mismatch", i)
		}
		if inv.ClaimHash != c.ClaimHash {
			t.Errorf("case %d: claimHash: got %s, want %s", i, inv.ClaimHash, c.ClaimHash)
		}
		gk, err := UnwrapGroupKey(c.WrappedGroupKey, inv.WrapKey)
		if err != nil {
			t.Fatalf("case %d: unwrap JS group key: %v", i, err)
		}
		if gk != c.GroupKey {
			t.Errorf("case %d: group key mismatch", i)
		}
	}
}

func TestVectorsThreadWraps(t *testing.T) {
	for i, c := range loadVectors(t).ThreadWraps {
		got, err := UnwrapThreadKey(c.Wrapped, c.PrivateKey)
		if err != nil {
			t.Fatalf("case %d: unwrap JS-wrapped thread key: %v", i, err)
		}
		if got != c.ThreadKey {
			t.Errorf("case %d: thread key mismatch", i)
		}
		// And the reverse direction: Go wraps to the JS-generated public key,
		// Go unwraps with the JS-exported pkcs8 private key. This exercises
		// Go's parse of WebCrypto's raw + pkcs8 exports on the wrap path too.
		rewrapped, err := WrapThreadKey(c.ThreadKey, c.PublicKey)
		if err != nil {
			t.Fatalf("case %d: wrap to JS pubkey: %v", i, err)
		}
		back, err := UnwrapThreadKey(rewrapped, c.PrivateKey)
		if err != nil {
			t.Fatalf("case %d: unwrap own wrap: %v", i, err)
		}
		if back != c.ThreadKey {
			t.Errorf("case %d: rewrap round-trip mismatch", i)
		}
	}
}

func TestVectorsSeals(t *testing.T) {
	for i, c := range loadVectors(t).Seals {
		var got json.RawMessage
		if err := OpenSealed(c.Sealed, c.PrivateKey, &got); err != nil {
			t.Fatalf("case %d: open JS seal: %v", i, err)
		}
		if !jsonEqual(t, got, c.Plaintext) {
			t.Errorf("case %d: got %s, want %s", i, got, c.Plaintext)
		}
	}
}

func TestRoundTripsAndWireShape(t *testing.T) {
	key, err := GenerateKey()
	if err != nil {
		t.Fatal(err)
	}
	blob, err := EncryptJSON(key, map[string]string{"text": "round trip"})
	if err != nil {
		t.Fatal(err)
	}
	var out map[string]string
	if err := DecryptJSON(key, blob, &out); err != nil {
		t.Fatal(err)
	}
	if out["text"] != "round trip" {
		t.Errorf("got %q", out["text"])
	}

	kp, err := GenerateUserKeypair()
	if err != nil {
		t.Fatal(err)
	}
	if raw, _ := B64uDecode(kp.PublicKey); len(raw) != 65 || raw[0] != 0x04 {
		t.Errorf("public key is not a 65-byte uncompressed P-256 point")
	}

	threadKey := B64uEncode(key)
	wrapped, err := WrapThreadKey(threadKey, kp.PublicKey)
	if err != nil {
		t.Fatal(err)
	}
	// The ECIES wire shape crypto.js expects: ephPub(65) || iv(12) || ct.
	if raw, _ := B64uDecode(wrapped); len(raw) < 65+12+len(key)+16 || raw[0] != 0x04 {
		t.Errorf("wrapped blob has wrong shape (len %d)", len(wrapped))
	}
	back, err := UnwrapThreadKey(wrapped, kp.PrivateKey)
	if err != nil {
		t.Fatal(err)
	}
	if back != threadKey {
		t.Error("thread key round trip mismatch")
	}

	// Wrong recipient must fail, not silently succeed.
	other, err := GenerateUserKeypair()
	if err != nil {
		t.Fatal(err)
	}
	if _, err := UnwrapThreadKey(wrapped, other.PrivateKey); err == nil {
		t.Error("unwrap with the wrong private key succeeded")
	}

	inv1, err := DeriveInvite("AAAAAAAAAAAAAAAAAAAAAA") // 16 zero bytes
	if err != nil {
		t.Fatal(err)
	}
	inv2, _ := DeriveInvite("AAAAAAAAAAAAAAAAAAAAAA")
	if inv1.ID != inv2.ID || inv1.ClaimHash != inv2.ClaimHash {
		t.Error("DeriveInvite is not deterministic")
	}
}