package agents

import (
	"context"
	"encoding/json"
	"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"
)

func TestOpenRedirectSweep_ProbesSinkParams(t *testing.T) {
	c := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, c, "http://t/go?next=/home") // next is a redirect sink
	writeEndpoint(ctx, board, c, "http://t/page?id=1")     // not a sink

	a := &ExploitAgent{campaignID: c, seenOpenRedir: make(map[string]struct{})}
	var probed []string
	a.SetOpenRedirectVerifier(func(_ context.Context, target string, spec *verify.ProbeSpec) (bool, string) {
		probed = append(probed, target)
		return true, "302 Location -> attacker host"
	})
	a.openRedirectSweep(ctx, board, nil)

	if len(probed) != 1 {
		t.Fatalf("expected exactly the sink endpoint probed, got %v", probed)
	}
	fs, _ := board.Query(ctx, blackboard.Predicate{Types: []blackboard.FindingType{blackboard.TypeMisconfig}})
	n := 0
	for _, f := range fs {
		var cf pipeline.ClassifiedFinding
		if json.Unmarshal(f.Data, &cf) == nil && cf.AttackCategory == "open redirect" && cf.Verified {
			n++
		}
	}
	if n != 1 {
		t.Fatalf("expected 1 VERIFIED open-redirect finding, got %d", n)
	}
}

func TestOpenRedirectSweep_NilVerifierNoop(t *testing.T) {
	c := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, c, "http://t/go?url=x")
	(&ExploitAgent{campaignID: c, seenOpenRedir: make(map[string]struct{})}).openRedirectSweep(ctx, board, nil)
	fs, _ := board.Query(ctx, blackboard.Predicate{Types: []blackboard.FindingType{blackboard.TypeMisconfig}})
	if len(fs) != 0 {
		t.Errorf("no verifier -> no finding, got %d", len(fs))
	}
}