// Package probe implements the Protocol Probe: a user-initiated handshake check
// against a Proxy Profile's endpoint that reports which proxy protocols (HTTP,
// SOCKS5) the endpoint speaks. It never forwards traffic through the proxy and
// never contacts anything beyond the endpoint itself.
package probe

import (
	"context"
	"errors"
	"io"
	"net"
	"strings"
	"time"

	"proxyctl/internal/profile"
)

// attemptTimeout bounds each handshake attempt; the two checks together keep
// one Protocol Probe under roughly three seconds.
const attemptTimeout = 1500 * time.Millisecond

// Capabilities reports which proxy protocols an endpoint speaks.
type Capabilities struct {
	HTTP   bool
	SOCKS5 bool
}

// URL maps a typed endpoint and the detected Capabilities to a normalized
// proxy URL. HTTP is preferred; a SOCKS-only endpoint maps to socks5h
// (remote-DNS semantics). Any typed scheme is replaced and any user:pass@
// userinfo is preserved exactly as typed (special characters are not
// percent-encoded).
func URL(raw string, caps Capabilities) (string, error) {
	scheme := ""
	switch {
	case caps.HTTP:
		scheme = "http"
	case caps.SOCKS5:
		scheme = "socks5h"
	default:
		return "", errors.New("no proxy protocol detected")
	}
	value := strings.TrimSpace(raw)

	// Extract the raw userinfo before parsing: passwords may contain
	// characters such as '@' or spaces that net/url rejects, and we want to
	// preserve the exact bytes the user typed rather than percent-encode them.
	normalized := value
	if !profile.SchemePrefixPattern.MatchString(normalized) {
		normalized = "http://" + normalized
	}
	schemeEnd := strings.Index(normalized, "://")
	rest := normalized[schemeEnd+len("://"):]
	atIdx := strings.LastIndex(rest, "@")
	var rawUser string
	hostOnly := value
	if atIdx != -1 {
		rawUser = rest[:atIdx]
		hostOnly = normalized[:schemeEnd+len("://")] + rest[atIdx+1:]
	}

	u, err := profile.ParseProxyURL(hostOnly)
	if err != nil {
		return "", errors.New("endpoint must include host:port")
	}
	result := scheme + "://"
	if rawUser != "" {
		result += rawUser + "@"
	}
	return result + u.Host, nil
}

// Detect handshakes with the endpoint at address (host:port) and reports the
// proxy protocols it speaks. Both checks always run so the result is truthful
// for mixed ports. A dial failure (refused, timeout) is returned as an error;
// a reachable endpoint that speaks neither protocol yields zero Capabilities
// and a nil error. Credentials play no part: an auth rejection already
// identifies the protocol.
func Detect(ctx context.Context, address string) (Capabilities, error) {
	var caps Capabilities
	socks5, socks5Err := checkSOCKS5(ctx, address)
	caps.SOCKS5 = socks5
	http, httpErr := checkHTTP(ctx, address)
	caps.HTTP = http
	if caps.HTTP || caps.SOCKS5 {
		return caps, nil
	}
	if err := ctx.Err(); err != nil {
		return Capabilities{}, err
	}
	if socks5Err != nil && httpErr != nil {
		return Capabilities{}, socks5Err
	}
	return Capabilities{}, nil
}

// checkHTTP asks the proxy to CONNECT its own loopback discard port. The
// proxy gets refused locally and answers with an HTTP error status line, so
// the "HTTP/" reply prefix identifies the protocol without contacting
// anything beyond the endpoint or forwarding traffic through it. A plain
// HTTP server answers the same way; that false positive matches the default
// assumption the form already makes.
func checkHTTP(ctx context.Context, address string) (bool, error) {
	conn, err := dial(ctx, address)
	if err != nil {
		return false, err
	}
	defer conn.Close()
	stop := context.AfterFunc(ctx, func() { conn.Close() })
	defer stop()
	if _, err := io.WriteString(conn, "CONNECT 127.0.0.1:9 HTTP/1.1\r\nHost: 127.0.0.1:9\r\n\r\n"); err != nil {
		return false, nil
	}
	// Only the status line prefix matters; reading exactly len("HTTP/")
	// bytes keeps the check bounded no matter what the endpoint streams.
	prefix := make([]byte, len("HTTP/"))
	if _, err := io.ReadFull(conn, prefix); err != nil {
		return false, nil
	}
	return string(prefix) == "HTTP/", nil
}

// checkSOCKS5 sends the version-5 greeting offering only the no-auth method.
// The version byte of the reply identifies SOCKS5 even when the server
// rejects every offered method (0xFF). Errors after a successful dial mean
// the endpoint spoke something else and are not reported.
func checkSOCKS5(ctx context.Context, address string) (bool, error) {
	conn, err := dial(ctx, address)
	if err != nil {
		return false, err
	}
	defer conn.Close()
	stop := context.AfterFunc(ctx, func() { conn.Close() })
	defer stop()
	if _, err := conn.Write([]byte{0x05, 0x01, 0x00}); err != nil {
		return false, nil
	}
	reply := make([]byte, 2)
	if _, err := io.ReadFull(conn, reply); err != nil {
		return false, nil
	}
	return reply[0] == 0x05, nil
}

func dial(ctx context.Context, address string) (net.Conn, error) {
	dialer := net.Dialer{Timeout: attemptTimeout}
	conn, err := dialer.DialContext(ctx, "tcp", address)
	if err != nil {
		return nil, err
	}
	conn.SetDeadline(time.Now().Add(attemptTimeout))
	return conn, nil
}