package redirect

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

	"git.okcu.io/root/drover-go/internal/socks5"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// fakeInjector captures injected packets for assertions.
type fakeInjector struct {
	mu      sync.Mutex
	packets [][]byte
	addrs   []UDPInjectAddr
}

func (f *fakeInjector) Send(buf []byte, addr UDPInjectAddr) (int, error) {
	f.mu.Lock()
	cp := make([]byte, len(buf))
	copy(cp, buf)
	f.packets = append(f.packets, cp)
	f.addrs = append(f.addrs, addr)
	f.mu.Unlock()
	return len(buf), nil
}

func (f *fakeInjector) packetsLen() int {
	f.mu.Lock()
	defer f.mu.Unlock()
	return len(f.packets)
}

func (f *fakeInjector) get(idx int) ([]byte, UDPInjectAddr) {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.packets[idx], f.addrs[idx]
}

// startUDPRelayProxy starts a fake SOCKS5 proxy with UDP ASSOCIATE
// support. It echoes any datagram it receives on the relay back to
// the sender, with the SOCKS5 envelope's DST.ADDR/DST.PORT preserved.
// The "echoOrigin" return-path is what the real upstream relay does:
// when an upstream UDP server responds, the proxy puts that server's
// addr in DST.ADDR/DST.PORT for the inbound envelope.
func startUDPRelayProxy(t *testing.T) (tcpAddr string, relay *net.UDPConn) {
	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() })

	go func() {
		for {
			c, err := tcpLn.Accept()
			if err != nil {
				return
			}
			go func(c net.Conn) {
				defer c.Close()
				buf := make([]byte, 256)
				// Greet
				io.ReadFull(c, buf[:2])
				nm := int(buf[1])
				io.ReadFull(c, buf[:nm])
				c.Write([]byte{0x05, 0x00})
				// UDP ASSOCIATE
				io.ReadFull(c, buf[:4])
				if buf[1] != 0x03 {
					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])
				// Reply with relay addr
				ra := relay.LocalAddr().(*net.UDPAddr)
				rep := []byte{0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0}
				copy(rep[4:8], ra.IP.To4())
				binary.BigEndian.PutUint16(rep[8:10], uint16(ra.Port))
				c.Write(rep)
				// Hold open
				io.Copy(io.Discard, c)
			}(c)
		}
	}()
	return tcpLn.Addr().String(), relay
}

func TestUDPProxy_ForwardEncapsulates(t *testing.T) {
	tcpAddr, relay := startUDPRelayProxy(t)
	inj := &fakeInjector{}

	u, err := NewUDP(UDPConfig{
		SOCKS5:   socks5.Config{ProxyAddr: tcpAddr},
		LocalIP:  net.IPv4(127, 0, 0, 1),
		Injector: inj,
	})
	require.NoError(t, err)
	t.Cleanup(func() { u.Close() })

	// Forward a packet and verify the relay receives it encapsulated.
	srcIP := net.IPv4(127, 0, 0, 1)
	dstIP := net.IPv4(140, 82, 121, 4)
	payload := []byte("hello voice")
	require.NoError(t, u.Forward(srcIP, 50100, dstIP, 50007, payload))

	// Read from the relay to verify the SOCKS5 envelope.
	buf := make([]byte, 1500)
	_ = relay.SetReadDeadline(time.Now().Add(2 * time.Second))
	n, _, err := relay.ReadFromUDP(buf)
	require.NoError(t, err)

	got := buf[:n]
	gotIP, gotPort, gotPayload, err := socks5.DecapUDPv4(got)
	require.NoError(t, err)
	assert.Equal(t, "140.82.121.4", gotIP.String())
	assert.Equal(t, uint16(50007), gotPort)
	assert.Equal(t, payload, gotPayload)

	assert.Equal(t, 1, u.FlowCount(), "should have one tracked flow")
}

func TestUDPProxy_RelayResponseInjectsBackToDiscord(t *testing.T) {
	tcpAddr, relay := startUDPRelayProxy(t)
	inj := &fakeInjector{}

	u, err := NewUDP(UDPConfig{
		SOCKS5:   socks5.Config{ProxyAddr: tcpAddr},
		LocalIP:  net.IPv4(127, 0, 0, 1),
		Injector: inj,
	})
	require.NoError(t, err)
	t.Cleanup(func() { u.Close() })

	// Establish a flow by forwarding one packet
	discordSrcIP := net.IPv4(127, 0, 0, 1)
	discordSrcPort := uint16(50100)
	realDstIP := net.IPv4(140, 82, 121, 4)
	realDstPort := uint16(50007)
	require.NoError(t, u.Forward(discordSrcIP, discordSrcPort, realDstIP, realDstPort, []byte("hi")))

	// Drain the encapsulated forward
	drainBuf := make([]byte, 1500)
	_ = relay.SetReadDeadline(time.Now().Add(2 * time.Second))
	_, clientRelayAddr, err := relay.ReadFromUDP(drainBuf)
	require.NoError(t, err)

	// Simulate upstream UDP server response: relay sends back an
	// envelope where DST.ADDR/DST.PORT = real upstream origin.
	respPayload := []byte("voice response")
	envelope, err := socks5.EncapUDPv4(realDstIP, realDstPort, respPayload)
	require.NoError(t, err)
	_, err = relay.WriteToUDP(envelope, clientRelayAddr)
	require.NoError(t, err)

	// The proxy's relayReadLoop should receive, decap, and inject.
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) && inj.packetsLen() == 0 {
		time.Sleep(20 * time.Millisecond)
	}
	require.Equal(t, 1, inj.packetsLen(), "expected one injected packet")

	pkt, addr := inj.get(0)
	assert.False(t, addr.Outbound, "injected as inbound")

	// Parse the fabricated IPv4+UDP packet
	require.GreaterOrEqual(t, len(pkt), 28)
	// Verify proto=UDP
	assert.Equal(t, byte(17), pkt[9], "IPv4 proto field")
	srcIP := net.IPv4(pkt[12], pkt[13], pkt[14], pkt[15])
	dstIP := net.IPv4(pkt[16], pkt[17], pkt[18], pkt[19])
	srcPort := binary.BigEndian.Uint16(pkt[20:22])
	dstPort := binary.BigEndian.Uint16(pkt[22:24])

	assert.Equal(t, "140.82.121.4", srcIP.String(), "fabricated src = real upstream origin")
	assert.Equal(t, "127.0.0.1", dstIP.String(), "fabricated dst = Discord-side IP")
	assert.Equal(t, realDstPort, srcPort, "fabricated src port = real upstream port")
	assert.Equal(t, discordSrcPort, dstPort, "fabricated dst port = Discord ephemeral port")

	// Payload after IPv4(20)+UDP(8) headers
	assert.Equal(t, respPayload, pkt[28:])
}

func TestUDPProxy_NoFlowDropsResponse(t *testing.T) {
	tcpAddr, relay := startUDPRelayProxy(t)
	inj := &fakeInjector{}

	u, err := NewUDP(UDPConfig{
		SOCKS5:   socks5.Config{ProxyAddr: tcpAddr},
		LocalIP:  net.IPv4(127, 0, 0, 1),
		Injector: inj,
	})
	require.NoError(t, err)
	t.Cleanup(func() { u.Close() })

	// Force association without registering any flow.
	require.NoError(t, u.ensureAssociated())

	// Read the local relay socket's port and substitute 127.0.0.1 for
	// 0.0.0.0 (kernel binds wildcard but Windows refuses to send TO
	// 0.0.0.0:N — it requires a routable destination).
	localAddr := u.relayConn.LocalAddr().(*net.UDPAddr)
	dst := &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: localAddr.Port}

	// Send a "stray" relay datagram with an origin we never registered.
	envelope, _ := socks5.EncapUDPv4(net.IPv4(8, 8, 8, 8), 53, []byte("dns"))
	_, err = relay.WriteToUDP(envelope, dst)
	require.NoError(t, err)

	// Give the reader time to process and drop.
	time.Sleep(200 * time.Millisecond)
	assert.Equal(t, 0, inj.packetsLen(), "stray response should be dropped, not injected")
}

func TestUDPProxy_RejectsIPv6(t *testing.T) {
	inj := &fakeInjector{}
	u, err := NewUDP(UDPConfig{
		SOCKS5:   socks5.Config{ProxyAddr: "127.0.0.1:0"},
		LocalIP:  net.IPv4(127, 0, 0, 1),
		Injector: inj,
	})
	require.NoError(t, err)
	t.Cleanup(func() { u.Close() })

	v6 := net.ParseIP("::1")
	err = u.Forward(net.IPv4(1, 2, 3, 4), 1000, v6, 80, []byte("x"))
	assert.Error(t, err)
}

func TestNewUDP_RejectsNilInjector(t *testing.T) {
	_, err := NewUDP(UDPConfig{
		LocalIP: net.IPv4(127, 0, 0, 1),
	})
	assert.Error(t, err)
}

func TestNewUDP_RejectsNonIPv4LocalIP(t *testing.T) {
	_, err := NewUDP(UDPConfig{
		LocalIP:  net.ParseIP("::1"),
		Injector: &fakeInjector{},
	})
	assert.Error(t, err)
}