package nip44

import (
	"bytes"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"os"
	"strings"
	"testing"

	"git.smesh.lol/nostrino/pkg/crypto/secp256k1"
)

// The test vectors are the official NIP-44 v2 vectors published at
// https://github.com/paulmillr/nip44/blob/main/nip44.vectors.json, whose
// SHA-256 is 269ed0f69e4c192512cc779e78c555090cebc7c785b609e338a62afc3ce25040
// as stated in the NIP.
const vectorsPath = "testdata/nip44.vectors.json"

type conversationKeyVector struct {
	Sec1            string `json:"sec1"`
	Pub2            string `json:"pub2"`
	ConversationKey string `json:"conversation_key"`
}

type messageKeysVector struct {
	Nonce       string `json:"nonce"`
	ChaChaKey   string `json:"chacha_key"`
	ChaChaNonce string `json:"chacha_nonce"`
	HMACKey     string `json:"hmac_key"`
}

type encryptDecryptVector struct {
	Sec1            string `json:"sec1"`
	Sec2            string `json:"sec2"`
	ConversationKey string `json:"conversation_key"`
	Nonce           string `json:"nonce"`
	Plaintext       string `json:"plaintext"`
	Payload         string `json:"payload"`
}

type longMessageVector struct {
	ConversationKey string `json:"conversation_key"`
	Nonce           string `json:"nonce"`
	Pattern         string `json:"pattern"`
	Repeat          int    `json:"repeat"`
	PlaintextSHA256 string `json:"plaintext_sha256"`
	PayloadSHA256   string `json:"payload_sha256"`
}

type invalidDecryptVector struct {
	ConversationKey string `json:"conversation_key"`
	Nonce           string `json:"nonce"`
	Plaintext       string `json:"plaintext"`
	Payload         string `json:"payload"`
	Note            string `json:"note"`
}

type invalidConversationKeyVector struct {
	Sec1 string `json:"sec1"`
	Pub2 string `json:"pub2"`
	Note string `json:"note"`
}

type vectors struct {
	V2 struct {
		Valid struct {
			GetConversationKey []conversationKeyVector `json:"get_conversation_key"`
			GetMessageKeys     struct {
				ConversationKey string              `json:"conversation_key"`
				Keys            []messageKeysVector `json:"keys"`
			} `json:"get_message_keys"`
			CalcPaddedLen         [][2]int               `json:"calc_padded_len"`
			EncryptDecrypt        []encryptDecryptVector `json:"encrypt_decrypt"`
			EncryptDecryptLongMsg []longMessageVector    `json:"encrypt_decrypt_long_msg"`
		} `json:"valid"`
		Invalid struct {
			EncryptMsgLengths  []int                          `json:"encrypt_msg_lengths"`
			GetConversationKey []invalidConversationKeyVector `json:"get_conversation_key"`
			Decrypt            []invalidDecryptVector         `json:"decrypt"`
		} `json:"invalid"`
	} `json:"v2"`
}

func loadVectors(t *testing.T) vectors {
	t.Helper()
	data, err := os.ReadFile(vectorsPath)
	if err != nil {
		t.Fatalf("read vectors: %v", err)
	}
	sum := sha256.Sum256(data)
	const wantSum = "269ed0f69e4c192512cc779e78c555090cebc7c785b609e338a62afc3ce25040"
	if got := hex.EncodeToString(sum[:]); got != wantSum {
		t.Fatalf("test vector checksum = %s, want %s", got, wantSum)
	}
	var v vectors
	if err := json.Unmarshal(data, &v); err != nil {
		t.Fatalf("parse vectors: %v", err)
	}
	return v
}

func mustHex(t *testing.T, s string, size int) []byte {
	t.Helper()
	b, err := hex.DecodeString(s)
	if err != nil {
		t.Fatalf("bad hex %q: %v", s, err)
	}
	if size >= 0 && len(b) != size {
		t.Fatalf("hex %q has length %d, want %d", s, len(b), size)
	}
	return b
}

func bytes32(t *testing.T, s string) [32]byte {
	t.Helper()
	var out [32]byte
	copy(out[:], mustHex(t, s, 32))
	return out
}

func TestOfficialVectorChecksum(t *testing.T) {
	// loadVectors verifies the checksum; this test makes the check appear in
	// the verbose test list and guards against a stale file.
	v := loadVectors(t)
	if len(v.V2.Valid.EncryptDecrypt) == 0 {
		t.Fatal("no encrypt_decrypt vectors")
	}
	if len(v.V2.Valid.EncryptDecrypt) < 10 || len(v.V2.Valid.GetConversationKey) < 35 {
		t.Fatalf("unexpected vector counts: %d/%d", len(v.V2.Valid.EncryptDecrypt), len(v.V2.Valid.GetConversationKey))
	}
}

func TestOfficialConversationKeys(t *testing.T) {
	v := loadVectors(t)
	for i, vec := range v.V2.Valid.GetConversationKey {
		sec := bytes32(t, vec.Sec1)
		pub := bytes32(t, vec.Pub2)
		got, err := ConversationKey(sec, pub)
		if err != nil {
			t.Fatalf("vector %d (%s): unexpected error: %v", i, vec.ConversationKey, err)
		}
		want := bytes32(t, vec.ConversationKey)
		if got != want {
			t.Fatalf("vector %d: conversation key = %x, want %x", i, got, want)
		}
	}
}

func TestOfficialMessageKeys(t *testing.T) {
	v := loadVectors(t)
	conv := bytes32(t, v.V2.Valid.GetMessageKeys.ConversationKey)
	if len(v.V2.Valid.GetMessageKeys.Keys) != 32 {
		t.Fatalf("got %d message key vectors, want 32", len(v.V2.Valid.GetMessageKeys.Keys))
	}
	for i, vec := range v.V2.Valid.GetMessageKeys.Keys {
		nonce := bytes32(t, vec.Nonce)
		key, chachaNonce, hmacKey := messageKeys(conv, nonce)
		if want := mustHex(t, vec.ChaChaKey, 32); !bytes.Equal(key[:], want) {
			t.Fatalf("vector %d: chacha key = %x, want %x", i, key, want)
		}
		if want := mustHex(t, vec.ChaChaNonce, 12); !bytes.Equal(chachaNonce[:], want) {
			t.Fatalf("vector %d: chacha nonce = %x, want %x", i, chachaNonce, want)
		}
		if want := mustHex(t, vec.HMACKey, 32); !bytes.Equal(hmacKey[:], want) {
			t.Fatalf("vector %d: hmac key = %x, want %x", i, hmacKey, want)
		}
	}
}

func TestOfficialCalcPaddedLen(t *testing.T) {
	v := loadVectors(t)
	if len(v.V2.Valid.CalcPaddedLen) != 24 {
		t.Fatalf("got %d padding vectors, want 24", len(v.V2.Valid.CalcPaddedLen))
	}
	for _, vec := range v.V2.Valid.CalcPaddedLen {
		got := calcPaddedLen(uint64(vec[0]))
		if got != uint64(vec[1]) {
			t.Fatalf("calcPaddedLen(%d) = %d, want %d", vec[0], got, vec[1])
		}
	}
}

func TestOfficialEncryptDecrypt(t *testing.T) {
	v := loadVectors(t)
	for i, vec := range v.V2.Valid.EncryptDecrypt {
		sec1 := bytes32(t, vec.Sec1)
		sec2 := bytes32(t, vec.Sec2)
		pub1, ok := secp256k1.PubKey(sec1)
		if !ok {
			t.Fatalf("vector %d: bad sec1", i)
		}
		pub2, ok := secp256k1.PubKey(sec2)
		if !ok {
			t.Fatalf("vector %d: bad sec2", i)
		}
		conv := bytes32(t, vec.ConversationKey)

		convA, err := ConversationKey(sec1, pub2)
		if err != nil {
			t.Fatalf("vector %d: conv(1, pub2): %v", i, err)
		}
		if convA != conv {
			t.Fatalf("vector %d: conv(1, pub2) = %x, want %x", i, convA, conv)
		}
		convB, err := ConversationKey(sec2, pub1)
		if err != nil {
			t.Fatalf("vector %d: conv(2, pub1): %v", i, err)
		}
		if convB != conv {
			t.Fatalf("vector %d: conv(2, pub1) = %x, want %x", i, convB, conv)
		}

		nonce := bytes32(t, vec.Nonce)
		payload, err := EncryptWithNonce(conv, nonce, []byte(vec.Plaintext))
		if err != nil {
			t.Fatalf("vector %d: encrypt: %v", i, err)
		}
		if payload != vec.Payload {
			t.Fatalf("vector %d: payload = %s, want %s", i, payload, vec.Payload)
		}
		got, err := Decrypt(conv, vec.Payload)
		if err != nil {
			t.Fatalf("vector %d: decrypt: %v", i, err)
		}
		if string(got) != vec.Plaintext {
			t.Fatalf("vector %d: plaintext = %q, want %q", i, got, vec.Plaintext)
		}
		// The other direction must decrypt with the same conversation key.
		got, err = Decrypt(convB, vec.Payload)
		if err != nil || string(got) != vec.Plaintext {
			t.Fatalf("vector %d: reverse decrypt: %q, %v", i, got, err)
		}
	}
}

func TestOfficialLongMessages(t *testing.T) {
	v := loadVectors(t)
	if len(v.V2.Valid.EncryptDecryptLongMsg) != 3 {
		t.Fatalf("got %d long message vectors, want 3", len(v.V2.Valid.EncryptDecryptLongMsg))
	}
	for i, vec := range v.V2.Valid.EncryptDecryptLongMsg {
		conv := bytes32(t, vec.ConversationKey)
		nonce := bytes32(t, vec.Nonce)
		plaintext := []byte(strings.Repeat(vec.Pattern, vec.Repeat))
		plainSum := sha256.Sum256(plaintext)
		if got, want := hex.EncodeToString(plainSum[:]), vec.PlaintextSHA256; got != want {
			t.Fatalf("vector %d: plaintext sha256 = %s, want %s", i, got, want)
		}
		payload, err := EncryptWithNonce(conv, nonce, plaintext)
		if err != nil {
			t.Fatalf("vector %d: encrypt: %v", i, err)
		}
		payloadSum := sha256.Sum256([]byte(payload))
		if got, want := hex.EncodeToString(payloadSum[:]), vec.PayloadSHA256; got != want {
			t.Fatalf("vector %d: payload sha256 = %s, want %s", i, got, want)
		}
		got, err := Decrypt(conv, payload)
		if err != nil {
			t.Fatalf("vector %d: decrypt: %v", i, err)
		}
		if !bytes.Equal(got, plaintext) {
			t.Fatalf("vector %d: plaintext mismatch", i)
		}
	}
}

// extendedPrefixVectors are the extended length prefix checks from the NIP
// itself (plaintext is "a" repeated plaintextLen times). Payloads are too
// large to embed, so only SHA-256 checksums are published; paddedLen excludes
// the length prefix.
var extendedPrefixVectors = []struct {
	plaintextLen    int
	paddedLen       int
	plaintextSHA256 string
	payloadSHA256   string
}{
	{65535, 65536, "6e1bebca6a8229364a162a72ef064826c4cd7457bf54f190ef782bd9deff3e42", "6d8c2810d1e870fbaa1f0a0937126cca837a15f9260e27060c331d70a3c0bc84"},
	{65536, 65536, "bf718b6f653bebc184e1479f1935b8da974d701b893afcf49e701f3e2f9f9c5a", "b7b4edb36ba92e267d322d56d9aebc22e7fa96ff52e3c12adc07f07a43cbc616"},
	{65537, 81920, "008ffc88d3c96a9f307524eb361e47c5222a887fc45fa0c1fb8d429c5c23b430", "eeb7c7c5373894ea2c1547cfd3ccb15d5a0b2d619da852e5c79df792dcc9e435"},
}

func TestExtendedPrefixVectors(t *testing.T) {
	const convHex = "c41c775356fd92eadc63ff5a0dc1da211b268cbea22316767095b2871ea1412d"
	const nonceHex = "0000000000000000000000000000000000000000000000000000000000000001"
	conv := bytes32(t, convHex)
	nonce := bytes32(t, nonceHex)

	for _, vec := range extendedPrefixVectors {
		plaintext := bytes.Repeat([]byte{'a'}, vec.plaintextLen)
		if got := hex.EncodeToString(sha256Sum(plaintext)); got != vec.plaintextSHA256 {
			t.Fatalf("len %d: plaintext sha256 = %s, want %s", vec.plaintextLen, got, vec.plaintextSHA256)
		}
		if got, want := calcPaddedLen(uint64(vec.plaintextLen)), uint64(vec.paddedLen); got != want {
			t.Fatalf("len %d: calcPaddedLen = %d, want %d", vec.plaintextLen, got, want)
		}
		payload, err := EncryptWithNonce(conv, nonce, plaintext)
		if err != nil {
			t.Fatalf("len %d: encrypt: %v", vec.plaintextLen, err)
		}
		if got := hex.EncodeToString(sha256Sum([]byte(payload))); got != vec.payloadSHA256 {
			t.Fatalf("len %d: payload sha256 = %s, want %s", vec.plaintextLen, got, vec.payloadSHA256)
		}
		// Check the wire sizes documented in the NIP.
		raw, err := base64.StdEncoding.DecodeString(payload)
		if err != nil {
			t.Fatalf("len %d: decode payload: %v", vec.plaintextLen, err)
		}
		prefixLen := 2
		if vec.plaintextLen >= 65536 {
			prefixLen = 6
		}
		if want := 1 + 32 + prefixLen + vec.paddedLen + 32; len(raw) != want {
			t.Fatalf("len %d: raw payload = %d bytes, want %d", vec.plaintextLen, len(raw), want)
		}
		got, err := Decrypt(conv, payload)
		if err != nil {
			t.Fatalf("len %d: decrypt: %v", vec.plaintextLen, err)
		}
		if !bytes.Equal(got, plaintext) {
			t.Fatalf("len %d: plaintext mismatch", vec.plaintextLen)
		}
	}
}

func TestOfficialInvalidConversationKeys(t *testing.T) {
	v := loadVectors(t)
	if len(v.V2.Invalid.GetConversationKey) != 8 {
		t.Fatalf("got %d invalid conversation key vectors, want 8", len(v.V2.Invalid.GetConversationKey))
	}
	for i, vec := range v.V2.Invalid.GetConversationKey {
		var sec, pub [32]byte
		copy(sec[:], mustHex(t, vec.Sec1, 32))
		copy(pub[:], mustHex(t, vec.Pub2, 32))
		if _, err := ConversationKey(sec, pub); err == nil {
			t.Fatalf("vector %d (%s): expected an error", i, vec.Note)
		}
	}
}

func TestOfficialInvalidDecrypts(t *testing.T) {
	v := loadVectors(t)
	if len(v.V2.Invalid.Decrypt) != 12 {
		t.Fatalf("got %d invalid decrypt vectors, want 12", len(v.V2.Invalid.Decrypt))
	}
	for i, vec := range v.V2.Invalid.Decrypt {
		conv := bytes32(t, vec.ConversationKey)
		if _, err := Decrypt(conv, vec.Payload); err == nil {
			t.Fatalf("vector %d (%s): expected an error", i, vec.Note)
		}
	}
}

// TestOfficialInvalidEncryptLengths covers invalid.encrypt_msg_lengths. The
// published vector file predates the 6-byte extended length prefix and still
// lists 65536, 100000 and 10000000, so only zero (an invalid empty plaintext)
// is checked here; extended lengths are covered by the NIP's own checksums in
// TestExtendedPrefixVectors.
func TestOfficialInvalidEncryptLengths(t *testing.T) {
	v := loadVectors(t)
	if len(v.V2.Invalid.EncryptMsgLengths) == 0 {
		t.Fatal("no invalid encrypt length vectors")
	}
	var conv [32]byte
	conv[0] = 1
	for _, n := range v.V2.Invalid.EncryptMsgLengths {
		if n != 0 {
			continue
		}
		if _, err := EncryptWithNonce(conv, [32]byte{}, make([]byte, n)); err == nil {
			t.Fatalf("encrypting %d bytes: expected an error", n)
		}
	}
}

func sha256Sum(b []byte) []byte {
	s := sha256.Sum256(b)
	return s[:]
}

func TestConversationKeySymmetry(t *testing.T) {
	pairs := [][2]string{
		{"0000000000000000000000000000000000000000000000000000000000000001", "0000000000000000000000000000000000000000000000000000000000000002"},
		{"5c0c523f52a5b6fad39ed2403092df8cebc36318b39383bca6c00808626fab3a", "4b22aa260e4acb7021e32f38a6cdf4b673c6a277755bfce287e370c924dc936d"},
		{"7a1c0b0f9b3f4e2d1c5a6b7c8d9e0f1a2b3c4d5e6f708192a3b4c5d6e7f80912", "112233445566778899aabbccddeeff00112233445566778899aabbccddeeff00"},
	}
	for _, pair := range pairs {
		secA := bytes32(t, pair[0])
		secB := bytes32(t, pair[1])
		pubA, ok := secp256k1.PubKey(secA)
		if !ok {
			t.Fatalf("bad secret a %s", pair[0])
		}
		pubB, ok := secp256k1.PubKey(secB)
		if !ok {
			t.Fatalf("bad secret b %s", pair[1])
		}
		k1, err := ConversationKey(secA, pubB)
		if err != nil {
			t.Fatal(err)
		}
		k2, err := ConversationKey(secB, pubA)
		if err != nil {
			t.Fatal(err)
		}
		if k1 != k2 {
			t.Fatalf("conversation keys differ: %x %x", k1, k2)
		}
	}
}

func TestConversationKeyRejectsBadInput(t *testing.T) {
	var zero, one [32]byte
	one[31] = 1
	if _, err := ConversationKey(zero, one); err == nil {
		t.Fatal("accepted a zero secret key")
	}
	pubOne, ok := secp256k1.PubKey(one)
	if !ok {
		t.Fatal("pubkey")
	}
	offCurve := [32]byte{0xff}
	if _, err := ConversationKey(one, offCurve); err == nil {
		t.Fatal("accepted an off-curve public key")
	}
	if _, err := ConversationKey(one, pubOne); err != nil {
		t.Fatalf("self ECDH: %v", err)
	}
}

func TestRoundTrip(t *testing.T) {
	conv := bytes32(t, "c41c775356fd92eadc63ff5a0dc1da211b268cbea22316767095b2871ea1412d")
	lengths := []int{1, 31, 32, 33, 100, 65535, 65536, 70000}
	for _, n := range lengths {
		plaintext := bytes.Repeat([]byte{'z'}, n)
		for i := range plaintext {
			plaintext[i] = byte(i * 31)
		}

		// Deterministic nonce.
		var nonce [32]byte
		nonce[0] = byte(n)
		nonce[1] = byte(n >> 8)
		nonce[31] = 0x5a
		payload, err := EncryptWithNonce(conv, nonce, plaintext)
		if err != nil {
			t.Fatalf("len %d: encrypt: %v", n, err)
		}
		got, err := Decrypt(conv, payload)
		if err != nil {
			t.Fatalf("len %d: decrypt: %v", n, err)
		}
		if !bytes.Equal(got, plaintext) {
			t.Fatalf("len %d: round trip mismatch", n)
		}

		// Random nonce through Encrypt.
		payload, err = Encrypt(conv, plaintext)
		if err != nil {
			t.Fatalf("len %d: Encrypt: %v", n, err)
		}
		got, err = Decrypt(conv, payload)
		if err != nil {
			t.Fatalf("len %d: Decrypt: %v", n, err)
		}
		if !bytes.Equal(got, plaintext) {
			t.Fatalf("len %d: random nonce round trip mismatch", n)
		}
	}
}

func TestEncryptRejectsEmptyPlaintext(t *testing.T) {
	var conv [32]byte
	if _, err := Encrypt(conv, nil); err == nil {
		t.Fatal("accepted nil plaintext")
	}
	if _, err := Encrypt(conv, []byte{}); err == nil {
		t.Fatal("accepted empty plaintext")
	}
}

// tamperDecode decodes a payload, applies mutate, and re-encodes it.
func tamperDecode(t *testing.T, payload string, mutate func([]byte) []byte) string {
	t.Helper()
	raw, err := base64.StdEncoding.DecodeString(payload)
	if err != nil {
		t.Fatal(err)
	}
	return base64.StdEncoding.EncodeToString(mutate(raw))
}

func TestDecryptTamper(t *testing.T) {
	conv := bytes32(t, "c41c775356fd92eadc63ff5a0dc1da211b268cbea22316767095b2871ea1412d")
	var nonce [32]byte
	nonce[0] = 7
	plaintext := bytes.Repeat([]byte("tamper me "), 32) // 320 bytes
	payload, err := EncryptWithNonce(conv, nonce, plaintext)
	if err != nil {
		t.Fatal(err)
	}
	if got, err := Decrypt(conv, payload); err != nil || !bytes.Equal(got, plaintext) {
		t.Fatalf("baseline decrypt failed: %v", err)
	}

	tests := []struct {
		name    string
		payload string
	}{
		{"modified MAC", tamperDecode(t, payload, func(raw []byte) []byte {
			raw[len(raw)-1] ^= 0x01
			return raw
		})},
		{"modified ciphertext", tamperDecode(t, payload, func(raw []byte) []byte {
			raw[50] ^= 0x80
			return raw
		})},
		{"truncated payload", payload[:len(payload)-4]},
		{"short payload", payload[:100]},
		{"truncated base64 decode", payload[:len(payload)-20]},
		{"hash prefix", "#" + payload[1:]},
		{"empty", ""},
		{"just hash", "#"},
		{"invalid base64", payload[:len(payload)-1] + "!"},
		{"bad version 0", tamperDecode(t, payload, func(raw []byte) []byte {
			raw[0] = 0
			return raw
		})},
		{"bad version 3", tamperDecode(t, payload, func(raw []byte) []byte {
			raw[0] = 3
			return raw
		})},
		{"bad version 1 with del", tamperDecode(t, payload, func(raw []byte) []byte {
			raw[0] = 1
			raw = append(raw[:1], raw[2:]...)
			return raw
		})},
		{"wrong conversation key", ""},
	}

	for _, tc := range tests {
		if tc.name == "wrong conversation key" {
			wrong := conv
			wrong[0] ^= 0xff
			if _, err := Decrypt(wrong, payload); err == nil {
				t.Fatal("wrong conversation key: expected an error")
			}
			continue
		}
		if _, err := Decrypt(conv, tc.payload); err == nil {
			t.Fatalf("%s: expected an error", tc.name)
		}
	}
}

func TestUnpad(t *testing.T) {
	// Valid small prefix.
	small := append([]byte{0, 5}, []byte("hello")...)
	small = append(small, make([]byte, 32-5)...)
	got, err := unpad(small)
	if err != nil || string(got) != "hello" {
		t.Fatalf("unpad small = %q, %v", got, err)
	}

	// Valid extended prefix: 65536 bytes at exactly the padded length.
	large := make([]byte, 6+65536)
	large[3] = 1 // 0x00010000 big-endian in bytes 2..6
	for i := 6; i < len(large); i++ {
		large[i] = 0x41
	}
	got, err = unpad(large)
	if err != nil || len(got) != 65536 || got[0] != 0x41 || got[65535] != 0x41 {
		t.Fatalf("unpad large: len=%d err=%v", len(got), err)
	}

	bad := []struct {
		name string
		in   []byte
	}{
		{"empty", nil},
		{"one byte", []byte{0}},
		{"zero length", []byte{0, 0, 0, 0, 0, 0}},
		{"short extended", []byte{0, 0, 0, 0}},
		{"extended below threshold", []byte{0, 0, 0, 0, 1, 0, 'a'}},
		{"length too small for buffer", []byte{0, 5, 'h', 'e', 'l', 'l', 'o'}},
		{"padded length wrong", append(append([]byte{0, 5}, []byte("hello")...), make([]byte, 10)...)},
		{"nonzero padding", append(append([]byte{0, 5}, []byte("hello")...), append([]byte{1}, make([]byte, 26)...)...)},
		{"u16 zero not followed by u32", []byte{0, 0}},
	}
	for _, tc := range bad {
		if _, err := unpad(tc.in); err == nil {
			t.Fatalf("%s: expected an error", tc.name)
		}
	}
}