package oracle

import (
	"context"
	"errors"
	"testing"
	"time"

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

// scriptProbe returns baseline latencies (for delay==0) and attack latencies
// (for delay>0), consumed in order per mode.
func scriptProbe(baseline, attack []time.Duration, err error) TimeProbe {
	bi, ai := 0, 0
	return func(_ context.Context, _ verify.Candidate, delay time.Duration) (time.Duration, error) {
		if err != nil {
			return 0, err
		}
		if delay == 0 {
			d := baseline[bi%len(baseline)]
			bi++
			return d, nil
		}
		d := attack[ai%len(attack)]
		ai++
		return d, nil
	}
}

func sqliTimingCand() verify.Candidate {
	return verify.Candidate{Type: blackboard.TypePotentialSQLI, Detail: "time-based blind?"}
}

func fastCfg() TimingConfig {
	// Short injected delay keeps tests instant; probes return scripted values.
	return TimingConfig{Delay: 200 * time.Millisecond, BaselineSamples: 5, AttackSamples: 3}
}

func ms(v ...int) []time.Duration {
	out := make([]time.Duration, len(v))
	for i, x := range v {
		out[i] = time.Duration(x) * time.Millisecond
	}
	return out
}

func TestTiming_Applies(t *testing.T) {
	o := NewTiming(scriptProbe(ms(1), ms(1), nil))
	if !o.Applies(sqliTimingCand()) {
		t.Error("blind SQLi should apply")
	}
	if !o.Applies(verify.Candidate{Type: blackboard.TypeCVEMatch, Detail: "race condition window"}) {
		t.Error("race condition should apply by keyword")
	}
	if o.Applies(verify.Candidate{Type: blackboard.TypeMisconfig, Detail: "cors"}) {
		t.Error("unrelated finding should not apply")
	}
}

func TestTiming_NilProbeNeverApplies(t *testing.T) {
	if NewTiming(nil).Applies(sqliTimingCand()) {
		t.Error("nil-probe oracle must never apply")
	}
}

// A clean, consistent induced delay matching the payload = Proven.
func TestTiming_InducedDelayIsProven(t *testing.T) {
	// baseline ~20ms low jitter; attack ~ baseline + 200ms payload.
	o := NewTiming(
		scriptProbe(ms(18, 20, 22, 19, 21), ms(221, 219, 223), nil),
		WithTimingConfig(fastCfg()),
	)
	v, conf, proof, err := o.Probe(context.Background(), sqliTimingCand())
	if err != nil || v != verify.Proven {
		t.Fatalf("verdict=%v err=%v, want proven; proof=%q", v, err, proof)
	}
	if conf < 0.8 {
		t.Errorf("conf=%.2f too low", conf)
	}
}

// Attack latency indistinguishable from baseline = Disproven.
func TestTiming_NoEffectIsDisproven(t *testing.T) {
	o := NewTiming(
		scriptProbe(ms(18, 20, 22, 19, 21), ms(20, 19, 21), nil),
		WithTimingConfig(fastCfg()),
	)
	v, _, _, _ := o.Probe(context.Background(), sqliTimingCand())
	if v != verify.Disproven {
		t.Fatalf("verdict = %v, want disproven (no induced delay)", v)
	}
}

// A partial/inconsistent signal (one slow outlier, rest baseline) is
// Inconclusive, not a false Proven.
func TestTiming_InconsistentIsInconclusive(t *testing.T) {
	// One attack sample is slow, two are baseline → induced median ~baseline,
	// but above the noise floor region: lands in the ambiguous band.
	o := NewTiming(
		scriptProbe(ms(18, 20, 22, 19, 21), ms(260, 60, 65), nil),
		WithTimingConfig(TimingConfig{Delay: 200 * time.Millisecond, BaselineSamples: 5, AttackSamples: 3, SigmaK: 3, MinInducedFrac: 0.6}),
	)
	v, _, _, _ := o.Probe(context.Background(), sqliTimingCand())
	if v != verify.Inconclusive {
		t.Fatalf("verdict = %v, want inconclusive (inconsistent samples)", v)
	}
}

// A probe error surfaces (→ Inconclusive at the corroborator).
func TestTiming_ProbeErrorSurfaced(t *testing.T) {
	o := NewTiming(scriptProbe(nil, nil, errors.New("timeout")), WithTimingConfig(fastCfg()))
	_, _, _, err := o.Probe(context.Background(), sqliTimingCand())
	if err == nil {
		t.Fatal("want probe error surfaced")
	}
}

func TestTiming_ContextCancel(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	o := NewTiming(scriptProbe(ms(20), ms(220), nil), WithTimingConfig(fastCfg()))
	if _, _, _, err := o.Probe(ctx, sqliTimingCand()); err == nil {
		t.Error("cancelled context should error")
	}
}

func TestTimingStats(t *testing.T) {
	if got := median(ms(10, 30, 20)); got != 20*time.Millisecond {
		t.Errorf("median = %v, want 20ms", got)
	}
	if got := median(ms(10, 20, 30, 40)); got != 25*time.Millisecond {
		t.Errorf("even median = %v, want 25ms", got)
	}
	if stddev(ms(20)) != 0 {
		t.Error("stddev of a single sample should be 0")
	}
	if stddev(ms(10, 10, 10)) != 0 {
		t.Error("stddev of constant samples should be 0")
	}
}

// Ladder mode: induced delay that scales with the injected delay = Proven.
func TestTiming_LadderScalesIsProven(t *testing.T) {
	// baseline ~20ms; @200ms → ~220ms (+200); @400ms → ~420ms (+400). Scales.
	probe := func(baseline, atk1, atk2 []time.Duration) TimeProbe {
		bi, a1, a2 := 0, 0, 0
		return func(_ context.Context, _ verify.Candidate, delay time.Duration) (time.Duration, error) {
			switch {
			case delay == 0:
				d := baseline[bi%len(baseline)]
				bi++
				return d, nil
			case delay == 200*time.Millisecond:
				d := atk1[a1%len(atk1)]
				a1++
				return d, nil
			default: // 400ms
				d := atk2[a2%len(atk2)]
				a2++
				return d, nil
			}
		}
	}
	cfg := TimingConfig{Delay: 200 * time.Millisecond, BaselineSamples: 5, AttackSamples: 3, Ladder: true}
	o := NewTiming(probe(ms(18, 20, 22, 19, 21), ms(221, 219, 223), ms(421, 419, 423)), WithTimingConfig(cfg))

	v, conf, _, err := o.Probe(context.Background(), sqliTimingCand())
	if err != nil || v != verify.Proven {
		t.Fatalf("verdict=%v err=%v, want proven (scales)", v, err)
	}
	if conf < 0.95 {
		t.Errorf("ladder-proven conf = %.2f, want high", conf)
	}
}

// Ladder mode: strong at one delay but no scaling (a merely-slow endpoint) =
// Inconclusive, not a false Proven.
func TestTiming_LadderNoScaleIsInconclusive(t *testing.T) {
	probe := func(_ context.Context, _ verify.Candidate, delay time.Duration) (time.Duration, error) {
		switch delay {
		case 0:
			return 20 * time.Millisecond, nil
		case 200 * time.Millisecond:
			return 230 * time.Millisecond, nil // +210, strong at D
		default:
			return 235 * time.Millisecond, nil // @400ms still ~+215: did NOT scale
		}
	}
	cfg := TimingConfig{Delay: 200 * time.Millisecond, BaselineSamples: 5, AttackSamples: 3, Ladder: true}
	o := NewTiming(probe, WithTimingConfig(cfg))

	v, _, _, _ := o.Probe(context.Background(), sqliTimingCand())
	if v != verify.Inconclusive {
		t.Fatalf("verdict = %v, want inconclusive (no scaling)", v)
	}
}

func TestTiming_ImplementsOracle(t *testing.T) {
	var _ verify.Oracle = NewTiming(scriptProbe(ms(1), ms(1), nil))
}

func TestRobustSigma(t *testing.T) {
	// Tight cluster: robust sigma is small.
	if s := robustSigma(ms(20, 21, 19, 20, 21)); s > 5*time.Millisecond {
		t.Errorf("robustSigma of a tight cluster = %v, want small", s)
	}
	// One wild outlier barely moves MAD (median of deviations), unlike stddev.
	tight := ms(20, 21, 19, 20, 21)
	withOutlier := ms(20, 21, 19, 20, 21, 5000)
	if robustSigma(withOutlier) > 10*robustSigma(tight)+time.Millisecond {
		// robust estimate should stay bounded; classical stddev would explode
		t.Logf("robust sigma stayed bounded under an outlier (good)")
	}
	if stddev(withOutlier) <= robustSigma(withOutlier) {
		t.Error("classical stddev should be far larger than robust sigma under an outlier")
	}
	if robustSigma(ms(10)) != 0 {
		t.Error("single sample robustSigma should be 0")
	}
}

// An outlier-heavy baseline must not let a non-injectable endpoint prove: the
// conservative noise floor keeps it Inconclusive/Disproven.
func TestTiming_OutlierBaselineStaysConservative(t *testing.T) {
	// baseline mostly ~20ms but one 400ms spike; attack ~same as baseline.
	o := NewTiming(
		scriptProbe(ms(20, 19, 21, 400, 20), ms(60, 62, 61), nil),
		WithTimingConfig(TimingConfig{Delay: 200 * time.Millisecond, BaselineSamples: 5, AttackSamples: 3}),
	)
	v, _, _, _ := o.Probe(context.Background(), sqliTimingCand())
	if v == verify.Proven {
		t.Error("a jittery baseline + weak induced delay must not prove")
	}
}