package agents

import (
	"context"
	"encoding/json"
	"strings"
	"testing"
	"time"

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

// newSSRFAgent builds a minimal exploit agent for sweep tests (no inner/executor
// needed — ssrfSweep only touches the board and the injected verifier).
func newSSRFAgent(campaignID uuid.UUID) *ExploitAgent {
	return &ExploitAgent{campaignID: campaignID, seenSSRF: make(map[string]struct{})}
}

// writeEndpoint stores an HTTP endpoint on the board (Write requires an agent
// name, which the trigger-only endpointFinding helper omits).
func writeEndpoint(ctx context.Context, board blackboard.Board, campaignID uuid.UUID, urlStr string) {
	data, _ := json.Marshal(pipeline.EndpointRecord{URL: urlStr, Method: "GET"})
	_, _ = board.Write(ctx, blackboard.Finding{
		CampaignID: campaignID, AgentName: "recon", Type: blackboard.TypeHTTPEndpoint,
		Target: urlStr, Data: data, PheromoneBase: 0.6, HalfLifeSec: 3600,
	})
}

func ssrfFindings(t *testing.T, ctx context.Context, board blackboard.Board) []pipeline.ClassifiedFinding {
	t.Helper()
	fs, err := board.Query(ctx, blackboard.Predicate{Types: []blackboard.FindingType{blackboard.TypeMisconfig}})
	if err != nil {
		t.Fatal(err)
	}
	var out []pipeline.ClassifiedFinding
	for _, f := range fs {
		var cf pipeline.ClassifiedFinding
		if json.Unmarshal(f.Data, &cf) == nil && cf.AttackCategory == "ssrf" {
			out = append(out, cf)
		}
	}
	return out
}

func TestSSRFSweep_PublishesVerifiedSinks(t *testing.T) {
	campaignID := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, campaignID, "http://t/fetch?url=http://seed")
	writeEndpoint(ctx, board, campaignID, "http://t/page?id=1") // no SSRF sink

	a := newSSRFAgent(campaignID)
	var probed []string
	a.SetSSRFVerifier(func(_ context.Context, target string, spec *verify.ProbeSpec) (bool, string) {
		probed = append(probed, target)
		return spec.Class == verify.ProbeOOB, "http interaction from 127.0.0.1"
	})

	a.ssrfSweep(ctx, board, nil)

	if len(probed) != 1 || !strings.Contains(probed[0], "/fetch") {
		t.Fatalf("expected exactly the sink endpoint probed, got %v", probed)
	}
	found := ssrfFindings(t, ctx, board)
	if len(found) != 1 {
		t.Fatalf("expected 1 verified SSRF finding, got %d", len(found))
	}
	if !found[0].Verified || found[0].Severity != pipeline.SeverityHigh {
		t.Errorf("finding should be VERIFIED high-severity: %+v", found[0])
	}
}

func TestSSRFSweep_SilencePublishesNothing(t *testing.T) {
	campaignID := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, campaignID, "http://t/fetch?url=http://seed")

	a := newSSRFAgent(campaignID)
	a.SetSSRFVerifier(func(context.Context, string, *verify.ProbeSpec) (bool, string) { return false, "" })
	a.ssrfSweep(ctx, board, nil)

	if n := len(ssrfFindings(t, ctx, board)); n != 0 {
		t.Errorf("no callback → no finding, got %d", n)
	}
}

func TestSSRFSweep_NilVerifierNoop(t *testing.T) {
	campaignID := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, campaignID, "http://t/fetch?url=http://seed")

	a := newSSRFAgent(campaignID) // no verifier
	a.ssrfSweep(ctx, board, nil)
	if n := len(ssrfFindings(t, ctx, board)); n != 0 {
		t.Errorf("without --verify the sweep is a no-op, got %d findings", n)
	}
}

func TestSSRFSweep_DedupesAcrossCalls(t *testing.T) {
	campaignID := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, campaignID, "http://t/fetch?url=http://seed")

	a := newSSRFAgent(campaignID)
	calls := 0
	a.SetSSRFVerifier(func(context.Context, string, *verify.ProbeSpec) (bool, string) {
		calls++
		return true, "callback"
	})
	a.ssrfSweep(ctx, board, nil)
	a.ssrfSweep(ctx, board, nil) // same sink — must not re-probe
	if calls != 1 {
		t.Errorf("sink should be probed once across calls, got %d", calls)
	}
}