package oracle

import (
	"context"
	"fmt"
	"math"
	"sort"
	"strings"
	"time"

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

// TimeProbe issues ONE request against the candidate with an injected delay of
// `delay` (delay == 0 means a clean baseline request) and returns the observed
// round-trip latency. Supplied by the exploit layer, which knows how to shape a
// time-based payload (e.g. SLEEP(delay), pg_sleep, WAITFOR) for the target. The
// oracle calls it repeatedly to build baseline and attack samples.
type TimeProbe func(ctx context.Context, c verify.Candidate, delay time.Duration) (time.Duration, error)

// TimingConfig tunes the statistical test. Defaults are conservative so network
// jitter does not get mistaken for an induced delay.
type TimingConfig struct {
	Delay           time.Duration // injected delay for attack samples (default 5s)
	BaselineSamples int           // clean requests to characterise normal latency (default 5)
	AttackSamples   int           // delayed requests (default 3)
	SigmaK          float64       // induced delay must clear K baseline std-devs (default 3)
	MinInducedFrac  float64       // and at least this fraction of Delay (default 0.6)
	Ladder          bool          // also probe at 2×Delay and require the induced delay to SCALE (linearity) before proving — the gold standard against a coincidental slow response
}

func (c TimingConfig) withDefaults() TimingConfig {
	if c.Delay <= 0 {
		c.Delay = 5 * time.Second
	}
	if c.BaselineSamples <= 0 {
		c.BaselineSamples = 5
	}
	if c.AttackSamples <= 0 {
		c.AttackSamples = 3
	}
	if c.SigmaK <= 0 {
		c.SigmaK = 3
	}
	if c.MinInducedFrac <= 0 {
		c.MinInducedFrac = 0.6
	}
	return c
}

// Timing is a statistical time-based oracle for blind SQLi and race windows.
// Its ground truth is wall-clock latency the model cannot fake: if injecting a
// known delay reliably moves response time beyond the target's own noise floor
// by a margin proportional to that delay, the injection executed — Proven. If
// the delayed request is indistinguishable from baseline, it did not —
// Disproven. A weak or inconsistent signal is Inconclusive (never a guess),
// which is the honest answer under heavy jitter.
type Timing struct {
	probe   TimeProbe
	cfg     TimingConfig
	applies func(verify.Candidate) bool
}

// TimingOption configures a Timing oracle.
type TimingOption func(*Timing)

// WithTimingConfig overrides the statistical parameters.
func WithTimingConfig(cfg TimingConfig) TimingOption {
	return func(t *Timing) { t.cfg = cfg.withDefaults() }
}

// WithTimingApplies overrides which candidates this oracle probes.
func WithTimingApplies(fn func(verify.Candidate) bool) TimingOption {
	return func(t *Timing) {
		if fn != nil {
			t.applies = fn
		}
	}
}

// NewTiming builds a timing oracle over the given probe.
func NewTiming(probe TimeProbe, opts ...TimingOption) *Timing {
	t := &Timing{
		probe:   probe,
		cfg:     TimingConfig{}.withDefaults(),
		applies: defaultTimingApplies,
	}
	for _, o := range opts {
		o(t)
	}
	return t
}

// Name implements verify.Oracle.
func (t *Timing) Name() string { return "timing" }

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

// Probe samples baseline then attack latencies and decides from their
// separation, controlling for baseline jitter.
func (t *Timing) Probe(ctx context.Context, c verify.Candidate) (verify.OracleVerdict, float64, string, error) {
	baseline, err := t.sample(ctx, c, 0, t.cfg.BaselineSamples)
	if err != nil {
		return verify.Inconclusive, 0, "", err
	}
	attack, err := t.sample(ctx, c, t.cfg.Delay, t.cfg.AttackSamples)
	if err != nil {
		return verify.Inconclusive, 0, "", err
	}

	baseMed := median(baseline)
	atkMed := median(attack)
	induced := atkMed - baseMed

	// Noise floor: K baseline std-devs, but estimated robustly. A single slow
	// baseline sample inflates the classical std-dev; MAD (median absolute
	// deviation, scaled to a std-dev by 1.4826 for normal jitter) resists that.
	// Take the LARGER of the two estimates — the conservative direction, so an
	// outlier can only make proving HARDER, never easier.
	baseSigma := maxDur(stddev(baseline), robustSigma(baseline))
	noiseFloor := time.Duration(t.cfg.SigmaK * float64(baseSigma))
	strong := maxDur(noiseFloor, time.Duration(t.cfg.MinInducedFrac*float64(t.cfg.Delay)))

	// Consistency: a majority of attack samples must clear baseline + half the
	// injected delay, so one slow outlier can't carry the verdict.
	cut := baseMed + t.cfg.Delay/2
	over := 0
	for _, d := range attack {
		if d >= cut {
			over++
		}
	}
	consistent := float64(over) >= math.Ceil(float64(len(attack))/2)

	switch {
	case induced >= strong && consistent:
		// Single-delay evidence is strong. In ladder mode, demand that the
		// induced delay SCALES with the injected delay before proving — a
		// coincidental slow endpoint won't track a second, larger delay.
		if t.cfg.Ladder {
			return t.ladder(ctx, c, baseMed, induced)
		}
		return verify.Proven, 0.9,
			fmt.Sprintf("baseline median %s, attack median %s → induced +%s on a %s payload (≥%s, %d/%d samples over threshold)",
				rt(baseMed), rt(atkMed), rt(induced), rt(t.cfg.Delay), rt(strong), over, len(attack)), nil
	case induced < noiseFloor:
		return verify.Disproven, 0.85,
			fmt.Sprintf("attack median %s ≈ baseline median %s (induced +%s < noise floor %s) — no time-based effect",
				rt(atkMed), rt(baseMed), rt(induced), rt(noiseFloor)), nil
	default:
		return verify.Inconclusive, 0, "", nil
	}
}

// ladder confirms time-based injection by a second probe at twice the delay:
// the induced latency must grow by at least MinInducedFrac·Delay more, i.e. it
// tracks the injected delay. A response that is merely slow (not injectable)
// does not scale, so this is the gold-standard guard against a false positive.
func (t *Timing) ladder(ctx context.Context, c verify.Candidate, baseMed, induced1 time.Duration) (verify.OracleVerdict, float64, string, error) {
	d2 := t.cfg.Delay * 2
	attack2, err := t.sample(ctx, c, d2, t.cfg.AttackSamples)
	if err != nil {
		return verify.Inconclusive, 0, "", err
	}
	induced2 := median(attack2) - baseMed
	step := time.Duration(t.cfg.MinInducedFrac * float64(t.cfg.Delay))
	if induced2 >= induced1+step {
		return verify.Proven, 0.97,
			fmt.Sprintf("induced delay scales with payload: +%s at %s, +%s at %s (Δ≥%s) — confirmed time-based",
				rt(induced1), rt(t.cfg.Delay), rt(induced2), rt(d2), rt(step)), nil
	}
	// Strong at one delay but no scaling: ambiguous, not proof.
	return verify.Inconclusive, 0,
		fmt.Sprintf("induced +%s at %s did not scale to %s at %s — not confirmed", rt(induced1), rt(t.cfg.Delay), rt(induced2), rt(d2)), nil
}

// sample runs n probes at the given injected delay and returns the latencies.
func (t *Timing) sample(ctx context.Context, c verify.Candidate, delay time.Duration, n int) ([]time.Duration, error) {
	out := make([]time.Duration, 0, n)
	for i := 0; i < n; i++ {
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		d, err := t.probe(ctx, c, delay)
		if err != nil {
			return nil, err
		}
		out = append(out, d)
	}
	return out, nil
}

func defaultTimingApplies(c verify.Candidate) bool {
	if c.Type == blackboard.TypePotentialSQLI {
		return true
	}
	detail := strings.ToLower(c.Detail)
	for _, kw := range timingKeywords {
		if strings.Contains(detail, kw) {
			return true
		}
	}
	return false
}

var timingKeywords = []string{"time-based", "time based", "timing", "sleep", "blind", "race", "waitfor", "pg_sleep"}

// ---- small stats helpers (over time.Duration) ----------------------------

func median(xs []time.Duration) time.Duration {
	if len(xs) == 0 {
		return 0
	}
	cp := append([]time.Duration(nil), xs...)
	sort.Slice(cp, func(i, j int) bool { return cp[i] < cp[j] })
	n := len(cp)
	if n%2 == 1 {
		return cp[n/2]
	}
	return (cp[n/2-1] + cp[n/2]) / 2
}

func stddev(xs []time.Duration) time.Duration {
	if len(xs) < 2 {
		return 0
	}
	var sum float64
	for _, x := range xs {
		sum += float64(x)
	}
	mean := sum / float64(len(xs))
	var acc float64
	for _, x := range xs {
		d := float64(x) - mean
		acc += d * d
	}
	return time.Duration(math.Sqrt(acc / float64(len(xs)-1)))
}

// robustSigma estimates the baseline std-dev from the median absolute deviation
// (MAD), scaled by the 1.4826 constant that makes MAD a consistent estimator of
// σ for normally-distributed jitter. Unlike the classical std-dev it is barely
// moved by a single outlier sample.
func robustSigma(xs []time.Duration) time.Duration {
	if len(xs) < 2 {
		return 0
	}
	med := median(xs)
	dev := make([]time.Duration, len(xs))
	for i, x := range xs {
		d := x - med
		if d < 0 {
			d = -d
		}
		dev[i] = d
	}
	return time.Duration(1.4826 * float64(median(dev)))
}

func maxDur(a, b time.Duration) time.Duration {
	if a > b {
		return a
	}
	return b
}

// rt rounds a duration to the millisecond for readable proof strings.
func rt(d time.Duration) time.Duration { return d.Round(time.Millisecond) }