package encryption

import (
	"bytes"
	"encoding/binary"
	"io"
	"strings"
	"testing"
)

func testKey(t *testing.T) []byte {
	t.Helper()
	return DeriveKey("test-encryption-key-for-tcp-frames")
}

func TestEncryptedFrameRoundTrip(t *testing.T) {
	key := testKey(t)
	original := []byte("Hello, framed TCP tunnel!")

	var buf bytes.Buffer
	if err := WriteEncryptedFrame(&buf, original, key); err != nil {
		t.Fatalf("WriteEncryptedFrame failed: %v", err)
	}

	got, err := ReadEncryptedFrame(&buf, key)
	if err != nil {
		t.Fatalf("ReadEncryptedFrame failed: %v", err)
	}
	if !bytes.Equal(got, original) {
		t.Errorf("round trip mismatch: got %q want %q", got, original)
	}
}

func TestEncryptedFrameMultipleChunks(t *testing.T) {
	key := testKey(t)
	chunks := [][]byte{
		[]byte("chunk-one"),
		[]byte("chunk-two-is-longer"),
		[]byte{0x00, 0x01, 0xff},
	}

	var buf bytes.Buffer
	for _, c := range chunks {
		if err := WriteEncryptedFrame(&buf, c, key); err != nil {
			t.Fatalf("WriteEncryptedFrame failed: %v", err)
		}
	}

	for i, want := range chunks {
		got, err := ReadEncryptedFrame(&buf, key)
		if err != nil {
			t.Fatalf("ReadEncryptedFrame chunk %d failed: %v", i, err)
		}
		if !bytes.Equal(got, want) {
			t.Errorf("chunk %d mismatch: got %q want %q", i, got, want)
		}
	}
}

func TestEncryptedFrameRejectsZeroLength(t *testing.T) {
	key := testKey(t)
	var buf bytes.Buffer
	if err := binary.Write(&buf, binary.BigEndian, uint32(0)); err != nil {
		t.Fatal(err)
	}
	if _, err := ReadEncryptedFrame(&buf, key); err == nil {
		t.Fatal("expected error for zero-length frame")
	}
}

func TestEncryptedFrameRejectsOversizeLength(t *testing.T) {
	key := testKey(t)
	var buf bytes.Buffer
	if err := binary.Write(&buf, binary.BigEndian, uint32(MaxFrameSize+1)); err != nil {
		t.Fatal(err)
	}
	if _, err := ReadEncryptedFrame(&buf, key); err == nil {
		t.Fatal("expected error for oversize frame length")
	} else if !strings.Contains(err.Error(), "too large") {
		t.Errorf("expected too-large error, got %v", err)
	}
}

func TestEncryptedFrameWrongKey(t *testing.T) {
	key1 := DeriveKey("tcp-frame-key-1")
	key2 := DeriveKey("tcp-frame-key-2")

	var buf bytes.Buffer
	if err := WriteEncryptedFrame(&buf, []byte("secret"), key1); err != nil {
		t.Fatal(err)
	}
	if _, err := ReadEncryptedFrame(&buf, key2); err == nil {
		t.Fatal("decrypting with the wrong key should fail")
	}
}

func TestHTTPPlaintextNotVisibleOnFramedTunnel(t *testing.T) {
	key := testKey(t)
	req := []byte("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n")

	var wire bytes.Buffer
	if err := WriteEncryptedFrame(&wire, req, key); err != nil {
		t.Fatalf("WriteEncryptedFrame failed: %v", err)
	}
	sniffed := wire.Bytes()

	if bytes.Contains(sniffed, []byte("GET /")) {
		t.Fatal("sniffer saw HTTP request line on the tunnel")
	}
	if bytes.Contains(sniffed, []byte("Host:")) {
		t.Fatal("sniffer saw HTTP Host header on the tunnel")
	}
	if bytes.Contains(sniffed, []byte("example.com")) {
		t.Fatal("sniffer saw HTTP hostname on the tunnel")
	}

	// Length prefix is 4 bytes; ciphertext must be longer than plaintext.
	if len(sniffed) < 4+len(req) {
		t.Fatalf("framed ciphertext too short: %d", len(sniffed))
	}

	got, err := ReadEncryptedFrame(bytes.NewReader(sniffed), key)
	if err != nil {
		t.Fatalf("decrypt failed: %v", err)
	}
	if !bytes.Equal(got, req) {
		t.Errorf("decrypted HTTP mismatch: got %q", got)
	}
}

func TestReadEncryptedFrameEOF(t *testing.T) {
	key := testKey(t)
	if _, err := ReadEncryptedFrame(bytes.NewReader(nil), key); err != io.EOF {
		t.Errorf("expected io.EOF, got %v", err)
	}
}