package exploit

import (
	"context"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/Armur-Ai/Pentest-Swarm-AI/internal/scope"
	"github.com/google/uuid"
)

func TestDeriveBOLATargets(t *testing.T) {
	in := []string{
		"http://t/identity/api/v2/vehicle/4bae9968-ec7f-4de3-a3a0-ba1b2ab5e5e5/location",
		"http://t/users/v1/42",
		"http://t/identity/api/v2/vehicle/{{victim_vehicle}}/location", // placeholder → same id slot, dedups with the uuid one
		"http://t/community/api/v2/community/posts/recent",             // no id segment → skipped
	}
	got := DeriveBOLATargets(in)
	want := map[string]bool{
		"http://t/identity/api/v2/vehicle/{id}/location": true,
		"http://t/users/v1/{id}":                         true,
	}
	if len(got) != 2 {
		t.Fatalf("expected 2 templated targets (uuid+placeholder vehicle collapse), got %d: %v", len(got), got)
	}
	for _, g := range got {
		if !want[g] {
			t.Errorf("unexpected target %q", g)
		}
	}
}

func TestSweepBOLA_DetectsCrossUserAccess(t *testing.T) {
	// Server returns 200 with data for any authenticated vehicle id — the BOLA.
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer jwt-xyz" {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(`{"latitude":"32.77","email":"victim@example.com"}`))
	}))
	defer srv.Close()

	e := NewExecutor(&scope.ScopeDefinition{AllowedCIDRs: []string{"127.0.0.1/32"}}, nil, false)
	templates := []string{srv.URL + "/vehicle/{id}/location"}
	refs := []ObjectRef{
		{Kind: "uuid", Value: "4bae9968-ec7f-4de3-a3a0-ba1b2ab5e5e5"},
		{Kind: "email", Value: "victim@example.com"}, // must be skipped (not path id)
	}
	hits := e.SweepBOLA(context.Background(), "Bearer jwt-xyz", templates, refs, uuid.New())
	if len(hits) != 1 {
		t.Fatalf("expected 1 BOLA hit, got %d: %+v", len(hits), hits)
	}
	if !contains(hits[0].Evidence, "32.77") {
		t.Errorf("hit should carry the leaked data as evidence: %q", hits[0].Evidence)
	}
}

// The probe callback must fire once per replay work-unit (skipping the email
// ref, which is never a path id) with the correct target + cross-user result,
// so the live dashboard/TUI can visualize the fan-out.
func TestSweepBOLAWithProbe_ReportsEachWorkUnit(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer jwt-xyz" {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(`{"latitude":"32.77"}`))
	}))
	defer srv.Close()

	e := NewExecutor(&scope.ScopeDefinition{AllowedCIDRs: []string{"127.0.0.1/32"}}, nil, false)
	templates := []string{srv.URL + "/vehicle/{id}/location"}
	refs := []ObjectRef{
		{Kind: "uuid", Value: "4bae9968-ec7f-4de3-a3a0-ba1b2ab5e5e5"},
		{Kind: "number", Value: "42"},
		{Kind: "email", Value: "victim@example.com"}, // skipped — not a path id
	}

	type probe struct {
		target string
		ok     bool
	}
	var probes []probe
	hits := e.SweepBOLAWithProbe(context.Background(), "Bearer jwt-xyz", templates, refs, uuid.New(),
		func(target string, ok bool) { probes = append(probes, probe{target, ok}) })

	// Two id refs replayed (email skipped) → two probe callbacks, both 200.
	if len(probes) != 2 {
		t.Fatalf("expected 2 probe callbacks (email ref skipped), got %d: %+v", len(probes), probes)
	}
	for _, p := range probes {
		if !p.ok {
			t.Errorf("probe %q should report ok=true (server returns 200 for any authed id)", p.target)
		}
		if !contains(p.target, "/vehicle/") {
			t.Errorf("probe target %q should be the templated endpoint", p.target)
		}
	}
	if len(hits) != 2 {
		t.Fatalf("expected 2 hits, got %d", len(hits))
	}
}

func contains(s, sub string) bool { return len(s) >= len(sub) && (stringIndex(s, sub) >= 0) }
func stringIndex(s, sub string) int {
	for i := 0; i+len(sub) <= len(s); i++ {
		if s[i:i+len(sub)] == sub {
			return i
		}
	}
	return -1
}