// Package probe executes one bounded Checker task and returns only a health
// fact. It has no Controller, Redis, or PostgreSQL dependency.
package probe

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net"
	"net/http"
	"net/netip"
	"net/url"
	"strconv"
	"strings"
	"time"

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

const (
	FailureInvalidTask      = "INVALID_TASK"
	FailureDeadlineExceeded = "DEADLINE_EXCEEDED"
	FailureProxyRequest     = "PROXY_REQUEST"
	FailureProxyAuth        = "PROXY_AUTH"
	FailureTargetHTTPStatus = "TARGET_HTTP_STATUS"
	FailureEgressIdentity   = "EGRESS_IDENTITY"
	FailureUnsupportedProxy = "UNSUPPORTED_PROXY_PROTOCOL"
	basicHandshakeProbeURL  = "http://example.invalid/"
	maximumEgressBodyBytes  = 16 << 10
)

var errUnsupportedProtocol = errors.New("unsupported proxy protocol")

type Result struct {
	Success          bool
	FailureClass     string
	ObservedEgressIP string
	Latency          time.Duration
}

// Executor constructs a short-lived HTTP transport for each task. That keeps
// credentials task-scoped and avoids keeping stale short-TTL proxy sessions
// alive after a task completes.
type Executor struct{}

func NewExecutor() *Executor { return &Executor{} }

// Execute reports protocol/authorization reachability for BASIC and requires
// a successful target response for TARGET. Every path returns one fact so a
// Checker can report failures without treating ordinary probe failures as
// control-plane errors.
func (executor *Executor) Execute(ctx context.Context, task *controlplanev1.CheckTask) Result {
	started := time.Now()
	if ctx == nil || executor == nil {
		return Result{FailureClass: FailureInvalidTask}
	}
	deadline, target, proxyURL, err := prepare(task, started.UTC())
	if err != nil {
		if errors.Is(err, errUnsupportedProtocol) {
			return Result{FailureClass: FailureUnsupportedProxy, Latency: probeLatency(started)}
		}
		return Result{FailureClass: FailureInvalidTask, Latency: probeLatency(started)}
	}
	probeContext, cancel := context.WithDeadline(ctx, deadline)
	defer cancel()
	transport := &http.Transport{
		DialContext:           (&net.Dialer{}).DialContext,
		ForceAttemptHTTP2:     false,
		TLSHandshakeTimeout:   minDuration(time.Until(deadline), 10*time.Second),
		ResponseHeaderTimeout: minDuration(time.Until(deadline), 10*time.Second),
	}
	if proxyURL.Scheme == "socks5" {
		transport.DialContext = newSOCKS5Dialer(proxyURL).DialContext
	} else {
		transport.Proxy = http.ProxyURL(proxyURL)
	}
	defer transport.CloseIdleConnections()
	client := &http.Client{
		Transport:     transport,
		CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse },
	}
	request, err := http.NewRequestWithContext(probeContext, http.MethodGet, target, nil)
	if err != nil {
		return Result{FailureClass: FailureInvalidTask, Latency: probeLatency(started)}
	}
	response, err := client.Do(request)
	latency := probeLatency(started)
	if err != nil {
		if errors.Is(probeContext.Err(), context.DeadlineExceeded) || errors.Is(ctx.Err(), context.DeadlineExceeded) {
			return Result{FailureClass: FailureDeadlineExceeded, Latency: latency}
		}
		if errors.Is(err, errSOCKS5Authentication) {
			return Result{FailureClass: FailureProxyAuth, Latency: latency}
		}
		return Result{FailureClass: FailureProxyRequest, Latency: latency}
	}
	defer response.Body.Close()
	if response.StatusCode == http.StatusProxyAuthRequired {
		return Result{FailureClass: FailureProxyAuth, Latency: latency}
	}
	if task.GetLevel() == controlplanev1.CheckLevel_CHECK_LEVEL_BASIC {
		return Result{Success: true, Latency: latency}
	}
	if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusBadRequest {
		return Result{FailureClass: FailureTargetHTTPStatus, Latency: latency}
	}
	if task.GetLevel() == controlplanev1.CheckLevel_CHECK_LEVEL_EGRESS {
		body, err := io.ReadAll(io.LimitReader(response.Body, maximumEgressBodyBytes+1))
		if err != nil || len(body) > maximumEgressBodyBytes {
			return Result{FailureClass: FailureEgressIdentity, Latency: latency}
		}
		observedIP, ok := parseEgressIP(body)
		if !ok {
			return Result{FailureClass: FailureEgressIdentity, Latency: latency}
		}
		return Result{Success: true, ObservedEgressIP: observedIP, Latency: latency}
	}
	return Result{Success: true, Latency: latency}
}

func parseEgressIP(body []byte) (string, bool) {
	if observed, ok := normalizeEgressIP(string(body)); ok {
		return observed, true
	}
	var payload map[string]json.RawMessage
	if err := json.Unmarshal(body, &payload); err != nil {
		return "", false
	}
	for _, key := range []string{"ip", "ip_address", "address", "origin", "query"} {
		raw, exists := payload[key]
		if !exists {
			continue
		}
		var value string
		if err := json.Unmarshal(raw, &value); err != nil {
			continue
		}
		if observed, ok := normalizeEgressIP(value); ok {
			return observed, true
		}
	}
	return "", false
}

func normalizeEgressIP(value string) (string, bool) {
	for _, item := range strings.FieldsFunc(strings.TrimSpace(value), func(character rune) bool {
		return character == ',' || character == ';' || character == ' ' || character == '\t' || character == '\r' || character == '\n'
	}) {
		address, err := netip.ParseAddr(item)
		if err == nil {
			return address.Unmap().String(), true
		}
	}
	return "", false
}

func prepare(task *controlplanev1.CheckTask, now time.Time) (time.Time, string, *url.URL, error) {
	if task == nil || task.GetTaskId() == "" || task.GetProxyId() == "" || task.GetHost() == "" || task.GetPort() == 0 ||
		task.GetTimeout() == nil || task.GetTimeout().CheckValid() != nil || task.GetTimeout().AsDuration() <= 0 ||
		task.GetDeadline() == nil || task.GetDeadline().CheckValid() != nil || task.GetDeadline().AsTime().UTC().Compare(now) <= 0 ||
		task.GetAttempt() == 0 || task.GetMaxAttempts() == 0 || task.GetAttempt() > task.GetMaxAttempts() ||
		(task.GetSecretRef() == "") != (task.GetCredentialVersion() == "") {
		return time.Time{}, "", nil, errors.New("invalid checker task")
	}
	proxyScheme := ""
	switch task.GetProtocol() {
	case controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTP:
		proxyScheme = "http"
	case controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTPS:
		proxyScheme = "https"
	case controlplanev1.ProxyProtocol_PROXY_PROTOCOL_SOCKS5:
		if task.GetUsername() != "" || task.GetPassword() != "" {
			if len(task.GetUsername()) == 0 || len(task.GetUsername()) > 255 || len(task.GetPassword()) == 0 || len(task.GetPassword()) > 255 {
				return time.Time{}, "", nil, errors.New("invalid SOCKS5 credentials")
			}
		}
		proxyScheme = "socks5"
	default:
		return time.Time{}, "", nil, errUnsupportedProtocol
	}
	deadline := task.GetDeadline().AsTime().UTC()
	timeoutDeadline := now.Add(task.GetTimeout().AsDuration())
	if timeoutDeadline.Before(deadline) {
		deadline = timeoutDeadline
	}
	if !deadline.After(now) {
		return time.Time{}, "", nil, errors.New("expired checker task")
	}
	proxyURL := &url.URL{Scheme: proxyScheme, Host: net.JoinHostPort(task.GetHost(), strconv.FormatUint(uint64(task.GetPort()), 10))}
	if task.GetUsername() != "" || task.GetPassword() != "" {
		proxyURL.User = url.UserPassword(task.GetUsername(), task.GetPassword())
	}
	target := basicHandshakeProbeURL
	switch task.GetLevel() {
	case controlplanev1.CheckLevel_CHECK_LEVEL_BASIC:
		if task.GetRoutingName() != "" || task.GetTargetUrl() != "" {
			return time.Time{}, "", nil, errors.New("invalid basic task")
		}
	case controlplanev1.CheckLevel_CHECK_LEVEL_TARGET:
		if task.GetRoutingName() == "" {
			return time.Time{}, "", nil, errors.New("target routing is required")
		}
		target = task.GetTargetUrl()
	case controlplanev1.CheckLevel_CHECK_LEVEL_EGRESS:
		if task.GetRoutingName() != "" || task.GetTargetUrl() == "" {
			return time.Time{}, "", nil, errors.New("invalid egress task")
		}
		target = task.GetTargetUrl()
	default:
		return time.Time{}, "", nil, errors.New("unsupported task level")
	}
	parsedTarget, err := url.Parse(target)
	if err != nil || parsedTarget.Scheme == "" || parsedTarget.Host == "" || parsedTarget.User != nil || parsedTarget.Fragment != "" ||
		(parsedTarget.Scheme != "http" && parsedTarget.Scheme != "https") || strings.TrimSpace(task.GetHost()) != task.GetHost() {
		return time.Time{}, "", nil, errors.New("invalid probe target")
	}
	return deadline, parsedTarget.String(), proxyURL, nil
}

func minDuration(left, right time.Duration) time.Duration {
	if left <= 0 || left < right {
		return left
	}
	return right
}

func probeLatency(started time.Time) time.Duration {
	if latency := time.Since(started); latency > 0 {
		return latency
	}
	return time.Nanosecond
}