package redirect

import (
	"context"
	"io"
	"net"
	"sync"
	"testing"
	"time"

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

// startEchoListener spins up a TCP server that echoes whatever it reads.
func startEchoListener(t *testing.T) string {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	t.Cleanup(func() { ln.Close() })
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			go func(c net.Conn) {
				defer c.Close()
				io.Copy(c, c)
			}(c)
		}
	}()
	return ln.Addr().String()
}

// startFakeSOCKS5 returns the addr of a no-auth SOCKS5 server that
// CONNECT-tunnels to the requested host:port. (Borrowed pattern from
// internal/socks5/client_test.go.)
func startFakeSOCKS5(t *testing.T) string {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	t.Cleanup(func() { ln.Close() })
	go func() {
		for {
			c, err := ln.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})
				// CONNECT
				io.ReadFull(c, buf[:4])
				atyp := buf[3]
				var host string
				switch atyp {
				case 1:
					io.ReadFull(c, buf[:4])
					host = net.IPv4(buf[0], buf[1], buf[2], buf[3]).String()
				case 3:
					io.ReadFull(c, buf[:1])
					hl := int(buf[0])
					io.ReadFull(c, buf[:hl])
					host = string(buf[:hl])
				}
				io.ReadFull(c, buf[:2])
				port := int(buf[0])<<8 | int(buf[1])
				c.Write([]byte{0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
				up, err := net.Dial("tcp", net.JoinHostPort(host, sportItoa(port)))
				if err != nil {
					return
				}
				defer up.Close()
				var wg sync.WaitGroup
				wg.Add(2)
				go func() { defer wg.Done(); io.Copy(up, c) }()
				go func() { defer wg.Done(); io.Copy(c, up) }()
				wg.Wait()
			}(c)
		}
	}()
	return ln.Addr().String()
}

func sportItoa(n int) string {
	if n == 0 {
		return "0"
	}
	out := []byte{}
	for n > 0 {
		out = append([]byte{byte('0' + n%10)}, out...)
		n /= 10
	}
	return string(out)
}

func TestRedirector_PipesEcho(t *testing.T) {
	echoAddr := startEchoListener(t)
	echoHost, echoPortStr, _ := net.SplitHostPort(echoAddr)
	echoPort := parseU16(echoPortStr)
	socksAddr := startFakeSOCKS5(t)

	r, err := New(Config{
		SOCKS5: socks5.Config{ProxyAddr: socksAddr},
		Bind:   "127.0.0.1:0",
	})
	require.NoError(t, err)
	t.Cleanup(func() { r.Close() })

	// Manually map: pretend a packet from src_port=12345 was destined to echo.
	r.SetMapping(12345, net.ParseIP(echoHost), echoPort)

	// Dial the redirector listener using src_port=12345 so it looks
	// up the mapping correctly.
	d := net.Dialer{LocalAddr: &net.TCPAddr{Port: 12345}}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	conn, err := d.DialContext(ctx, "tcp", r.LocalAddr())
	require.NoError(t, err)
	defer conn.Close()

	conn.Write([]byte("ping"))
	conn.SetReadDeadline(time.Now().Add(time.Second))
	buf := make([]byte, 4)
	io.ReadFull(conn, buf)
	require.Equal(t, "ping", string(buf))
}

func parseU16(s string) uint16 {
	var n int
	for _, c := range s {
		n = n*10 + int(c-'0')
	}
	return uint16(n)
}