package lspool

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"strings"
	"testing"

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

func TestPrepareLSProxyURLPassesThroughHTTPProxy(t *testing.T) {
	t.Cleanup(closeAllLSProxyBridgesForTest)

	got, err := prepareLSProxyURL("http://proxy.example.com:8080")
	require.NoError(t, err)
	require.Equal(t, "http://proxy.example.com:8080", got)
}

func TestPrepareLSProxyURLBridgesSOCKS5ForLS(t *testing.T) {
	t.Cleanup(closeAllLSProxyBridgesForTest)

	targetAddr, closeTarget := startBridgeEchoServer(t)
	defer closeTarget()

	socksURL, closeSOCKS := startBridgeSOCKS5Server(t)
	defer closeSOCKS()

	bridgeURL, err := prepareLSProxyURL(socksURL)
	require.NoError(t, err)
	require.True(t, strings.HasPrefix(bridgeURL, "http://127.0.0.1:"))

	// Same SOCKS upstream should reuse the same local bridge.
	reusedURL, err := prepareLSProxyURL(socksURL)
	require.NoError(t, err)
	require.Equal(t, bridgeURL, reusedURL)

	bridgeAddr := strings.TrimPrefix(bridgeURL, "http://")
	conn, err := net.Dial("tcp", bridgeAddr)
	require.NoError(t, err)
	defer conn.Close()

	_, err = fmt.Fprintf(conn, "CONNECT %s HTTP/1.1\r\nHost: %s\r\n\r\n", targetAddr, targetAddr)
	require.NoError(t, err)

	reader := bufio.NewReader(conn)
	resp, err := readConnectResponse(reader)
	require.NoError(t, err)
	require.Equal(t, 200, resp.StatusCode)

	_, err = conn.Write([]byte("ping"))
	require.NoError(t, err)

	reply := make([]byte, 4)
	_, err = io.ReadFull(reader, reply)
	require.NoError(t, err)
	require.Equal(t, "pong", string(reply))
}

func startBridgeEchoServer(t *testing.T) (string, func()) {
	t.Helper()

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)

	done := make(chan struct{})
	go func() {
		defer close(done)
		for {
			conn, err := ln.Accept()
			if err != nil {
				return
			}
			go func(c net.Conn) {
				defer c.Close()
				buf := make([]byte, 4)
				if _, err := io.ReadFull(c, buf); err != nil {
					return
				}
				if string(buf) == "ping" {
					_, _ = c.Write([]byte("pong"))
				}
			}(conn)
		}
	}()

	return ln.Addr().String(), func() {
		_ = ln.Close()
		<-done
	}
}

func startBridgeSOCKS5Server(t *testing.T) (string, func()) {
	t.Helper()

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)

	done := make(chan struct{})
	go func() {
		defer close(done)
		for {
			conn, err := ln.Accept()
			if err != nil {
				return
			}
			go handleBridgeSOCKS5Conn(conn)
		}
	}()

	return "socks5://" + ln.Addr().String(), func() {
		_ = ln.Close()
		<-done
	}
}

func handleBridgeSOCKS5Conn(conn net.Conn) {
	header := make([]byte, 2)
	if _, err := io.ReadFull(conn, header); err != nil {
		_ = conn.Close()
		return
	}
	methods := make([]byte, int(header[1]))
	if _, err := io.ReadFull(conn, methods); err != nil {
		_ = conn.Close()
		return
	}
	_, _ = conn.Write([]byte{0x05, 0x00})

	reqHeader := make([]byte, 4)
	if _, err := io.ReadFull(conn, reqHeader); err != nil {
		_ = conn.Close()
		return
	}
	if reqHeader[0] != 0x05 || reqHeader[1] != 0x01 {
		_ = conn.Close()
		return
	}

	targetHost, ok := readSOCKS5Addr(conn, reqHeader[3])
	if !ok {
		_ = conn.Close()
		return
	}
	portBuf := make([]byte, 2)
	if _, err := io.ReadFull(conn, portBuf); err != nil {
		_ = conn.Close()
		return
	}
	targetAddr := fmt.Sprintf("%s:%d", targetHost, binary.BigEndian.Uint16(portBuf))

	targetConn, err := net.Dial("tcp", targetAddr)
	if err != nil {
		_, _ = conn.Write([]byte{0x05, 0x01, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
		_ = conn.Close()
		return
	}

	_, _ = conn.Write([]byte{0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
	tunnelConns(conn, targetConn)
}

func readSOCKS5Addr(conn net.Conn, atyp byte) (string, bool) {
	switch atyp {
	case 0x01:
		buf := make([]byte, 4)
		if _, err := io.ReadFull(conn, buf); err != nil {
			return "", false
		}
		return net.IP(buf).String(), true
	case 0x03:
		lenBuf := make([]byte, 1)
		if _, err := io.ReadFull(conn, lenBuf); err != nil {
			return "", false
		}
		buf := make([]byte, int(lenBuf[0]))
		if _, err := io.ReadFull(conn, buf); err != nil {
			return "", false
		}
		return string(buf), true
	case 0x04:
		buf := make([]byte, 16)
		if _, err := io.ReadFull(conn, buf); err != nil {
			return "", false
		}
		return net.IP(buf).String(), true
	default:
		return "", false
	}
}