package probe

import (
	"context"
	"encoding/binary"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"testing"
	"time"

	controlplanev1 "proxy-pool/gen/controlplane/v1"

	"google.golang.org/protobuf/types/known/durationpb"
	"google.golang.org/protobuf/types/known/timestamppb"
)

func TestExecutorBasicConfirmsProxyHandshakeEvenWhenProbeTargetFails(t *testing.T) {
	proxy := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
		if request.URL.String() != "http://example.invalid/" {
			t.Errorf("proxy request URL = %q", request.URL)
		}
		if value := request.Header.Get("Proxy-Authorization"); value != "Basic dXNlcjpzZWNyZXQ=" {
			t.Errorf("Proxy-Authorization = %q", value)
		}
		response.WriteHeader(http.StatusBadGateway)
	}))
	defer proxy.Close()
	task := validTask(t, proxy.URL, controlplanev1.CheckLevel_CHECK_LEVEL_BASIC)
	result := NewExecutor().Execute(context.Background(), task)
	if !result.Success || result.FailureClass != "" || result.Latency <= 0 {
		t.Fatalf("Execute(BASIC) = %+v", result)
	}
}

func splitProxyAddress(t *testing.T, address string) (string, uint32) {
	t.Helper()
	host, rawPort, err := net.SplitHostPort(address[len("http://"):])
	if err != nil {
		t.Fatalf("net.SplitHostPort(%q): %v", address, err)
	}
	port, err := strconv.ParseUint(rawPort, 10, 16)
	if err != nil {
		t.Fatalf("strconv.ParseUint(%q): %v", rawPort, err)
	}
	return host, uint32(port)
}

func TestExecutorTargetReportsHTTPFailureAsFact(t *testing.T) {
	proxy := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) {
		response.WriteHeader(http.StatusServiceUnavailable)
	}))
	defer proxy.Close()
	task := validTask(t, proxy.URL, controlplanev1.CheckLevel_CHECK_LEVEL_TARGET)
	task.RoutingName = "route-a"
	task.TargetUrl = "http://target.example/check"
	result := NewExecutor().Execute(context.Background(), task)
	if result.Success || result.FailureClass != FailureTargetHTTPStatus || result.Latency <= 0 {
		t.Fatalf("Execute(TARGET) = %+v", result)
	}
}

func TestExecutorRejectsEgressTaskWithRoutingProfile(t *testing.T) {
	proxy := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) {
		response.WriteHeader(http.StatusOK)
	}))
	defer proxy.Close()
	task := validTask(t, proxy.URL, controlplanev1.CheckLevel_CHECK_LEVEL_EGRESS)
	task.RoutingName = "route-a"
	task.TargetUrl = "https://egress.example/identity"
	result := NewExecutor().Execute(context.Background(), task)
	if result.Success || result.FailureClass != FailureInvalidTask {
		t.Fatalf("Execute(EGRESS with routing) = %+v", result)
	}
}

func TestExecutorReportsUnsupportedProtocolAsFact(t *testing.T) {
	now := time.Now().UTC()
	result := NewExecutor().Execute(context.Background(), &controlplanev1.CheckTask{
		TaskId: "task-a", ProxyId: "proxy-a", Protocol: controlplanev1.ProxyProtocol_PROXY_PROTOCOL_UNSPECIFIED,
		Host: "proxy.example", Port: 1080, Level: controlplanev1.CheckLevel_CHECK_LEVEL_BASIC,
		Timeout: durationpb.New(time.Second), Attempt: 1, MaxAttempts: 1, Deadline: timestamppb.New(now.Add(time.Second)),
	})
	if result.Success || result.FailureClass != FailureUnsupportedProxy {
		t.Fatalf("Execute(SOCKS5) = %+v", result)
	}
}

func TestExecutorBasicSupportsAuthenticatedSOCKS5Proxy(t *testing.T) {
	target := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
		if request.Host != "example.invalid" {
			t.Errorf("request host = %q", request.Host)
		}
		response.WriteHeader(http.StatusBadGateway)
	}))
	defer target.Close()
	proxy := newSOCKS5TestProxy(t, target.URL, "user", "secret")
	task := validSOCKS5Task(t, proxy.address, "user", "secret")

	result := NewExecutor().Execute(context.Background(), task)
	if !result.Success || result.FailureClass != "" || result.Latency <= 0 {
		t.Fatalf("Execute(SOCKS5 BASIC) = %+v", result)
	}
	select {
	case address := <-proxy.requested:
		if address != "example.invalid:80" {
			t.Fatalf("SOCKS5 requested address = %q", address)
		}
	case <-time.After(time.Second):
		t.Fatal("SOCKS5 proxy did not receive CONNECT request")
	}
}

func TestExecutorReportsSOCKS5AuthenticationFailure(t *testing.T) {
	target := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
	defer target.Close()
	proxy := newSOCKS5TestProxy(t, target.URL, "user", "secret")
	result := NewExecutor().Execute(context.Background(), validSOCKS5Task(t, proxy.address, "user", "wrong"))
	if result.Success || result.FailureClass != FailureProxyAuth || result.Latency <= 0 {
		t.Fatalf("Execute(SOCKS5 auth failure) = %+v", result)
	}
}

func validTask(t *testing.T, proxyAddress string, level controlplanev1.CheckLevel) *controlplanev1.CheckTask {
	t.Helper()
	host, port := splitProxyAddress(t, proxyAddress)
	now := time.Now().UTC()
	return &controlplanev1.CheckTask{
		TaskId: "task-a", ProxyId: "proxy-a", Protocol: controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTP,
		Host: host, Port: port, Level: level, Username: "user", Password: "secret",
		Timeout: durationpb.New(time.Second), Attempt: 1, MaxAttempts: 1,
		Deadline: timestamppb.New(now.Add(time.Second)),
	}
}

func validSOCKS5Task(t *testing.T, address, username, password string) *controlplanev1.CheckTask {
	t.Helper()
	host, rawPort, err := net.SplitHostPort(address)
	if err != nil {
		t.Fatalf("net.SplitHostPort(%q): %v", address, err)
	}
	port, err := strconv.ParseUint(rawPort, 10, 16)
	if err != nil {
		t.Fatalf("strconv.ParseUint(%q): %v", rawPort, err)
	}
	now := time.Now().UTC()
	return &controlplanev1.CheckTask{
		TaskId: "task-socks", ProxyId: "proxy-socks", Protocol: controlplanev1.ProxyProtocol_PROXY_PROTOCOL_SOCKS5,
		Host: host, Port: uint32(port), Level: controlplanev1.CheckLevel_CHECK_LEVEL_BASIC, Username: username, Password: password,
		Timeout: durationpb.New(time.Second), Attempt: 1, MaxAttempts: 1, Deadline: timestamppb.New(now.Add(time.Second)),
	}
}

type socks5TestProxy struct {
	address   string
	target    string
	username  string
	password  string
	requested chan string
	listener  net.Listener
}

func newSOCKS5TestProxy(t *testing.T, target, username, password string) *socks5TestProxy {
	t.Helper()
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("net.Listen(): %v", err)
	}
	proxy := &socks5TestProxy{
		address: listener.Addr().String(), target: target, username: username, password: password,
		requested: make(chan string, 1), listener: listener,
	}
	go func() {
		for {
			connection, err := listener.Accept()
			if err != nil {
				return
			}
			go proxy.serve(connection)
		}
	}()
	t.Cleanup(func() { _ = listener.Close() })
	return proxy
}

func (proxy *socks5TestProxy) serve(connection net.Conn) {
	defer connection.Close()
	if _, err := readSOCKS5Greeting(connection); err != nil {
		return
	}
	method := byte(0x00)
	if proxy.username != "" || proxy.password != "" {
		method = 0x02
	}
	if _, err := connection.Write([]byte{0x05, method}); err != nil {
		return
	}
	if method == 0x02 {
		username, password, err := readSOCKS5Credentials(connection)
		if err != nil {
			return
		}
		if username != proxy.username || password != proxy.password {
			_, _ = connection.Write([]byte{0x01, 0x01})
			return
		}
		if _, err := connection.Write([]byte{0x01, 0x00}); err != nil {
			return
		}
	}
	address, err := readSOCKS5ConnectRequest(connection)
	if err != nil {
		return
	}
	proxy.requested <- address
	targetAddress := strings.TrimPrefix(proxy.target, "http://")
	target, err := net.Dial("tcp", targetAddress)
	if err != nil {
		_, _ = connection.Write([]byte{0x05, 0x01, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
		return
	}
	defer target.Close()
	if _, err := connection.Write([]byte{0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0}); err != nil {
		return
	}
	done := make(chan struct{}, 2)
	go func() { _, _ = io.Copy(target, connection); done <- struct{}{} }()
	go func() { _, _ = io.Copy(connection, target); done <- struct{}{} }()
	<-done
}

func readSOCKS5Greeting(reader io.Reader) ([]byte, error) {
	header := make([]byte, 2)
	if _, err := io.ReadFull(reader, header); err != nil || header[0] != 0x05 || header[1] == 0 {
		return nil, io.ErrUnexpectedEOF
	}
	methods := make([]byte, header[1])
	_, err := io.ReadFull(reader, methods)
	return methods, err
}

func readSOCKS5Credentials(reader io.Reader) (string, string, error) {
	header := make([]byte, 2)
	if _, err := io.ReadFull(reader, header); err != nil || header[0] != 0x01 || header[1] == 0 {
		return "", "", io.ErrUnexpectedEOF
	}
	username := make([]byte, header[1])
	if _, err := io.ReadFull(reader, username); err != nil {
		return "", "", err
	}
	passwordLength := make([]byte, 1)
	if _, err := io.ReadFull(reader, passwordLength); err != nil || passwordLength[0] == 0 {
		return "", "", io.ErrUnexpectedEOF
	}
	password := make([]byte, passwordLength[0])
	if _, err := io.ReadFull(reader, password); err != nil {
		return "", "", err
	}
	return string(username), string(password), nil
}

func readSOCKS5ConnectRequest(reader io.Reader) (string, error) {
	header := make([]byte, 4)
	if _, err := io.ReadFull(reader, header); err != nil || header[0] != 0x05 || header[1] != 0x01 || header[2] != 0x00 {
		return "", io.ErrUnexpectedEOF
	}
	var host string
	switch header[3] {
	case 0x01:
		address := make([]byte, net.IPv4len)
		if _, err := io.ReadFull(reader, address); err != nil {
			return "", err
		}
		host = net.IP(address).String()
	case 0x03:
		length := make([]byte, 1)
		if _, err := io.ReadFull(reader, length); err != nil || length[0] == 0 {
			return "", io.ErrUnexpectedEOF
		}
		address := make([]byte, length[0])
		if _, err := io.ReadFull(reader, address); err != nil {
			return "", err
		}
		host = string(address)
	case 0x04:
		address := make([]byte, net.IPv6len)
		if _, err := io.ReadFull(reader, address); err != nil {
			return "", err
		}
		host = net.IP(address).String()
	default:
		return "", io.ErrUnexpectedEOF
	}
	port := make([]byte, 2)
	if _, err := io.ReadFull(reader, port); err != nil {
		return "", err
	}
	return net.JoinHostPort(host, strconv.Itoa(int(binary.BigEndian.Uint16(port)))), nil
}