package agents

import (
	"context"
	"testing"

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

func sqlmapVuln(injections ...pipeline.SQLiInjectionPoint) pipeline.VulnerabilityRecord {
	return pipeline.VulnerabilityRecord{
		Tool: "sqlmap", Title: "SQL Injection", URL: "https://t/items?id=1",
		SQLiInjections: injections,
	}
}

func TestVerifySQLi_StampsWhenProven(t *testing.T) {
	a := &ReconAgent{}
	var gotSpec *verify.ProbeSpec
	a.SetSQLiVerifier(func(_ context.Context, target string, spec *verify.ProbeSpec) (bool, string) {
		gotSpec = spec
		return true, "induced +4.9s on a 5s payload"
	})

	cf := pipeline.ClassifiedFinding{Title: "SQL Injection"}
	a.verifySQLi(context.Background(), &cf, sqlmapVuln(pipeline.SQLiInjectionPoint{
		Place: "GET", Parameter: "id", DBMS: "MySQL",
	}))

	if !cf.Verified || cf.VerifiedBy == "" {
		t.Fatalf("expected VERIFIED stamp, got %+v", cf)
	}
	if gotSpec == nil || gotSpec.Class != verify.ProbeTiming {
		t.Errorf("verifier got spec %+v, want a timing spec", gotSpec)
	}
}

func TestVerifySQLi_NoStampWhenUnproven(t *testing.T) {
	a := &ReconAgent{}
	a.SetSQLiVerifier(func(context.Context, string, *verify.ProbeSpec) (bool, string) { return false, "" })
	cf := pipeline.ClassifiedFinding{Title: "SQL Injection"}
	a.verifySQLi(context.Background(), &cf, sqlmapVuln(pipeline.SQLiInjectionPoint{Parameter: "id", DBMS: "MySQL"}))
	if cf.Verified {
		t.Error("unproven verdict must not stamp VERIFIED")
	}
}

func TestVerifySQLi_NilVerifierNoop(t *testing.T) {
	a := &ReconAgent{} // no verifier
	cf := pipeline.ClassifiedFinding{Title: "SQL Injection"}
	a.verifySQLi(context.Background(), &cf, sqlmapVuln(pipeline.SQLiInjectionPoint{Parameter: "id", DBMS: "MySQL"}))
	if cf.Verified {
		t.Error("no verifier must be a no-op")
	}
}

func TestVerifySQLi_NoInjectionPointsNoop(t *testing.T) {
	a := &ReconAgent{}
	called := false
	a.SetSQLiVerifier(func(context.Context, string, *verify.ProbeSpec) (bool, string) { called = true; return true, "x" })
	cf := pipeline.ClassifiedFinding{Title: "SQL Injection"}
	a.verifySQLi(context.Background(), &cf, sqlmapVuln()) // no injection points
	if cf.Verified || called {
		t.Error("a finding with no injection points must not invoke the verifier or stamp")
	}
}

func TestVerifySQLi_SkipsUnbuildableAndUsesNext(t *testing.T) {
	a := &ReconAgent{}
	calls := 0
	a.SetSQLiVerifier(func(_ context.Context, _ string, spec *verify.ProbeSpec) (bool, string) {
		calls++
		return true, "proven"
	})
	cf := pipeline.ClassifiedFinding{Title: "SQL Injection"}
	// First point has an unknown DBMS (no spec builds); the second is usable.
	a.verifySQLi(context.Background(), &cf, sqlmapVuln(
		pipeline.SQLiInjectionPoint{Parameter: "id", DBMS: "CouchDB"},
		pipeline.SQLiInjectionPoint{Place: "GET", Parameter: "id", DBMS: "MySQL"},
	))
	if !cf.Verified {
		t.Fatal("should fall through to the buildable injection point")
	}
	if calls != 1 {
		t.Errorf("verifier called %d times, want 1 (unbuildable point skipped before the call)", calls)
	}
}