package oracle

import (
	"context"
	"fmt"
	"strings"

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

// Response is the minimal shape of an HTTP response the differential oracle
// reasons over. The exploit layer fills it from a real request.
type Response struct {
	Status int    // HTTP status code
	Body   string // response body (already read)
}

// DiffProbe performs the two requests a differential verdict needs and returns
// their responses: the attack response (the privileged / injected request) and
// the control response (the negative control — the same shape that SHOULD be
// denied, or a known-good baseline). Supplied by the exploit layer, which knows
// the request shapes per class; the oracle only compares the two. An error
// means the pair could not be gathered (→ Inconclusive at the corroborator).
type DiffProbe func(ctx context.Context, c verify.Candidate) (attack, control Response, err error)

// Differential is a differential-response oracle for authorization and
// injection classes (BOLA/IDOR, BFLA, auth bypass, data-returning injection).
// Its ground-truth signal is divergence between two responses the model never
// touches: if a request that SHOULD be denied returns privileged success while
// the control is denied, that boundary was really crossed — Proven. If the
// attack response collapses onto the control, the differential vanished and the
// finding is refuted — Disproven. Genuine ambiguity is Inconclusive.
type Differential struct {
	probe      DiffProbe
	applies    func(verify.Candidate) bool
	similarity float64 // bodies at/above this ratio count as "the same" (default 0.95)
}

// DiffOption configures a Differential oracle.
type DiffOption func(*Differential)

// WithSimilarity sets the body-similarity ratio at/above which two bodies are
// treated as equivalent (0..1; default 0.95).
func WithSimilarity(r float64) DiffOption {
	return func(d *Differential) {
		if r > 0 && r <= 1 {
			d.similarity = r
		}
	}
}

// WithDiffApplies overrides which candidates this oracle probes.
func WithDiffApplies(fn func(verify.Candidate) bool) DiffOption {
	return func(d *Differential) {
		if fn != nil {
			d.applies = fn
		}
	}
}

// NewDifferential builds a differential oracle over the given probe.
func NewDifferential(probe DiffProbe, opts ...DiffOption) *Differential {
	d := &Differential{
		probe:      probe,
		applies:    defaultDiffApplies,
		similarity: 0.95,
	}
	for _, o := range opts {
		o(d)
	}
	return d
}

// Name implements verify.Oracle.
func (d *Differential) Name() string { return "differential-response" }

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

// Probe gathers the attack/control pair and returns a verdict from their
// divergence. The decision rule, in order:
//
//  1. Attack succeeds (2xx) while control is denied (401/403/404) → Proven:
//     the authorization boundary was crossed for real.
//  2. Attack and control are equivalent (same status, near-identical body) →
//     Disproven: the attack had no privileged effect; the differential
//     collapsed.
//  3. Attack succeeds and its body diverges materially from the control →
//     Proven: a data-returning divergence (e.g. injection yielding different
//     rows).
//  4. Anything else → Inconclusive.
func (d *Differential) Probe(ctx context.Context, c verify.Candidate) (verify.OracleVerdict, float64, string, error) {
	attack, control, err := d.probe(ctx, c)
	if err != nil {
		return verify.Inconclusive, 0, "", err
	}

	attackOK := attack.Status >= 200 && attack.Status < 300
	controlDenied := control.Status == 401 || control.Status == 403 || control.Status == 404
	// Compare bodies with the injected payloads themselves normalised out: an
	// endpoint that merely ECHOES the id would otherwise look "divergent" just
	// because the attack body says "2" and the control says "999999" — a
	// trivial difference, not evidence of privileged data. Stripping both
	// payloads keeps only real content divergence.
	aBody, cBody := attack.Body, control.Body
	if c.Spec != nil {
		aBody = normalizePayload(aBody, c.Spec.AttackPayload)
		cBody = normalizePayload(cBody, c.Spec.ControlPayload)
	}
	sim := bodySimilarity(aBody, cBody)
	equivalent := attack.Status == control.Status && sim >= d.similarity

	// Strongest, most specific signal: a known success indicator that appears
	// in the attack response but NOT the control. This is the discriminator the
	// finding's author named, so it outranks the generic status/body heuristics.
	if ind := candidateIndicator(c); ind != "" {
		inAttack := strings.Contains(attack.Body, ind)
		inControl := strings.Contains(control.Body, ind)
		switch {
		case inAttack && !inControl:
			return verify.Proven, 0.98,
				fmt.Sprintf("indicator %q present in attack (%d), absent in control (%d)", ind, attack.Status, control.Status), nil
		case !inAttack && !inControl:
			// The named success marker never appeared — the exploit didn't land.
			return verify.Disproven, 0.9,
				fmt.Sprintf("indicator %q absent from both attack (%d) and control (%d)", ind, attack.Status, control.Status), nil
		}
	}

	switch {
	case attackOK && controlDenied:
		return verify.Proven, 0.97,
			fmt.Sprintf("attack %d vs control %d — authorization boundary crossed", attack.Status, control.Status), nil
	case equivalent:
		return verify.Disproven, 0.9,
			fmt.Sprintf("attack %d ≈ control %d (body similarity %.2f) — no privileged divergence", attack.Status, control.Status, sim), nil
	case attackOK && sim < d.similarity:
		return verify.Proven, 0.85,
			fmt.Sprintf("attack %d diverges from control %d (body similarity %.2f) — differential response", attack.Status, control.Status, sim), nil
	default:
		return verify.Inconclusive, 0, "", nil
	}
}

// bodySimilarity returns a cheap 0..1 similarity between two bodies: 1.0 when
// equal, otherwise a token-Jaccard over whitespace-split fields (robust to
// per-request noise like timestamps/CSRF tokens that length-only would miss).
// Two empty bodies are considered identical.
func bodySimilarity(a, b string) float64 {
	if a == b {
		return 1.0
	}
	fa := strings.Fields(a)
	fb := strings.Fields(b)
	if len(fa) == 0 && len(fb) == 0 {
		return 1.0
	}
	if len(fa) == 0 || len(fb) == 0 {
		return 0.0
	}
	a1 := make(map[string]struct{}, len(fa))
	for _, t := range fa {
		a1[t] = struct{}{}
	}
	b1 := make(map[string]struct{}, len(fb))
	for _, t := range fb {
		b1[t] = struct{}{}
	}
	inter := 0
	for t := range a1 {
		if _, ok := b1[t]; ok {
			inter++
		}
	}
	union := len(a1) + len(b1) - inter // |A ∪ B| by inclusion–exclusion
	if union == 0 {
		return 1.0
	}
	return float64(inter) / float64(union)
}

// normalizePayload removes the echoed injected value from a response body so a
// reflected id doesn't masquerade as content divergence. Empty payload is a
// no-op.
func normalizePayload(body, payload string) string {
	if payload == "" {
		return body
	}
	return strings.ReplaceAll(body, payload, "§ID§")
}

// candidateIndicator returns the success-indicator substring the finding's spec
// named, if any — the discriminator that most specifically confirms the class.
func candidateIndicator(c verify.Candidate) string {
	if c.Spec != nil {
		return c.Spec.Indicator
	}
	return ""
}

// defaultDiffApplies matches authorization and injection classes by type or by
// a detail-keyword heuristic.
func defaultDiffApplies(c verify.Candidate) bool {
	if c.Type == blackboard.TypePotentialSQLI {
		return true
	}
	detail := strings.ToLower(c.Detail)
	for _, kw := range diffKeywords {
		if strings.Contains(detail, kw) {
			return true
		}
	}
	return false
}

var diffKeywords = []string{
	"idor", "bola", "bfla", "authorization", "authz", "privilege",
	"access control", "auth bypass", "broken access", "mass assignment",
	"injection", "sqli", "sql injection",
}