//go:build linux || freebsd
// +build linux freebsd

package transport

import (
	"bufio"
	"context"
	"encoding/base64"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
)

// ProxyType defines the proxy type
type ProxyType string

const (
	ProxyHTTP   ProxyType = "http"
	ProxyHTTPS  ProxyType = "https"
	ProxySOCKS5 ProxyType = "socks5"
)

// ProxyConfig configures the proxy
type ProxyConfig struct {
	// Proxy server URL
	URL *url.URL

	// Proxy type
	Type ProxyType

	// Authentication
	Username string
	Password string

	// SOCKS5LocalResolve controls DNS resolution behavior for SOCKS5 proxies.
	// When true (socks5:// scheme), hostnames are resolved locally before being
	// sent to the proxy. When false (socks5h:// scheme), the hostname is sent
	// to the proxy for remote resolution.
	SOCKS5LocalResolve bool
}

// ParseProxyConfig parses proxy URL into configuration
func ParseProxyConfig(proxyURL string) (*ProxyConfig, error) {
	if proxyURL == "" {
		return nil, nil
	}

	u, err := url.Parse(proxyURL)
	if err != nil {
		return nil, fmt.Errorf("invalid proxy url: %w", err)
	}

	var proxyType ProxyType
	socks5LocalResolve := false
	switch u.Scheme {
	case "http", "https":
		proxyType = ProxyHTTP
	case "socks5":
		proxyType = ProxySOCKS5
		socks5LocalResolve = true
	case "socks5h":
		proxyType = ProxySOCKS5
		// socks5LocalResolve stays false
	default:
		proxyType = ProxyHTTP
	}

	username := u.User.Username()
	password, _ := u.User.Password()

	// Remove authentication from URL for further use
	u.User = nil

	return &ProxyConfig{
		URL:                u,
		Type:               proxyType,
		Username:           username,
		Password:           password,
		SOCKS5LocalResolve: socks5LocalResolve,
	}, nil
}

// NewProxyTransport creates HTTP transport with proxy
func NewProxyTransport(cfg *ProxyConfig, baseTransport *http.Transport) (*http.Transport, error) {
	if cfg == nil {
		return baseTransport, nil
	}

	// SOCKS5 must be implemented as a custom DialContext — Go's stdlib
	// http.ProxyURL only supports HTTP CONNECT proxies.
	if cfg.Type == ProxySOCKS5 {
		return newSOCKS5Transport(cfg, baseTransport)
	}

	// Create a copy of the transport
	transport := baseTransport.Clone()

	// Set up the proxy
	proxyURL := cfg.URL.String()
	if cfg.Username != "" {
		// Add authentication back for http.ProxyURL
		if cfg.Password != "" {
			proxyURL = fmt.Sprintf("%s://%s:%s@%s",
				cfg.URL.Scheme,
				cfg.Username,
				cfg.Password,
				cfg.URL.Host)
		} else {
			proxyURL = fmt.Sprintf("%s://%s@%s",
				cfg.URL.Scheme,
				cfg.Username,
				cfg.URL.Host)
		}
	}

	parsedProxy, err := url.Parse(proxyURL)
	if err != nil {
		return nil, fmt.Errorf("failed to parse proxy url: %w", err)
	}

	transport.Proxy = http.ProxyURL(parsedProxy)

	return transport, nil
}

// newSOCKS5Transport creates an http.Transport that routes connections
// through a SOCKS5 proxy by overriding DialContext. The transport's Proxy
// field is cleared so the stdlib HTTP proxy machinery does not interfere.
func newSOCKS5Transport(cfg *ProxyConfig, baseTransport *http.Transport) (*http.Transport, error) {
	transport := baseTransport.Clone()
	transport.Proxy = nil

	// Preserve the existing dialer if one was configured (e.g., bind interface).
	// SOCKS5 dialing uses net.Dialer directly because the proxy itself handles
	// connection establishment to the target.
	prevDial := transport.DialContext
	transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
		// When the caller already provided a custom dialer (e.g., bind-to-device),
		// use it to reach the SOCKS5 proxy. The proxy then negotiates the
		// connection to addr.
		if prevDial != nil {
			return prevDial(ctx, network, cfg.URL.Host)
		}
		return (&net.Dialer{}).DialContext(ctx, network, cfg.URL.Host)
	}

	// Wrap so the actual target connection happens via SOCKS5 handshake.
	transport.DialContext = socks5DialContext(transport.DialContext, cfg)
	return transport, nil
}

// socks5DialContext wraps a base dialer (which connects to the SOCKS5 proxy
// host) and returns a DialContext that performs the SOCKS5 handshake to reach
// the requested target address.
func socks5DialContext(proxyDialer DialContextFunc, cfg *ProxyConfig) DialContextFunc {
	return func(ctx context.Context, network, addr string) (net.Conn, error) {
		proxyConn, err := proxyDialer(ctx, network, cfg.URL.Host)
		if err != nil {
			return nil, fmt.Errorf("connect to socks5 proxy: %w", err)
		}
		// Perform SOCKS5 handshake on top of the proxy connection.
		return socks5Handshake(ctx, proxyConn, network, addr, cfg)
	}
}

// DialContextFunc is the shape of net.Dialer.DialContext — used internally
// for wrapping proxy dialers.
type DialContextFunc func(ctx context.Context, network, addr string) (net.Conn, error)

// DialProxy creates a connection through the proxy
func DialProxy(ctx context.Context, network, addr string, cfg *ProxyConfig) (net.Conn, error) {
	if cfg == nil {
		return (&net.Dialer{}).DialContext(ctx, network, addr)
	}

	switch cfg.Type {
	case ProxyHTTP, ProxyHTTPS:
		return dialHTTPProxy(ctx, network, addr, cfg)
	case ProxySOCKS5:
		return dialSOCKS5Proxy(ctx, network, addr, cfg)
	default:
		return nil, fmt.Errorf("unsupported proxy type: %s", cfg.Type)
	}
}

// dialHTTPProxy creates a connection via HTTP CONNECT proxy
func dialHTTPProxy(ctx context.Context, network, addr string, cfg *ProxyConfig) (net.Conn, error) {
	if network != "tcp" && network != "tcp4" && network != "tcp6" {
		return nil, fmt.Errorf("unsupported network: %s", network)
	}

	// Connect to the proxy server
	proxyAddr := cfg.URL.Host
	conn, err := (&net.Dialer{}).DialContext(ctx, "tcp", proxyAddr)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to proxy: %w", err)
	}

	// Validate addr for CRLF injection
	if strings.ContainsAny(addr, "\r\n") {
		return nil, fmt.Errorf("invalid proxy target address")
	}

	// Send HTTP CONNECT request
	connectReq := fmt.Sprintf("CONNECT %s HTTP/1.1\r\nHost: %s\r\n", addr, addr)

	if cfg.Username != "" {
		auth := base64.StdEncoding.EncodeToString([]byte(cfg.Username + ":" + cfg.Password))
		connectReq += fmt.Sprintf("Proxy-Authorization: Basic %s\r\n", auth)
	}

	connectReq += "\r\n"

	_, err = conn.Write([]byte(connectReq))
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("failed to send connect request: %w", err)
	}

	// Read the proxy response
	resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("failed to read proxy response: %w", err)
	}
	resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		conn.Close()
		return nil, fmt.Errorf("proxy connect returned %s", resp.Status)
	}

	return conn, nil
}

// dialSOCKS5Proxy creates a connection via SOCKS5 proxy
func dialSOCKS5Proxy(ctx context.Context, network, addr string, cfg *ProxyConfig) (net.Conn, error) {
	conn, err := (&net.Dialer{}).DialContext(ctx, "tcp", cfg.URL.Host)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to socks5 proxy: %w", err)
	}
	return socks5Handshake(ctx, conn, network, addr, cfg)
}

// socks5Handshake performs the SOCKS5 negotiation on top of an already-open
// connection to the proxy. It implements RFC 1928 (SOCKS5) and RFC 1929
// (username/password auth). If cfg.SOCKS5LocalResolve is true and addr contains
// a hostname, the hostname is resolved locally before being sent to the proxy.
func socks5Handshake(ctx context.Context, conn net.Conn, network, addr string, cfg *ProxyConfig) (retConn net.Conn, retErr error) {
	// Ensure the connection is closed on any error path below.
	defer func() {
		if retErr != nil {
			conn.Close()
		}
	}()

	// SOCKS5 greeting
	authMethods := []byte{0x00} // No auth
	if cfg.Username != "" {
		authMethods = []byte{0x02} // Username/password auth
	}
	greeting := append([]byte{0x05, byte(len(authMethods))}, authMethods...)

	if _, err := conn.Write(greeting); err != nil {
		return nil, fmt.Errorf("socks5 greeting failed: %w", err)
	}

	// Read greeting response
	resp := make([]byte, 2)
	if _, err := io.ReadFull(conn, resp); err != nil {
		return nil, fmt.Errorf("socks5 greeting response failed: %w", err)
	}

	if resp[0] != 0x05 {
		return nil, fmt.Errorf("socks5: invalid version: %d", resp[0])
	}

	if resp[1] == 0x02 && cfg.Username != "" {
		// Username/password auth
		authMsg := []byte{0x01, byte(len(cfg.Username))}
		authMsg = append(authMsg, []byte(cfg.Username)...)
		authMsg = append(authMsg, byte(len(cfg.Password)))
		authMsg = append(authMsg, []byte(cfg.Password)...)

		if _, err := conn.Write(authMsg); err != nil {
			return nil, fmt.Errorf("socks5 auth failed: %w", err)
		}

		authResp := make([]byte, 2)
		if _, err := io.ReadFull(conn, authResp); err != nil {
			return nil, fmt.Errorf("socks5 auth response failed: %w", err)
		}
		if authResp[1] != 0x00 {
			return nil, fmt.Errorf("socks5 authentication failed")
		}
	} else if resp[1] != 0x00 {
		return nil, fmt.Errorf("socks5: no acceptable auth method")
	}

	// SOCKS5 CONNECT request
	host, portStr, err := net.SplitHostPort(addr)
	if err != nil {
		return nil, fmt.Errorf("invalid target address: %w", err)
	}

	port, _ := strconv.Atoi(portStr)

	var connectReq []byte
	connectReq = append(connectReq, 0x05, 0x01, 0x00) // VER=5, CMD=CONNECT, RSV=0

	// Decide what to send: IPv4 literal, IPv6 literal, or domain name.
	// When SOCKS5LocalResolve is true, hostnames are resolved locally and
	// sent as IPs; otherwise the hostname is passed through to the proxy.
	target := host
	if cfg.SOCKS5LocalResolve {
		if ip := net.ParseIP(host); ip == nil {
			resolved, err := resolveSOCKS5Hostname(ctx, host)
			if err != nil {
				return nil, fmt.Errorf("socks5: resolve %s: %w", host, err)
			}
			target = resolved
		}
	}

	if ip := net.ParseIP(target); ip != nil {
		if ip4 := ip.To4(); ip4 != nil {
			connectReq = append(connectReq, 0x01) // IPv4
			connectReq = append(connectReq, ip4...)
		} else {
			connectReq = append(connectReq, 0x04) // IPv6
			connectReq = append(connectReq, ip.To16()...)
		}
	} else {
		connectReq = append(connectReq, 0x03) // Domain name
		connectReq = append(connectReq, byte(len(target)))
		connectReq = append(connectReq, []byte(target)...)
	}

	connectReq = append(connectReq, byte(port>>8), byte(port))

	if _, err := conn.Write(connectReq); err != nil {
		return nil, fmt.Errorf("socks5 connect request failed: %w", err)
	}

	// Read CONNECT response
	resp6 := make([]byte, 4)
	if _, err := io.ReadFull(conn, resp6); err != nil {
		return nil, fmt.Errorf("socks5 connect response failed: %w", err)
	}

	if resp6[0] != 0x05 || resp6[1] != 0x00 {
		return nil, fmt.Errorf("socks5 connection rejected: code=%d", resp6[1])
	}

	// Read remaining address from response (bnd.addr + bnd.port)
	addrType := resp6[3]
	switch addrType {
	case 0x01: // IPv4
		_, err = io.ReadFull(conn, make([]byte, 4+2))
	case 0x03: // Domain name
		var nameLen [1]byte
		_, err = io.ReadFull(conn, nameLen[:])
		if err == nil {
			_, err = io.ReadFull(conn, make([]byte, int(nameLen[0])+2))
		}
	case 0x04: // IPv6
		_, err = io.ReadFull(conn, make([]byte, 16+2))
	}
	if err != nil {
		return nil, fmt.Errorf("socks5: failed to read bind address: %w", err)
	}

	return conn, nil
}

// resolveSOCKS5Hostname resolves a hostname to a single IP literal. It uses
// the system resolver — the SOCKS5 proxy is the bottleneck for this path, so
// IPv6-first dialer integration is not applied here.
func resolveSOCKS5Hostname(ctx context.Context, host string) (string, error) {
	ips, err := (&net.Resolver{}).LookupIPAddr(ctx, host)
	if err != nil {
		return "", err
	}
	if len(ips) == 0 {
		return "", fmt.Errorf("no addresses for %s", host)
	}
	return ips[0].IP.String(), nil
}

// GetProxyFromEnv gets proxy from environment
func GetProxyFromEnv() string {
	// Order: HTTPS_PROXY > HTTP_PROXY > http_proxy
	for _, env := range []string{"HTTPS_PROXY", "https_proxy", "HTTP_PROXY", "http_proxy"} {
		if proxy := strings.TrimSpace(os.Getenv(env)); proxy != "" {
			return proxy
		}
	}
	return ""
}