package oracle

import (
	"context"
	"errors"
	"math"
	"testing"

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

func idorCand() verify.Candidate {
	return verify.Candidate{Target: "https://t/api/user/2", Type: blackboard.TypeCVEMatch, Detail: "possible IDOR on /user/{id}"}
}

func diffProbe(attack, control Response, err error) DiffProbe {
	return func(context.Context, verify.Candidate) (Response, Response, error) { return attack, control, err }
}

func TestDifferential_Applies(t *testing.T) {
	d := NewDifferential(diffProbe(Response{}, Response{}, nil))
	cases := []struct {
		c    verify.Candidate
		want bool
	}{
		{idorCand(), true},
		{verify.Candidate{Type: blackboard.TypePotentialSQLI}, true},
		{verify.Candidate{Type: blackboard.TypeExploitChain, Detail: "BFLA on admin route"}, true},
		{verify.Candidate{Type: blackboard.TypeMisconfig, Detail: "missing security headers"}, false},
	}
	for _, tc := range cases {
		if got := d.Applies(tc.c); got != tc.want {
			t.Errorf("Applies(%q) = %v, want %v", tc.c.Detail, got, tc.want)
		}
	}
}

func TestDifferential_NilProbeNeverApplies(t *testing.T) {
	if NewDifferential(nil).Applies(idorCand()) {
		t.Error("a nil-probe oracle must never apply")
	}
}

// Success against a denied control = crossed authorization boundary = Proven.
func TestDifferential_BoundaryCrossedIsProven(t *testing.T) {
	d := NewDifferential(diffProbe(
		Response{Status: 200, Body: `{"user":"victim","ssn":"123"}`},
		Response{Status: 403, Body: `{"error":"forbidden"}`},
		nil,
	))
	v, conf, proof, err := d.Probe(context.Background(), idorCand())
	if err != nil || v != verify.Proven {
		t.Fatalf("verdict=%v err=%v, want proven", v, err)
	}
	if conf < 0.9 {
		t.Errorf("conf = %.2f, want high", conf)
	}
	if proof == "" {
		t.Error("expected a proof receipt")
	}
}

// Attack response collapses onto the control = no divergence = Disproven.
func TestDifferential_CollapseOntoControlIsDisproven(t *testing.T) {
	body := `{"error":"forbidden"}`
	d := NewDifferential(diffProbe(
		Response{Status: 403, Body: body},
		Response{Status: 403, Body: body},
		nil,
	))
	v, _, _, _ := d.Probe(context.Background(), idorCand())
	if v != verify.Disproven {
		t.Fatalf("verdict = %v, want disproven (attack == control)", v)
	}
}

// Both 2xx but materially different bodies = differential response = Proven.
func TestDifferential_BodyDivergenceIsProven(t *testing.T) {
	d := NewDifferential(diffProbe(
		Response{Status: 200, Body: "alice bob carol dave erin frank grace heidi"},
		Response{Status: 200, Body: "zoe yan xavier wendy victor ulysses tara sam"},
		nil,
	))
	v, _, _, _ := d.Probe(context.Background(), idorCand())
	if v != verify.Proven {
		t.Fatalf("verdict = %v, want proven (bodies diverge)", v)
	}
}

// Both 2xx with near-identical bodies = equivalent = Disproven.
func TestDifferential_EquivalentBodiesDisproven(t *testing.T) {
	d := NewDifferential(diffProbe(
		Response{Status: 200, Body: "the same list of fields one two three four five"},
		Response{Status: 200, Body: "the same list of fields one two three four five"},
		nil,
	))
	v, _, _, _ := d.Probe(context.Background(), idorCand())
	if v != verify.Disproven {
		t.Fatalf("verdict = %v, want disproven (equivalent)", v)
	}
}

// A probe error is surfaced (→ Inconclusive at the corroborator).
func TestDifferential_ProbeErrorSurfaced(t *testing.T) {
	d := NewDifferential(diffProbe(Response{}, Response{}, errors.New("network")))
	v, _, _, err := d.Probe(context.Background(), idorCand())
	if err == nil || v != verify.Inconclusive {
		t.Fatalf("verdict=%v err=%v, want inconclusive+err", v, err)
	}
}

// An ambiguous pair (both denied, different reasons) is Inconclusive, not a guess.
func TestDifferential_AmbiguousIsInconclusive(t *testing.T) {
	d := NewDifferential(diffProbe(
		Response{Status: 500, Body: "stacktrace A"},
		Response{Status: 400, Body: "bad request B"},
		nil,
	))
	v, _, _, _ := d.Probe(context.Background(), idorCand())
	if v != verify.Inconclusive {
		t.Fatalf("verdict = %v, want inconclusive", v)
	}
}

func TestBodySimilarity(t *testing.T) {
	cases := []struct {
		a, b string
		want float64
	}{
		{"", "", 1.0},
		{"a b c", "a b c", 1.0},
		{"a b c d", "", 0.0},
		{"a b c d", "a b c e", 3.0 / 5.0}, // inter=3 {a,b,c}, union=5 {a,b,c,d,e}
	}
	for _, tc := range cases {
		if got := bodySimilarity(tc.a, tc.b); math.Abs(got-tc.want) > 1e-9 {
			t.Errorf("bodySimilarity(%q,%q) = %.4f, want %.4f", tc.a, tc.b, got, tc.want)
		}
	}
}

func idorCandWithIndicator(ind string) verify.Candidate {
	c := idorCand()
	c.Spec = &verify.ProbeSpec{Class: verify.ProbeDifferential, Indicator: ind}
	return c
}

// Indicator present in attack but not control = strongest Proven signal.
func TestDifferential_IndicatorPresentOnlyInAttackIsProven(t *testing.T) {
	d := NewDifferential(diffProbe(
		Response{Status: 200, Body: `{"ssn":"555-12-3456"}`},
		Response{Status: 200, Body: `{"ssn":"000-00-0000"}`},
		nil,
	))
	v, conf, _, _ := d.Probe(context.Background(), idorCandWithIndicator("555-12-3456"))
	if v != verify.Proven || conf < 0.95 {
		t.Fatalf("verdict=%v conf=%.2f, want strong proven via indicator", v, conf)
	}
}

// Indicator absent from both = the named marker never fired = Disproven.
func TestDifferential_IndicatorAbsentFromBothIsDisproven(t *testing.T) {
	d := NewDifferential(diffProbe(
		Response{Status: 200, Body: "nothing useful"},
		Response{Status: 403, Body: "forbidden"},
		nil,
	))
	v, _, _, _ := d.Probe(context.Background(), idorCandWithIndicator("SECRET-TOKEN"))
	if v != verify.Disproven {
		t.Fatalf("verdict = %v, want disproven (indicator never appeared)", v)
	}
}

// Indicator in both responses falls through to the status/body heuristics (not
// a discriminator), so a boundary-cross still proves.
func TestDifferential_IndicatorInBothFallsThrough(t *testing.T) {
	d := NewDifferential(diffProbe(
		Response{Status: 200, Body: "shared banner X"},
		Response{Status: 403, Body: "shared banner X"},
		nil,
	))
	v, _, _, _ := d.Probe(context.Background(), idorCandWithIndicator("shared banner"))
	if v != verify.Proven { // boundary crossed (200 vs 403)
		t.Fatalf("verdict = %v, want proven via boundary heuristic", v)
	}
}

func TestDifferential_ImplementsOracle(t *testing.T) {
	var _ verify.Oracle = NewDifferential(diffProbe(Response{}, Response{}, nil))
}

// An endpoint that merely echoes the id must NOT be proven by body divergence:
// after normalising the payloads out, the two bodies are identical.
func TestDifferential_EchoedIDIsNotDivergence(t *testing.T) {
	c := idorCand()
	c.Spec = &verify.ProbeSpec{Class: verify.ProbeDifferential, AttackPayload: "2", ControlPayload: "999999"}
	d := NewDifferential(diffProbe(
		Response{Status: 200, Body: `{"id":"2","note":"your order"}`},
		Response{Status: 200, Body: `{"id":"999999","note":"your order"}`},
		nil,
	))
	v, _, _, _ := d.Probe(context.Background(), c)
	if v == verify.Proven {
		t.Error("an endpoint that only echoes the id must not be proven as a differential")
	}
	if v != verify.Disproven {
		t.Errorf("verdict = %v, want disproven (bodies equivalent once id normalised)", v)
	}
}

// Real privileged-data divergence (beyond the echoed id) still proves.
func TestDifferential_RealDivergenceStillProves(t *testing.T) {
	c := idorCand()
	c.Spec = &verify.ProbeSpec{Class: verify.ProbeDifferential, AttackPayload: "2", ControlPayload: "999999"}
	d := NewDifferential(diffProbe(
		Response{Status: 200, Body: `{"id":"2","ssn":"555-12-3456","card":"4111"}`},
		Response{Status: 200, Body: `{"id":"999999","error":"not found"}`},
		nil,
	))
	v, _, _, _ := d.Probe(context.Background(), c)
	if v != verify.Proven {
		t.Errorf("verdict = %v, want proven (real content divergence beyond the id)", v)
	}
}