package oracle

import (
	"context"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/Armur-Ai/Pentest-Swarm-AI/internal/swarm/verify"
)

// RedirectProbe injects the candidate's attack payload (an attacker-controlled
// URL) and returns the response status and Location header, WITHOUT following
// the redirect. Supplied by the live layer.
type RedirectProbe func(ctx context.Context, c verify.Candidate) (status int, location string, err error)

// Redirect is an open-redirect oracle. Ground truth: if a parameter's value is
// reflected into a 3xx Location header pointing at an attacker-controlled host,
// the endpoint redirects wherever the attacker says — a real open redirect, not
// a model's guess. A non-3xx response, or a Location that stays on the target's
// own host, refutes it.
type Redirect struct {
	probe   RedirectProbe
	applies func(verify.Candidate) bool
}

// RedirectOption configures a Redirect oracle.
type RedirectOption func(*Redirect)

// WithRedirectApplies overrides which candidates this oracle probes.
func WithRedirectApplies(fn func(verify.Candidate) bool) RedirectOption {
	return func(o *Redirect) {
		if fn != nil {
			o.applies = fn
		}
	}
}

// NewRedirect builds an open-redirect oracle over the given probe.
func NewRedirect(probe RedirectProbe, opts ...RedirectOption) *Redirect {
	o := &Redirect{probe: probe, applies: RequireSpec(verify.ProbeDifferential)}
	for _, opt := range opts {
		opt(o)
	}
	return o
}

// Name implements verify.Oracle.
func (o *Redirect) Name() string { return "open-redirect" }

// Applies implements verify.Oracle.
func (o *Redirect) Applies(c verify.Candidate) bool {
	if o.probe == nil {
		return false
	}
	return o.applies(c)
}

// Probe injects the attack URL and proves the finding iff the response is a 3xx
// whose Location points at the attacker-controlled host from the attack payload.
func (o *Redirect) Probe(ctx context.Context, c verify.Candidate) (verify.OracleVerdict, float64, string, error) {
	status, loc, err := o.probe(ctx, c)
	if err != nil {
		return verify.Inconclusive, 0, "", err
	}
	attackHost := hostOf(c.Spec.AttackPayload)
	if attackHost == "" {
		return verify.Inconclusive, 0, "", fmt.Errorf("open-redirect: attack payload has no host")
	}
	is3xx := status >= 300 && status < 400
	if is3xx && locationPointsAt(loc, attackHost) {
		return verify.Proven, 0.95,
			fmt.Sprintf("%d redirect with Location → attacker-controlled host %q (%s)", status, attackHost, loc), nil
	}
	if !is3xx || loc == "" {
		return verify.Disproven, 0.8,
			fmt.Sprintf("no attacker-controlled redirect (status %d, Location %q)", status, loc), nil
	}
	return verify.Inconclusive, 0, "", nil
}

// hostOf extracts the host from an attack payload that is (or starts with) a URL.
func hostOf(payload string) string {
	p := strings.TrimSpace(payload)
	if strings.HasPrefix(p, "//") {
		p = "http:" + p
	}
	u, err := url.Parse(p)
	if err != nil {
		return ""
	}
	return u.Hostname()
}

// locationPointsAt reports whether a Location header sends the browser to host
// (absolute URL to that host, or scheme-relative //host/...). A relative path or
// a Location on a different (non-attacker) host does not count.
func locationPointsAt(loc, host string) bool {
	loc = strings.TrimSpace(loc)
	if loc == "" || host == "" {
		return false
	}
	if strings.HasPrefix(loc, "//") {
		loc = "http:" + loc
	}
	u, err := url.Parse(loc)
	if err != nil {
		return false
	}
	return strings.EqualFold(u.Hostname(), host)
}

// DefaultClientNoRedirect is DefaultClient but does not follow redirects, so a
// 3xx and its Location are observable (required by the open-redirect probe).
func DefaultClientNoRedirect() *http.Client {
	return &http.Client{
		Timeout:       30 * time.Second,
		CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse },
	}
}