package encoder

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"os"
	"strings"
	"testing"

	"golang.org/x/crypto/nacl/box"
)

func testKeyring(t *testing.T) string {
	t.Helper()
	path := t.TempDir() + "/pubring.mix"
	var content strings.Builder
	for i, name := range []string{"entry", "middle", "exit"} {
		keyID := strings.Repeat(string(rune('a'+i)), 32)
		publicKey := strings.Repeat(string(rune('d'+i)), 64)
		content.WriteString(name + " " + name + "@example.org " + keyID + " 4:0.2c E 2025-01-01 2099-12-31\n\n")
		content.WriteString("-----Begin Mix Key-----\n" + keyID + "\n" + publicKey + "\n-----End Mix Key-----\n")
	}
	if err := os.WriteFile(path, []byte(content.String()), 0600); err != nil {
		t.Fatal(err)
	}
	return path
}

func TestValidateEmail(t *testing.T) {
	r := Request{Kind: Email, PublicKeyring: "/tmp/pubring.mix", Entry: "entry", Chain: []string{"entry"}, To: "user@example.org", Body: "hello"}
	if err := Validate(r); err != nil {
		t.Fatal(err)
	}
}

func TestValidateRejectsPlainIncompleteRequest(t *testing.T) {
	r := Request{Kind: Email, PublicKeyring: "/tmp/pubring.mix", Entry: "entry", Chain: []string{"entry"}, To: "bad", Body: "hello"}
	if err := Validate(r); err == nil {
		t.Fatal("expected invalid recipient")
	}
}

func TestValidateUsenetRequiresGatewayRecipient(t *testing.T) {
	r := Request{Kind: Usenet, PublicKeyring: "/tmp/pubring.mix", Chain: []string{"entry"}, Newsgroup: "misc.test", Body: "hello"}
	if err := Validate(r); err == nil {
		t.Fatal("expected missing Usenet gateway recipient to be rejected")
	}
	r.To = "mail2news@example.org"
	if err := Validate(r); err != nil {
		t.Fatalf("expected valid Usenet request: %v", err)
	}
}

func TestComposeUsenetGatewayHeaders(t *testing.T) {
	plain, err := composeMessage(Request{
		Kind: Usenet, To: "mail2news@example.org", Subject: "test", Newsgroup: "misc.test", Body: "hello",
		References: "<parent@example.org>",
	})
	if err != nil {
		t.Fatal(err)
	}
	text := string(plain)
	if !strings.Contains(text, "To: mail2news@example.org\n") ||
		!strings.Contains(text, "Newsgroups: misc.test\n") ||
		!strings.Contains(text, "References: <parent@example.org>\n") ||
		!strings.Contains(text, "In-Reply-To: <parent@example.org>\n") {
		t.Fatalf("missing Usenet delivery headers: %q", text)
	}
}

func TestComposeReferencesChainUsesLastIDAsParent(t *testing.T) {
	plain, err := composeMessage(Request{
		Kind: Usenet, To: "mail2news@example.org", Subject: "test", Newsgroup: "misc.test", Body: "hello",
		References: "  <root@example.org>   <parent@example.org>  ",
	})
	if err != nil {
		t.Fatal(err)
	}
	text := string(plain)
	if !strings.Contains(text, "References: <root@example.org> <parent@example.org>\n") ||
		!strings.Contains(text, "In-Reply-To: <parent@example.org>\n") {
		t.Fatalf("thread headers were not normalized: %q", text)
	}
}

func TestValidateReferences(t *testing.T) {
	tests := []struct {
		name       string
		references string
		valid      bool
	}{
		{name: "empty", valid: true},
		{name: "whitespace only", references: "   "},
		{name: "parent", references: "<parent@example.org>", valid: true},
		{name: "thread chain", references: "<root@example.org> <parent@example.org>", valid: true},
		{name: "missing brackets", references: "parent@example.org"},
		{name: "missing local part", references: "<@example.org>"},
		{name: "multiple at signs", references: "<parent@example@org>"},
		{name: "embedded newline", references: "<root@example.org>\n<parent@example.org>"},
		{name: "header label", references: "References: <parent@example.org>"},
		{name: "too long", references: "<" + strings.Repeat("a", maxReferencesBytes) + "@example.org>"},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			request := Request{
				Kind: Email, PublicKeyring: "/tmp/pubring.mix", Entry: "entry", Chain: []string{"entry"},
				To: "user@example.org", Body: "hello", References: test.references,
			}
			err := Validate(request)
			if test.valid && err != nil {
				t.Fatalf("expected valid References, got %v", err)
			}
			if !test.valid && err == nil {
				t.Fatal("expected invalid References to be rejected")
			}
		})
	}
}

func TestValidateReplyTo(t *testing.T) {
	request := Request{
		Kind: Email, PublicKeyring: "/tmp/pubring.mix", Entry: "entry", Chain: []string{"entry"},
		To: "user@example.org", Body: "hello", ReplyTo: "Pseudonym <reply@example.org>",
	}
	if err := Validate(request); err != nil {
		t.Fatalf("expected valid Reply-To, got %v", err)
	}
	request.ReplyTo = "not an address"
	if err := Validate(request); err == nil {
		t.Fatal("expected invalid Reply-To to be rejected")
	}
}

func TestEncodeProducesYAMNArmor(t *testing.T) {
	keyring := testKeyring(t)
	result, err := Encode(Request{
		Kind: Email, PublicKeyring: keyring, Entry: "entry", Chain: []string{"entry"},
		From: "Anonymous <anon@example.org>", To: "user@example.org", Subject: "test", Body: "hello",
	})
	if err != nil {
		t.Fatal(err)
	}
	if result.EntryAddress != "entry@example.org" {
		t.Fatalf("unexpected entry address: %q", result.EntryAddress)
	}
	text := string(result.Envelope)
	if !strings.Contains(text, "-----BEGIN REMAILER MESSAGE-----") || !strings.Contains(text, "-----END REMAILER MESSAGE-----") {
		t.Fatal("missing YAMN armor markers")
	}
	lines := strings.Split(text, "\n")
	start := 0
	for i, line := range lines {
		if line == "-----BEGIN REMAILER MESSAGE-----" {
			start = i + 3
			break
		}
	}
	var encoded strings.Builder
	for _, line := range lines[start:] {
		if line == "" || strings.HasPrefix(line, "-----END") {
			break
		}
		encoded.WriteString(line)
	}
	packet, err := base64.StdEncoding.DecodeString(encoded.String())
	if err != nil {
		t.Fatal(err)
	}
	if len(packet) != messageBytes {
		t.Fatalf("unexpected packet size: got %d, want %d", len(packet), messageBytes)
	}
}

func TestEncodeMultiHop(t *testing.T) {
	result, err := Encode(Request{
		Kind: Email, PublicKeyring: testKeyring(t), Entry: "entry",
		Chain: []string{"entry", "middle", "exit"}, To: "user@example.org", Body: "hello",
	})
	if err != nil {
		t.Fatal(err)
	}
	if len(result.Envelope) == 0 || result.EntryAddress != "entry@example.org" {
		t.Fatal("multi-hop envelope was not produced")
	}
}

func TestEncodeMultiHopDecodesAtEveryRemailer(t *testing.T) {
	keyring, secretKeys := testDecodeKeyring(t)
	request := Request{
		Kind: Email, PublicKeyring: keyring, Entry: "entry",
		Chain: []string{"entry", "middle", "exit"}, From: "Anonymous <anon@example.org>",
		To: "user@example.org", Subject: "three-hop compatibility", Body: "hello through three remailers",
	}
	result, err := Encode(request)
	if err != nil {
		t.Fatal(err)
	}
	packet := decodeArmoredPacket(t, result.Envelope)
	wantPlain, err := composeMessage(request)
	if err != nil {
		t.Fatal(err)
	}

	for hop, name := range request.Chain {
		data := openHeaderForTest(t, packet[:headerBytes], secretKeys[name])
		if data[0] != 2 {
			t.Fatalf("hop %d (%s): unexpected packet version %d", hop, name, data[0])
		}
		if !bytes.Equal(data[117:149], antiTag(packet)) {
			t.Fatalf("hop %d (%s): anti-tag mismatch", hop, name)
		}

		switch data[1] {
		case 0:
			if hop == len(request.Chain)-1 {
				t.Fatalf("hop %d (%s): exit encoded as intermediate", hop, name)
			}
			shiftHeadersUpForTest(packet)
			encryptAll(packet, data[19:51], data[53:65])
		case 1:
			if hop != len(request.Chain)-1 {
				t.Fatalf("hop %d (%s): intermediate encoded as exit", hop, name)
			}
			bodyLength := int(binary.LittleEndian.Uint32(data[87:91]))
			if bodyLength != len(wantPlain) {
				t.Fatalf("exit body length: got %d, want %d", bodyLength, len(wantPlain))
			}
			body := aesCTR(packet[maxChainLength*headerBytes:], data[19:51], data[53:69])
			if !bytes.Equal(body[:bodyLength], wantPlain) {
				t.Fatal("exit plaintext does not match the encoded message")
			}
		default:
			t.Fatalf("hop %d (%s): unknown packet type %d", hop, name, data[1])
		}
	}
}

func testDecodeKeyring(t *testing.T) (string, map[string]*[32]byte) {
	t.Helper()
	path := t.TempDir() + "/pubring.mix"
	secretKeys := make(map[string]*[32]byte)
	var content strings.Builder
	for i, name := range []string{"entry", "middle", "exit"} {
		publicKey, secretKey, err := box.GenerateKey(rand.Reader)
		if err != nil {
			t.Fatal(err)
		}
		keyID := bytes.Repeat([]byte{byte(i + 1)}, 16)
		content.WriteString(name + " " + name + "@example.org " + hex.EncodeToString(keyID) + " 4:0.2c E 2025-01-01 2099-12-31\n\n")
		content.WriteString("-----Begin Mix Key-----\n" + hex.EncodeToString(keyID) + "\n" + hex.EncodeToString(publicKey[:]) + "\n-----End Mix Key-----\n")
		secretKeys[name] = secretKey
	}
	if err := os.WriteFile(path, []byte(content.String()), 0600); err != nil {
		t.Fatal(err)
	}
	return path, secretKeys
}

func decodeArmoredPacket(t *testing.T, envelope []byte) []byte {
	t.Helper()
	lines := strings.Split(string(envelope), "\n")
	start := -1
	for i, line := range lines {
		if line == "-----BEGIN REMAILER MESSAGE-----" {
			start = i + 3
			break
		}
	}
	if start < 0 {
		t.Fatal("missing armored packet start")
	}
	var encoded strings.Builder
	for _, line := range lines[start:] {
		if line == "" || strings.HasPrefix(line, "-----END") {
			break
		}
		encoded.WriteString(line)
	}
	packet, err := base64.StdEncoding.DecodeString(encoded.String())
	if err != nil {
		t.Fatal(err)
	}
	if len(packet) != messageBytes {
		t.Fatalf("unexpected packet size: got %d, want %d", len(packet), messageBytes)
	}
	return packet
}

func openHeaderForTest(t *testing.T, header []byte, recipientSecret *[32]byte) []byte {
	t.Helper()
	var senderPublic [32]byte
	copy(senderPublic[:], header[16:48])
	var nonce [24]byte
	copy(nonce[:], header[48:72])
	data, ok := box.Open(nil, header[72:248], &nonce, &senderPublic, recipientSecret)
	if !ok {
		t.Fatal("header authentication failed")
	}
	if len(data) != encHeadBytes {
		t.Fatalf("unexpected decoded header size: got %d, want %d", len(data), encHeadBytes)
	}
	return data
}

func shiftHeadersUpForTest(packet []byte) {
	headersEnd := maxChainLength * headerBytes
	copy(packet, packet[headerBytes:headersEnd])
	clear(packet[headersEnd-headerBytes : headersEnd])
}