package proxyproto

import (
	"bytes"
	"encoding/binary"
	"net"
	"net/netip"
	"testing"
	"time"
)

// pipeWithDeadline returns a net.Conn pair where the reader side supports deadlines.
func pipeWithDeadline(t *testing.T) (reader net.Conn, writer net.Conn) {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	t.Cleanup(func() { _ = ln.Close() })

	ch := make(chan net.Conn, 1)
	go func() {
		c, err := ln.Accept()
		if err != nil {
			return
		}
		ch <- c
	}()

	w, err := net.Dial("tcp", ln.Addr().String())
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	t.Cleanup(func() { _ = w.Close() })

	r := <-ch
	t.Cleanup(func() { _ = r.Close() })
	return r, w
}

func TestParseV1TCP4(t *testing.T) {
	reader, writer := pipeWithDeadline(t)

	go func() {
		_, _ = writer.Write([]byte("PROXY TCP4 192.168.1.1 10.0.0.1 56324 443\r\n"))
	}()

	hdr, err := Parse(reader, time.Now().Add(5*time.Second))
	if err != nil {
		t.Fatalf("Parse: %v", err)
	}

	if hdr.Version != 1 {
		t.Fatalf("version = %d, want 1", hdr.Version)
	}
	if hdr.Command != CommandProxy {
		t.Fatalf("command = %d, want CommandProxy", hdr.Command)
	}
	if got := hdr.SrcAddr.String(); got != "192.168.1.1:56324" {
		t.Fatalf("src = %s, want 192.168.1.1:56324", got)
	}
	if got := hdr.DstAddr.String(); got != "10.0.0.1:443" {
		t.Fatalf("dst = %s, want 10.0.0.1:443", got)
	}
}

func TestParseV1TCP6(t *testing.T) {
	reader, writer := pipeWithDeadline(t)

	go func() {
		_, _ = writer.Write([]byte("PROXY TCP6 2001:db8::1 2001:db8::2 56324 8443\r\n"))
	}()

	hdr, err := Parse(reader, time.Now().Add(5*time.Second))
	if err != nil {
		t.Fatalf("Parse: %v", err)
	}

	if hdr.Version != 1 {
		t.Fatalf("version = %d, want 1", hdr.Version)
	}
	wantSrc := "[2001:db8::1]:56324"
	if got := hdr.SrcAddr.String(); got != wantSrc {
		t.Fatalf("src = %s, want %s", got, wantSrc)
	}
}

func TestParseV2TCP4(t *testing.T) {
	reader, writer := pipeWithDeadline(t)

	go func() {
		var buf []byte
		buf = append(buf, v2Signature[:]...)
		buf = append(buf, 0x21) // version 2, PROXY command
		buf = append(buf, 0x11) // AF_INET, STREAM
		buf = binary.BigEndian.AppendUint16(buf, 12)
		buf = append(buf, 192, 168, 1, 100) // src IP
		buf = append(buf, 10, 0, 0, 1)      // dst IP
		buf = binary.BigEndian.AppendUint16(buf, 12345)
		buf = binary.BigEndian.AppendUint16(buf, 443)
		_, _ = writer.Write(buf)
	}()

	hdr, err := Parse(reader, time.Now().Add(5*time.Second))
	if err != nil {
		t.Fatalf("Parse: %v", err)
	}

	if hdr.Version != 2 {
		t.Fatalf("version = %d, want 2", hdr.Version)
	}
	if hdr.Command != CommandProxy {
		t.Fatalf("command = %d, want CommandProxy", hdr.Command)
	}
	if got := hdr.SrcAddr.String(); got != "192.168.1.100:12345" {
		t.Fatalf("src = %s, want 192.168.1.100:12345", got)
	}
	if got := hdr.DstAddr.String(); got != "10.0.0.1:443" {
		t.Fatalf("dst = %s, want 10.0.0.1:443", got)
	}
}

func TestParseV2TCP6(t *testing.T) {
	reader, writer := pipeWithDeadline(t)

	go func() {
		var buf []byte
		buf = append(buf, v2Signature[:]...)
		buf = append(buf, 0x21) // version 2, PROXY command
		buf = append(buf, 0x21) // AF_INET6, STREAM
		buf = binary.BigEndian.AppendUint16(buf, 36)
		// src: 2001:db8::1
		src := netip.MustParseAddr("2001:db8::1").As16()
		buf = append(buf, src[:]...)
		// dst: 2001:db8::2
		dst := netip.MustParseAddr("2001:db8::2").As16()
		buf = append(buf, dst[:]...)
		buf = binary.BigEndian.AppendUint16(buf, 56324)
		buf = binary.BigEndian.AppendUint16(buf, 8443)
		_, _ = writer.Write(buf)
	}()

	hdr, err := Parse(reader, time.Now().Add(5*time.Second))
	if err != nil {
		t.Fatalf("Parse: %v", err)
	}

	if hdr.Version != 2 {
		t.Fatalf("version = %d, want 2", hdr.Version)
	}
	wantSrc := "[2001:db8::1]:56324"
	if got := hdr.SrcAddr.String(); got != wantSrc {
		t.Fatalf("src = %s, want %s", got, wantSrc)
	}
}

func TestParseV2Local(t *testing.T) {
	reader, writer := pipeWithDeadline(t)

	go func() {
		var buf []byte
		buf = append(buf, v2Signature[:]...)
		buf = append(buf, 0x20) // version 2, LOCAL command
		buf = append(buf, 0x00) // unspec family, unspec protocol
		buf = binary.BigEndian.AppendUint16(buf, 0)
		_, _ = writer.Write(buf)
	}()

	hdr, err := Parse(reader, time.Now().Add(5*time.Second))
	if err != nil {
		t.Fatalf("Parse: %v", err)
	}

	if hdr.Version != 2 {
		t.Fatalf("version = %d, want 2", hdr.Version)
	}
	if hdr.Command != CommandLocal {
		t.Fatalf("command = %d, want CommandLocal", hdr.Command)
	}
}

func TestParseV1Malformed(t *testing.T) {
	tests := []struct {
		name string
		data string
	}{
		{"wrong prefix", "XROXY TCP4 1.2.3.4 5.6.7.8 1 2\r\n"},
		{"missing fields", "PROXY TCP4 1.2.3.4\r\n"},
		{"bad protocol", "PROXY UDP4 1.2.3.4 5.6.7.8 1 2\r\n"},
		{"bad src IP", "PROXY TCP4 not-an-ip 5.6.7.8 1 2\r\n"},
		{"bad src port", "PROXY TCP4 1.2.3.4 5.6.7.8 notport 2\r\n"},
		{"missing CRLF", "PROXY TCP4 1.2.3.4 5.6.7.8 1 2\n"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			reader, writer := pipeWithDeadline(t)
			go func() {
				_, _ = writer.Write([]byte(tt.data))
			}()
			_, err := Parse(reader, time.Now().Add(2*time.Second))
			if err == nil {
				t.Fatal("expected error")
			}
		})
	}
}

func TestParseV2Malformed(t *testing.T) {
	t.Run("bad signature", func(t *testing.T) {
		reader, writer := pipeWithDeadline(t)
		go func() {
			bad := make([]byte, v2HeaderLen)
			bad[0] = v2Signature[0] // first byte matches but rest doesn't
			_, _ = writer.Write(bad)
		}()
		_, err := Parse(reader, time.Now().Add(2*time.Second))
		if err == nil {
			t.Fatal("expected error for bad signature")
		}
	})

	t.Run("bad version", func(t *testing.T) {
		reader, writer := pipeWithDeadline(t)
		go func() {
			var buf []byte
			buf = append(buf, v2Signature[:]...)
			buf = append(buf, 0x31) // version 3, PROXY command
			buf = append(buf, 0x11)
			buf = binary.BigEndian.AppendUint16(buf, 0)
			_, _ = writer.Write(buf)
		}()
		_, err := Parse(reader, time.Now().Add(2*time.Second))
		if err == nil {
			t.Fatal("expected error for bad version")
		}
	})

	t.Run("truncated address", func(t *testing.T) {
		reader, writer := pipeWithDeadline(t)
		go func() {
			var buf []byte
			buf = append(buf, v2Signature[:]...)
			buf = append(buf, 0x21) // version 2, PROXY
			buf = append(buf, 0x11) // AF_INET, STREAM
			buf = binary.BigEndian.AppendUint16(buf, 12)
			buf = append(buf, 1, 2, 3) // only 3 bytes, need 12
			_, _ = writer.Write(buf)
			_ = writer.Close()
		}()
		_, err := Parse(reader, time.Now().Add(2*time.Second))
		if err == nil {
			t.Fatal("expected error for truncated address")
		}
	})

	t.Run("unsupported family", func(t *testing.T) {
		reader, writer := pipeWithDeadline(t)
		go func() {
			var buf []byte
			buf = append(buf, v2Signature[:]...)
			buf = append(buf, 0x21) // version 2, PROXY
			buf = append(buf, 0x31) // AF_UNIX (3), STREAM
			buf = binary.BigEndian.AppendUint16(buf, 0)
			_, _ = writer.Write(buf)
		}()
		_, err := Parse(reader, time.Now().Add(2*time.Second))
		if err == nil {
			t.Fatal("expected error for unsupported family")
		}
	})
}

func TestParseTimeout(t *testing.T) {
	reader, _ := pipeWithDeadline(t)

	// No data sent — should hit deadline.
	_, err := Parse(reader, time.Now().Add(100*time.Millisecond))
	if err == nil {
		t.Fatal("expected timeout error")
	}
}

func TestWriteV2IPv4(t *testing.T) {
	src := netip.MustParseAddrPort("192.168.1.1:56324")
	dst := netip.MustParseAddrPort("10.0.0.1:443")

	var buf bytes.Buffer
	if err := WriteV2(&buf, src, dst); err != nil {
		t.Fatalf("WriteV2: %v", err)
	}

	b := buf.Bytes()
	// Signature (12) + ver/cmd (1) + fam (1) + len (2) + 4+4+2+2 = 28
	if len(b) != 28 {
		t.Fatalf("wrote %d bytes, want 28", len(b))
	}
	// Verify signature.
	if [12]byte(b[:12]) != v2Signature {
		t.Fatal("bad signature")
	}
	if b[12] != 0x21 {
		t.Fatalf("ver/cmd = 0x%02x, want 0x21", b[12])
	}
	if b[13] != 0x11 {
		t.Fatalf("fam/proto = 0x%02x, want 0x11", b[13])
	}
}

func TestWriteV2IPv6(t *testing.T) {
	src := netip.MustParseAddrPort("[2001:db8::1]:56324")
	dst := netip.MustParseAddrPort("[2001:db8::2]:443")

	var buf bytes.Buffer
	if err := WriteV2(&buf, src, dst); err != nil {
		t.Fatalf("WriteV2: %v", err)
	}

	b := buf.Bytes()
	// Signature (12) + ver/cmd (1) + fam (1) + len (2) + 16+16+2+2 = 52
	if len(b) != 52 {
		t.Fatalf("wrote %d bytes, want 52", len(b))
	}
	if b[13] != 0x21 {
		t.Fatalf("fam/proto = 0x%02x, want 0x21", b[13])
	}
}

func TestRoundTripV2IPv4(t *testing.T) {
	src := netip.MustParseAddrPort("203.0.113.50:12345")
	dst := netip.MustParseAddrPort("198.51.100.1:443")

	reader, writer := pipeWithDeadline(t)

	go func() {
		_ = WriteV2(writer, src, dst)
	}()

	hdr, err := Parse(reader, time.Now().Add(5*time.Second))
	if err != nil {
		t.Fatalf("Parse: %v", err)
	}

	if hdr.Version != 2 {
		t.Fatalf("version = %d, want 2", hdr.Version)
	}
	if hdr.Command != CommandProxy {
		t.Fatalf("command = %d, want CommandProxy", hdr.Command)
	}
	if hdr.SrcAddr != src {
		t.Fatalf("src = %s, want %s", hdr.SrcAddr, src)
	}
	if hdr.DstAddr != dst {
		t.Fatalf("dst = %s, want %s", hdr.DstAddr, dst)
	}
}

func TestRoundTripV2IPv6(t *testing.T) {
	src := netip.MustParseAddrPort("[2001:db8:cafe::1]:40000")
	dst := netip.MustParseAddrPort("[2001:db8:beef::2]:8443")

	reader, writer := pipeWithDeadline(t)

	go func() {
		_ = WriteV2(writer, src, dst)
	}()

	hdr, err := Parse(reader, time.Now().Add(5*time.Second))
	if err != nil {
		t.Fatalf("Parse: %v", err)
	}

	if hdr.SrcAddr != src {
		t.Fatalf("src = %s, want %s", hdr.SrcAddr, src)
	}
	if hdr.DstAddr != dst {
		t.Fatalf("dst = %s, want %s", hdr.DstAddr, dst)
	}
}

func TestParseGarbageFirstByte(t *testing.T) {
	reader, writer := pipeWithDeadline(t)
	go func() {
		_, _ = writer.Write([]byte{0xFF, 0x00, 0x01})
	}()
	_, err := Parse(reader, time.Now().Add(2*time.Second))
	if err == nil {
		t.Fatal("expected error for garbage first byte")
	}
}