package checker

import (
	"encoding/binary"
	"errors"
	"net"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// mkXorMappedV4 builds a synthetic STUN binding success response carrying an
// XOR-MAPPED-ADDRESS attribute for an IPv4 endpoint. Used by several test
// cases to keep byte-construction DRY.
func mkXorMappedV4(t *testing.T, ip net.IP, port uint16, txID [12]byte) []byte {
	t.Helper()
	ip4 := ip.To4()
	require.NotNil(t, ip4, "ip must be IPv4")

	// Attribute value: 1B reserved + 1B family + 2B xPort + 4B xAddr = 8 bytes.
	attrVal := make([]byte, 8)
	attrVal[0] = 0
	attrVal[1] = stunAddressFamilyIPv4
	binary.BigEndian.PutUint16(attrVal[2:4], port^uint16(stunMagicCookie>>16))
	xAddr := binary.BigEndian.Uint32(ip4) ^ stunMagicCookie
	binary.BigEndian.PutUint32(attrVal[4:8], xAddr)

	// Attribute header (4B) + value (8B) = 12 bytes total, no padding needed.
	attr := make([]byte, 4+len(attrVal))
	binary.BigEndian.PutUint16(attr[0:2], stunAttrXORMappedAddress)
	binary.BigEndian.PutUint16(attr[2:4], uint16(len(attrVal)))
	copy(attr[4:], attrVal)

	// 20B header + attrs.
	resp := make([]byte, 20+len(attr))
	binary.BigEndian.PutUint16(resp[0:2], stunBindingSuccessResponse)
	binary.BigEndian.PutUint16(resp[2:4], uint16(len(attr)))
	binary.BigEndian.PutUint32(resp[4:8], stunMagicCookie)
	copy(resp[8:20], txID[:])
	copy(resp[20:], attr)
	return resp
}

// mkXorMappedV6 builds a synthetic response with an IPv6 XOR-MAPPED-ADDRESS.
func mkXorMappedV6(t *testing.T, ip net.IP, port uint16, txID [12]byte) []byte {
	t.Helper()
	ip6 := ip.To16()
	require.NotNil(t, ip6, "ip must be IPv6")
	require.Equal(t, net.IPv6len, len(ip6))

	attrVal := make([]byte, 4+16)
	attrVal[0] = 0
	attrVal[1] = stunAddressFamilyIPv6
	binary.BigEndian.PutUint16(attrVal[2:4], port^uint16(stunMagicCookie>>16))

	var key [16]byte
	binary.BigEndian.PutUint32(key[0:4], stunMagicCookie)
	copy(key[4:16], txID[:])
	for i := 0; i < 16; i++ {
		attrVal[4+i] = ip6[i] ^ key[i]
	}

	attr := make([]byte, 4+len(attrVal))
	binary.BigEndian.PutUint16(attr[0:2], stunAttrXORMappedAddress)
	binary.BigEndian.PutUint16(attr[2:4], uint16(len(attrVal)))
	copy(attr[4:], attrVal)

	resp := make([]byte, 20+len(attr))
	binary.BigEndian.PutUint16(resp[0:2], stunBindingSuccessResponse)
	binary.BigEndian.PutUint16(resp[2:4], uint16(len(attr)))
	binary.BigEndian.PutUint32(resp[4:8], stunMagicCookie)
	copy(resp[8:20], txID[:])
	copy(resp[20:], attr)
	return resp
}

func TestNewTransactionID(t *testing.T) {
	a, err := NewTransactionID()
	require.NoError(t, err)
	b, err := NewTransactionID()
	require.NoError(t, err)
	assert.NotEqual(t, a, b, "two consecutive transaction IDs should differ (cryptographic randomness)")
	assert.Len(t, a[:], 12)
}

func TestEncodeBindingRequest(t *testing.T) {
	txID := [12]byte{0xDE, 0xAD, 0xBE, 0xEF, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08}
	got := EncodeBindingRequest(txID)

	require.Len(t, got, 20)
	assert.Equal(t, byte(0x00), got[0])
	assert.Equal(t, byte(0x01), got[1], "type LSB = 0x01 (binding request)")
	assert.Equal(t, byte(0x00), got[2])
	assert.Equal(t, byte(0x00), got[3], "attribute length = 0 (empty body)")
	assert.Equal(t, []byte{0x21, 0x12, 0xA4, 0x42}, got[4:8], "magic cookie")
	assert.Equal(t, txID[:], got[8:20], "transaction id")
}

func TestParseBindingResponse_HappyV4(t *testing.T) {
	txID := [12]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}
	ip := net.IPv4(198, 51, 100, 1).To4()
	const port uint16 = 42
	resp := mkXorMappedV4(t, ip, port, txID)

	// Sanity-check the bytes match the worked example in the task description.
	// xPort = 42 ^ 0x2112 = 0x2138
	assert.Equal(t, byte(0x21), resp[20+4+2])
	assert.Equal(t, byte(0x38), resp[20+4+3])
	// xIP = 0xC6336401 ^ 0x2112A442 = 0xE721C043
	assert.Equal(t, []byte{0xE7, 0x21, 0xC0, 0x43}, resp[20+4+4:20+4+8])

	gotIP, gotPort, err := ParseBindingResponse(resp, txID)
	require.NoError(t, err)
	assert.True(t, gotIP.Equal(ip), "got %s want %s", gotIP, ip)
	assert.Equal(t, port, gotPort)
	assert.Len(t, gotIP, 4, "IPv4 result should be 4-byte slice")
}

func TestParseBindingResponse_HappyV6(t *testing.T) {
	txID := [12]byte{0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xAA, 0xBB, 0xCC}
	ip := net.ParseIP("2001:db8::1")
	require.NotNil(t, ip)
	const port uint16 = 0x1234
	resp := mkXorMappedV6(t, ip, port, txID)

	gotIP, gotPort, err := ParseBindingResponse(resp, txID)
	require.NoError(t, err)
	assert.True(t, gotIP.Equal(ip), "got %s want %s", gotIP, ip)
	assert.Equal(t, port, gotPort)
	assert.Len(t, gotIP, 16, "IPv6 result should be 16-byte slice")
}

func TestParseBindingResponse_MultipleUnknownAttributesThenMapped(t *testing.T) {
	txID := [12]byte{9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9}
	ip := net.IPv4(8, 8, 8, 8).To4()
	const port uint16 = 53

	// Build header with three attributes:
	//   1. unknown type 0x8022 (SOFTWARE), value="abc" -> 3 bytes value + 1 byte pad
	//   2. unknown type 0x8023, value="hi"  -> 2 bytes + 2 bytes pad
	//   3. real XOR-MAPPED-ADDRESS
	var attrs []byte

	addAttr := func(t uint16, val []byte) {
		hdr := make([]byte, 4)
		binary.BigEndian.PutUint16(hdr[0:2], t)
		binary.BigEndian.PutUint16(hdr[2:4], uint16(len(val)))
		attrs = append(attrs, hdr...)
		attrs = append(attrs, val...)
		// pad to 4-byte boundary
		for len(attrs)%4 != 0 {
			attrs = append(attrs, 0)
		}
	}
	addAttr(0x8022, []byte("abc"))
	addAttr(0x8023, []byte("hi"))

	// XOR-MAPPED-ADDRESS attribute
	xmAttrVal := make([]byte, 8)
	xmAttrVal[1] = stunAddressFamilyIPv4
	binary.BigEndian.PutUint16(xmAttrVal[2:4], port^uint16(stunMagicCookie>>16))
	xAddr := binary.BigEndian.Uint32(ip) ^ stunMagicCookie
	binary.BigEndian.PutUint32(xmAttrVal[4:8], xAddr)
	addAttr(stunAttrXORMappedAddress, xmAttrVal)

	resp := make([]byte, 20+len(attrs))
	binary.BigEndian.PutUint16(resp[0:2], stunBindingSuccessResponse)
	binary.BigEndian.PutUint16(resp[2:4], uint16(len(attrs)))
	binary.BigEndian.PutUint32(resp[4:8], stunMagicCookie)
	copy(resp[8:20], txID[:])
	copy(resp[20:], attrs)

	gotIP, gotPort, err := ParseBindingResponse(resp, txID)
	require.NoError(t, err)
	assert.True(t, gotIP.Equal(ip))
	assert.Equal(t, port, gotPort)
}

func TestParseBindingResponse_Errors(t *testing.T) {
	txID := [12]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}
	otherTxID := [12]byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}

	t.Run("truncated", func(t *testing.T) {
		buf := make([]byte, 10)
		_, _, err := ParseBindingResponse(buf, txID)
		assert.ErrorIs(t, err, ErrSTUNTooShort)
	})

	t.Run("bad_magic_cookie", func(t *testing.T) {
		resp := mkXorMappedV4(t, net.IPv4(1, 2, 3, 4), 80, txID)
		copy(resp[4:8], []byte{0xAA, 0xBB, 0xCC, 0xDD})
		_, _, err := ParseBindingResponse(resp, txID)
		assert.ErrorIs(t, err, ErrSTUNBadMagicCookie)
	})

	t.Run("not_success_request_type", func(t *testing.T) {
		resp := mkXorMappedV4(t, net.IPv4(1, 2, 3, 4), 80, txID)
		binary.BigEndian.PutUint16(resp[0:2], stunBindingRequest) // 0x0001
		_, _, err := ParseBindingResponse(resp, txID)
		assert.ErrorIs(t, err, ErrSTUNNotSuccess)
	})

	t.Run("not_success_error_response_type", func(t *testing.T) {
		resp := mkXorMappedV4(t, net.IPv4(1, 2, 3, 4), 80, txID)
		binary.BigEndian.PutUint16(resp[0:2], 0x0111) // binding error response
		_, _, err := ParseBindingResponse(resp, txID)
		assert.ErrorIs(t, err, ErrSTUNNotSuccess)
	})

	t.Run("no_xor_mapped_address", func(t *testing.T) {
		// 20-byte header + zero attributes
		resp := make([]byte, 20)
		binary.BigEndian.PutUint16(resp[0:2], stunBindingSuccessResponse)
		binary.BigEndian.PutUint16(resp[2:4], 0)
		binary.BigEndian.PutUint32(resp[4:8], stunMagicCookie)
		copy(resp[8:20], txID[:])
		_, _, err := ParseBindingResponse(resp, txID)
		assert.ErrorIs(t, err, ErrSTUNNoMappedAddress)
	})

	t.Run("unsupported_family", func(t *testing.T) {
		resp := mkXorMappedV4(t, net.IPv4(1, 2, 3, 4), 80, txID)
		// Flip family byte (offset 20 + 4 + 1 = 25) to 0x03.
		resp[25] = 0x03
		_, _, err := ParseBindingResponse(resp, txID)
		assert.ErrorIs(t, err, ErrSTUNUnsupportedFamily)
	})

	t.Run("attribute_length_overflow", func(t *testing.T) {
		// Build a header claiming 24 bytes of attrs, but only put one bogus
		// attribute of declared length 100 inside.
		resp := make([]byte, 20+24)
		binary.BigEndian.PutUint16(resp[0:2], stunBindingSuccessResponse)
		binary.BigEndian.PutUint16(resp[2:4], 24)
		binary.BigEndian.PutUint32(resp[4:8], stunMagicCookie)
		copy(resp[8:20], txID[:])
		// attribute: type=0x0020, length=100 (lies — only 20 bytes of value follow)
		binary.BigEndian.PutUint16(resp[20:22], stunAttrXORMappedAddress)
		binary.BigEndian.PutUint16(resp[22:24], 100)
		_, _, err := ParseBindingResponse(resp, txID)
		require.Error(t, err)
		assert.Contains(t, err.Error(), "claims length 100")
	})

	t.Run("attribute_section_length_overflow", func(t *testing.T) {
		// Header says attrLen=200 but buffer only has 20 bytes after header.
		resp := make([]byte, 20+20)
		binary.BigEndian.PutUint16(resp[0:2], stunBindingSuccessResponse)
		binary.BigEndian.PutUint16(resp[2:4], 200)
		binary.BigEndian.PutUint32(resp[4:8], stunMagicCookie)
		copy(resp[8:20], txID[:])
		_, _, err := ParseBindingResponse(resp, txID)
		require.Error(t, err)
		assert.Contains(t, err.Error(), "exceeds buffer")
	})

	t.Run("tx_id_mismatch", func(t *testing.T) {
		resp := mkXorMappedV4(t, net.IPv4(1, 2, 3, 4), 80, txID)
		_, _, err := ParseBindingResponse(resp, otherTxID)
		assert.ErrorIs(t, err, ErrSTUNTxIDMismatch)
	})
}

// TestRoundTripLocalhost stands up a tiny STUN server on loopback that
// handles exactly one binding request and replies with the client's own
// address as XOR-MAPPED-ADDRESS. Verifies the encode/parse pair end-to-end
// against a real UDP socket.
func TestRoundTripLocalhost(t *testing.T) {
	server, err := net.ListenPacket("udp", "127.0.0.1:0")
	require.NoError(t, err, "net.ListenPacket must succeed for round-trip test (real-network requirement)")
	t.Cleanup(func() { _ = server.Close() })

	// Server goroutine: read one request, parse minimally, reply.
	serverDone := make(chan struct{})
	go func() {
		defer close(serverDone)
		buf := make([]byte, 1500)
		_ = server.SetReadDeadline(time.Now().Add(3 * time.Second))
		n, from, rerr := server.ReadFrom(buf)
		if rerr != nil {
			return
		}
		if n < 20 {
			return
		}
		// Verify it's a binding request with right magic.
		if binary.BigEndian.Uint16(buf[0:2]) != stunBindingRequest {
			return
		}
		if binary.BigEndian.Uint32(buf[4:8]) != stunMagicCookie {
			return
		}
		var txID [12]byte
		copy(txID[:], buf[8:20])

		udpFrom, ok := from.(*net.UDPAddr)
		if !ok {
			return
		}
		reply := mkXorMappedV4(t, udpFrom.IP.To4(), uint16(udpFrom.Port), txID)
		_, _ = server.WriteTo(reply, from)
	}()

	// Client side.
	serverAddr := server.LocalAddr().(*net.UDPAddr)
	conn, err := net.DialUDP("udp", nil, serverAddr)
	require.NoError(t, err)
	t.Cleanup(func() { _ = conn.Close() })

	txID, err := NewTransactionID()
	require.NoError(t, err)

	start := time.Now()
	_, err = conn.Write(EncodeBindingRequest(txID))
	require.NoError(t, err)

	require.NoError(t, conn.SetReadDeadline(time.Now().Add(time.Second)))
	respBuf := make([]byte, 1500)
	n, err := conn.Read(respBuf)
	require.NoError(t, err)
	rtt := time.Since(start)

	gotIP, gotPort, err := ParseBindingResponse(respBuf[:n], txID)
	require.NoError(t, err)

	clientLocal := conn.LocalAddr().(*net.UDPAddr)
	assert.True(t, gotIP.Equal(net.IPv4(127, 0, 0, 1)), "got %s want 127.0.0.1", gotIP)
	assert.Equal(t, uint16(clientLocal.Port), gotPort, "port should match client local port")
	assert.Less(t, rtt, 200*time.Millisecond, "loopback RTT should be under 200ms (got %s)", rtt)

	// Make sure the server goroutine exits cleanly.
	select {
	case <-serverDone:
	case <-time.After(2 * time.Second):
		t.Fatal("server goroutine did not exit")
	}
}

// Sanity: errors.Is chain works for wrapped sentinels.
func TestSentinelsAreUnique(t *testing.T) {
	all := []error{
		ErrSTUNTooShort,
		ErrSTUNBadMagicCookie,
		ErrSTUNNotSuccess,
		ErrSTUNTxIDMismatch,
		ErrSTUNNoMappedAddress,
		ErrSTUNUnsupportedFamily,
	}
	for i, a := range all {
		for j, b := range all {
			if i == j {
				continue
			}
			assert.False(t, errors.Is(a, b), "sentinel %d should not match sentinel %d", i, j)
		}
	}
}