package engine

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/google/uuid"
)

func writeCF(ctx context.Context, t *testing.T, board blackboard.Board, campaignID uuid.UUID, cf pipeline.ClassifiedFinding) {
	t.Helper()
	data, _ := json.Marshal(cf)
	if _, err := board.Write(ctx, blackboard.Finding{
		CampaignID: campaignID, AgentName: "exploit", Type: blackboard.TypeMisconfig,
		Target: cf.Target, Data: data, PheromoneBase: 1.0, HalfLifeSec: 3600,
	}); err != nil {
		t.Fatal(err)
	}
}

func TestCountVerifiedFindings(t *testing.T) {
	ctx := context.Background()
	campaignID := uuid.New()
	board := blackboard.NewMemoryBoard(time.Now)

	id1, id2, id3 := uuid.New(), uuid.New(), uuid.New()
	writeCF(ctx, t, board, campaignID, pipeline.ClassifiedFinding{ID: id1, Target: "http://t/a", AttackCategory: "ssrf", Verified: true})
	writeCF(ctx, t, board, campaignID, pipeline.ClassifiedFinding{ID: id2, Target: "http://t/b", AttackCategory: "ssti", Verified: true})
	writeCF(ctx, t, board, campaignID, pipeline.ClassifiedFinding{ID: id3, Target: "http://t/c", AttackCategory: "ssrf", Verified: false}) // not verified

	n, summary := countVerifiedFindings(ctx, board)
	if n != 2 {
		t.Fatalf("verified count = %d, want 2", n)
	}
	// Deterministic, sorted "class=count".
	if summary != "ssrf=1, ssti=1" {
		t.Errorf("summary = %q, want %q", summary, "ssrf=1, ssti=1")
	}
}

func TestCountVerifiedFindings_DedupesByID(t *testing.T) {
	ctx := context.Background()
	campaignID := uuid.New()
	board := blackboard.NewMemoryBoard(time.Now)

	id := uuid.New()
	// Same finding written twice (e.g. re-published) must count once.
	cf := pipeline.ClassifiedFinding{ID: id, Target: "http://t/a", AttackCategory: "command_injection", Verified: true}
	writeCF(ctx, t, board, campaignID, cf)
	writeCF(ctx, t, board, campaignID, cf)

	n, summary := countVerifiedFindings(ctx, board)
	if n != 1 || summary != "command_injection=1" {
		t.Errorf("n=%d summary=%q, want 1 / command_injection=1", n, summary)
	}
}

func TestCountVerifiedFindings_NoneVerified(t *testing.T) {
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	n, summary := countVerifiedFindings(ctx, board)
	if n != 0 || summary != "" {
		t.Errorf("empty board: n=%d summary=%q, want 0 / empty", n, summary)
	}
}