package socks5

import (
	"context"
	"encoding/binary"
	"io"
	"net"
	"testing"
	"time"

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

// fakeUDPProxy is a minimal SOCKS5 server that handles greet+(optional auth)
// then UDP ASSOCIATE — replying with a relay endpoint we control.
type fakeUDPProxy struct {
	tcpAddr  string
	relay    *net.UDPConn // bound on 127.0.0.1, ephemeral port
	useAuth  bool
	login    string
	password string
}

func startFakeUDPProxy(t *testing.T, useAuth bool, login, password string) *fakeUDPProxy {
	tcpLn, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	t.Cleanup(func() { tcpLn.Close() })

	relay, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0})
	require.NoError(t, err)
	t.Cleanup(func() { relay.Close() })

	p := &fakeUDPProxy{
		tcpAddr: tcpLn.Addr().String(),
		relay:   relay,
		useAuth: useAuth, login: login, password: password,
	}

	go func() {
		for {
			c, err := tcpLn.Accept()
			if err != nil {
				return
			}
			go p.handle(c)
		}
	}()
	return p
}

func (p *fakeUDPProxy) handle(c net.Conn) {
	defer c.Close()
	_ = c.SetReadDeadline(time.Now().Add(5 * time.Second))
	buf := make([]byte, 256)

	// Greet
	io.ReadFull(c, buf[:2])
	nm := int(buf[1])
	io.ReadFull(c, buf[:nm])
	if p.useAuth {
		c.Write([]byte{0x05, 0x02})
		// Auth subneg: 01 ULEN UNAME PLEN PASS
		io.ReadFull(c, buf[:2])
		ulen := int(buf[1])
		io.ReadFull(c, buf[:ulen])
		login := string(buf[:ulen])
		io.ReadFull(c, buf[:1])
		plen := int(buf[0])
		io.ReadFull(c, buf[:plen])
		pwd := string(buf[:plen])
		if login != p.login || pwd != p.password {
			c.Write([]byte{0x01, 0x01})
			return
		}
		c.Write([]byte{0x01, 0x00})
	} else {
		c.Write([]byte{0x05, 0x00})
	}

	// UDP ASSOCIATE: 05 03 00 ATYP ...
	io.ReadFull(c, buf[:4])
	if buf[1] != 0x03 {
		// Not UDP ASSOCIATE; reject.
		c.Write([]byte{0x05, 0x07, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
		return
	}
	atyp := buf[3]
	switch atyp {
	case 1:
		io.ReadFull(c, buf[:4])
	case 3:
		io.ReadFull(c, buf[:1])
		io.ReadFull(c, buf[:int(buf[0])])
	}
	io.ReadFull(c, buf[:2]) // port

	// Reply with relay's local addr
	relayAddr := p.relay.LocalAddr().(*net.UDPAddr)
	rep := []byte{0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0}
	v4 := relayAddr.IP.To4()
	copy(rep[4:8], v4)
	binary.BigEndian.PutUint16(rep[8:10], uint16(relayAddr.Port))
	c.Write(rep)
	_ = c.SetReadDeadline(time.Time{})

	// Hold the conn open until peer closes (RFC 1928 §6 — control TCP
	// must remain open for the relay to stay valid).
	io.Copy(io.Discard, c)
}

func TestAssociateUDP_NoAuth(t *testing.T) {
	p := startFakeUDPProxy(t, false, "", "")
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	relay, ctrl, err := AssociateUDP(ctx, Config{ProxyAddr: p.tcpAddr})
	require.NoError(t, err)
	defer ctrl.Close()

	expected := p.relay.LocalAddr().(*net.UDPAddr)
	assert.Equal(t, expected.Port, relay.Port)
	assert.Equal(t, "127.0.0.1", relay.IP.String())
}

func TestAssociateUDP_WithAuth(t *testing.T) {
	p := startFakeUDPProxy(t, true, "user", "pass")
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	relay, ctrl, err := AssociateUDP(ctx, Config{
		ProxyAddr: p.tcpAddr,
		UseAuth:   true,
		Login:     "user",
		Password:  "pass",
	})
	require.NoError(t, err)
	defer ctrl.Close()

	require.NotNil(t, relay)
	assert.Greater(t, relay.Port, 0)
}

func TestAssociateUDP_BadAuth(t *testing.T) {
	p := startFakeUDPProxy(t, true, "user", "pass")
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	_, _, err := AssociateUDP(ctx, Config{
		ProxyAddr: p.tcpAddr,
		UseAuth:   true,
		Login:     "wrong",
		Password:  "wrong",
	})
	require.Error(t, err)
}

func TestEncapDecapUDPv4_Roundtrip(t *testing.T) {
	dstIP := net.IPv4(140, 82, 121, 4)
	payload := []byte("voice payload bytes")

	envelope, err := EncapUDPv4(dstIP, 50007, payload)
	require.NoError(t, err)

	// Verify wire layout (RFC 1928 §7)
	assert.Equal(t, byte(0x00), envelope[0], "RSV[0]")
	assert.Equal(t, byte(0x00), envelope[1], "RSV[1]")
	assert.Equal(t, byte(0x00), envelope[2], "FRAG")
	assert.Equal(t, byte(0x01), envelope[3], "ATYP=IPv4")
	assert.Equal(t, []byte{140, 82, 121, 4}, envelope[4:8])
	assert.Equal(t, uint16(50007), binary.BigEndian.Uint16(envelope[8:10]))
	assert.Equal(t, payload, envelope[10:])

	// Round-trip via DecapUDPv4
	srcIP, srcPort, gotPayload, err := DecapUDPv4(envelope)
	require.NoError(t, err)
	assert.Equal(t, "140.82.121.4", srcIP.String())
	assert.Equal(t, uint16(50007), srcPort)
	assert.Equal(t, payload, gotPayload)
}

func TestEncapUDPv4_NotIPv4(t *testing.T) {
	v6 := net.ParseIP("::1")
	_, err := EncapUDPv4(v6, 1, []byte("x"))
	assert.Error(t, err)
}

func TestDecapUDPv4_Errors(t *testing.T) {
	cases := []struct {
		name string
		buf  []byte
	}{
		{"too_short", []byte{0, 0, 0, 1, 1, 2, 3}},
		{"frag_nonzero", []byte{0, 0, 1 /* frag */, 1, 1, 2, 3, 4, 0, 80, 'x'}},
		{"atyp_not_ipv4", []byte{0, 0, 0, 4 /* IPv6 */, 1, 2, 3, 4, 0, 80, 'x'}},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			_, _, _, err := DecapUDPv4(c.buf)
			assert.Error(t, err)
		})
	}
}