package proxy

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"testing"
	"time"

	xproxy "golang.org/x/net/proxy"
)

// startStubConnectProxy runs a minimal HTTP CONNECT proxy for the
// duration of the test: it accepts one CONNECT request, replies with the
// given status, and if status is 200 splices the tunnel through to a
// real dial of the requested host. Returns the proxy's address.
func startStubConnectProxy(t *testing.T, status int) string {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	t.Cleanup(func() { ln.Close() })

	go func() {
		c, err := ln.Accept()
		if err != nil {
			return
		}
		defer c.Close()
		br := bufio.NewReader(c)
		req, err := http.ReadRequest(br)
		if err != nil {
			return
		}
		if req.Method != http.MethodConnect {
			c.Write([]byte("HTTP/1.1 405 Method Not Allowed\r\n\r\n"))
			return
		}
		if status != http.StatusOK {
			c.Write([]byte("HTTP/1.1 403 Forbidden\r\n\r\n"))
			return
		}
		target, err := net.Dial("tcp", req.Host)
		if err != nil {
			c.Write([]byte("HTTP/1.1 502 Bad Gateway\r\n\r\n"))
			return
		}
		defer target.Close()
		c.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n"))
		done := make(chan struct{}, 2)
		go func() { io.Copy(target, br); done <- struct{}{} }()
		go func() { io.Copy(c, target); done <- struct{}{} }()
		<-done
	}()
	return ln.Addr().String()
}

// startStubSOCKS5Proxy runs a minimal RFC 1928 SOCKS5 server for the
// duration of the test: negotiates no-auth or username/password (per
// requireAuth), accepts one CONNECT request, and splices the tunnel
// through to a real dial of the requested address. Deliberately minimal
// (IPv4/domain address types only, one connection) - enough to exercise
// dialSOCKS5's actual wire behavior against a real server, not a mock of
// golang.org/x/net/proxy's own client logic.
func startStubSOCKS5Proxy(t *testing.T, requireAuth bool, user, pass string) string {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	t.Cleanup(func() { ln.Close() })

	go func() {
		c, err := ln.Accept()
		if err != nil {
			return
		}
		defer c.Close()
		br := bufio.NewReader(c)

		// Greeting: VER NMETHODS METHODS...
		hdr := make([]byte, 2)
		if _, err := io.ReadFull(br, hdr); err != nil {
			return
		}
		methods := make([]byte, hdr[1])
		if _, err := io.ReadFull(br, methods); err != nil {
			return
		}
		want := byte(0x00) // no auth
		if requireAuth {
			want = 0x02 // username/password
		}
		found := false
		for _, m := range methods {
			if m == want {
				found = true
			}
		}
		if !found {
			c.Write([]byte{0x05, 0xff})
			return
		}
		c.Write([]byte{0x05, want})

		if requireAuth {
			// VER ULEN UNAME PLEN PASSWD
			authHdr := make([]byte, 2)
			if _, err := io.ReadFull(br, authHdr); err != nil {
				return
			}
			uname := make([]byte, authHdr[1])
			if _, err := io.ReadFull(br, uname); err != nil {
				return
			}
			plenBuf := make([]byte, 1)
			if _, err := io.ReadFull(br, plenBuf); err != nil {
				return
			}
			passwd := make([]byte, plenBuf[0])
			if _, err := io.ReadFull(br, passwd); err != nil {
				return
			}
			if string(uname) != user || string(passwd) != pass {
				c.Write([]byte{0x01, 0x01}) // auth failure
				return
			}
			c.Write([]byte{0x01, 0x00}) // auth success
		}

		// Request: VER CMD RSV ATYP DST.ADDR DST.PORT
		req := make([]byte, 4)
		if _, err := io.ReadFull(br, req); err != nil {
			return
		}
		if req[1] != 0x01 { // CONNECT only
			c.Write([]byte{0x05, 0x07, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
			return
		}
		var targetHost string
		switch req[3] {
		case 0x01: // IPv4
			ip := make([]byte, 4)
			if _, err := io.ReadFull(br, ip); err != nil {
				return
			}
			targetHost = net.IP(ip).String()
		case 0x03: // domain
			l := make([]byte, 1)
			if _, err := io.ReadFull(br, l); err != nil {
				return
			}
			d := make([]byte, l[0])
			if _, err := io.ReadFull(br, d); err != nil {
				return
			}
			targetHost = string(d)
		default:
			c.Write([]byte{0x05, 0x08, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
			return
		}
		portBuf := make([]byte, 2)
		if _, err := io.ReadFull(br, portBuf); err != nil {
			return
		}
		targetPort := int(portBuf[0])<<8 | int(portBuf[1])
		targetAddr := net.JoinHostPort(targetHost, fmt.Sprintf("%d", targetPort))

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

		done := make(chan struct{}, 2)
		go func() { io.Copy(target, br); done <- struct{}{} }()
		go func() { io.Copy(c, target); done <- struct{}{} }()
		<-done
	}()
	return ln.Addr().String()
}

// startEchoServer runs a TCP server that echoes back whatever it reads,
// standing in for "the origin" on the far side of a CONNECT tunnel.
func startEchoServer(t *testing.T) string {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", 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()
}

func TestDialViaProxyDirect(t *testing.T) {
	echoAddr := startEchoServer(t)
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	conn, err := dialViaProxy(ctx, echoAddr, "")
	if err != nil {
		t.Fatalf("dialViaProxy direct: %v", err)
	}
	defer conn.Close()

	if _, err := conn.Write([]byte("hello")); err != nil {
		t.Fatalf("write: %v", err)
	}
	buf := make([]byte, 5)
	if _, err := io.ReadFull(conn, buf); err != nil {
		t.Fatalf("read: %v", err)
	}
	if string(buf) != "hello" {
		t.Errorf("got %q, want %q", buf, "hello")
	}
}

func TestDialViaProxyTunneled(t *testing.T) {
	echoAddr := startEchoServer(t)
	proxyAddr := startStubConnectProxy(t, http.StatusOK)
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	conn, err := dialViaProxy(ctx, echoAddr, proxyAddr)
	if err != nil {
		t.Fatalf("dialViaProxy via proxy: %v", err)
	}
	defer conn.Close()

	if _, err := conn.Write([]byte("world")); err != nil {
		t.Fatalf("write: %v", err)
	}
	buf := make([]byte, 5)
	if _, err := io.ReadFull(conn, buf); err != nil {
		t.Fatalf("read: %v", err)
	}
	if string(buf) != "world" {
		t.Errorf("got %q, want %q", buf, "world")
	}
}

func TestDialViaProxyRejected(t *testing.T) {
	proxyAddr := startStubConnectProxy(t, http.StatusForbidden)
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	_, err := dialViaProxy(ctx, "example.invalid:443", proxyAddr)
	if err == nil {
		t.Fatal("expected an error when the upstream proxy refuses CONNECT, got nil")
	}
}

func TestParseSOCKS5NotSOCKS5(t *testing.T) {
	for _, in := range []string{"", "127.0.0.1:8080", "http://127.0.0.1:8080"} {
		addr, auth, err := parseSOCKS5(in)
		if err != nil {
			t.Errorf("parseSOCKS5(%q): unexpected error: %v", in, err)
		}
		if addr != "" || auth != nil {
			t.Errorf("parseSOCKS5(%q) = %q, %+v, want empty/nil (not a socks5 upstream)", in, addr, auth)
		}
	}
}

func TestParseSOCKS5NoAuth(t *testing.T) {
	addr, auth, err := parseSOCKS5("socks5://127.0.0.1:1080")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if addr != "127.0.0.1:1080" {
		t.Errorf("addr = %q, want %q", addr, "127.0.0.1:1080")
	}
	if auth != nil {
		t.Errorf("auth = %+v, want nil (no credentials in the URL)", auth)
	}
}

func TestParseSOCKS5WithAuth(t *testing.T) {
	addr, auth, err := parseSOCKS5("socks5://alice:s3cret@127.0.0.1:1080")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if addr != "127.0.0.1:1080" {
		t.Errorf("addr = %q, want %q", addr, "127.0.0.1:1080")
	}
	if auth == nil || auth.User != "alice" || auth.Password != "s3cret" {
		t.Errorf("auth = %+v, want User=alice Password=s3cret", auth)
	}
}

func TestParseSOCKS5InvalidURL(t *testing.T) {
	// A raw control character is enough to make url.Parse fail.
	if _, _, err := parseSOCKS5("socks5://\x7f"); err == nil {
		t.Error("expected an error for a malformed socks5 URL")
	}
}

func TestDialSOCKS5NoAuth(t *testing.T) {
	echoAddr := startEchoServer(t)
	proxyAddr := startStubSOCKS5Proxy(t, false, "", "")
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	conn, err := dialSOCKS5(ctx, proxyAddr, nil, echoAddr)
	if err != nil {
		t.Fatalf("dialSOCKS5: %v", err)
	}
	defer conn.Close()

	if _, err := conn.Write([]byte("via s5")); err != nil {
		t.Fatalf("write: %v", err)
	}
	buf := make([]byte, 6)
	if _, err := io.ReadFull(conn, buf); err != nil {
		t.Fatalf("read: %v", err)
	}
	if string(buf) != "via s5" {
		t.Errorf("got %q, want %q", buf, "via s5")
	}
}

func TestDialSOCKS5WithAuth(t *testing.T) {
	echoAddr := startEchoServer(t)
	proxyAddr := startStubSOCKS5Proxy(t, true, "alice", "s3cret")
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	auth := &xproxy.Auth{User: "alice", Password: "s3cret"}
	conn, err := dialSOCKS5(ctx, proxyAddr, auth, echoAddr)
	if err != nil {
		t.Fatalf("dialSOCKS5 with correct credentials: %v", err)
	}
	conn.Close()
}

func TestDialSOCKS5WrongAuthRejected(t *testing.T) {
	proxyAddr := startStubSOCKS5Proxy(t, true, "alice", "s3cret")
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	auth := &xproxy.Auth{User: "alice", Password: "wrong"}
	if _, err := dialSOCKS5(ctx, proxyAddr, auth, "example.invalid:443"); err == nil {
		t.Fatal("expected an error for wrong SOCKS5 credentials, got nil")
	}
}

func TestDialViaProxySOCKS5(t *testing.T) {
	echoAddr := startEchoServer(t)
	proxyAddr := startStubSOCKS5Proxy(t, false, "", "")
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	conn, err := dialViaProxy(ctx, echoAddr, "socks5://"+proxyAddr)
	if err != nil {
		t.Fatalf("dialViaProxy via socks5: %v", err)
	}
	defer conn.Close()

	if _, err := conn.Write([]byte("hi")); err != nil {
		t.Fatalf("write: %v", err)
	}
	buf := make([]byte, 2)
	if _, err := io.ReadFull(conn, buf); err != nil {
		t.Fatalf("read: %v", err)
	}
	if string(buf) != "hi" {
		t.Errorf("got %q, want %q", buf, "hi")
	}
}