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 xxeFindings(t *testing.T, ctx context.Context, b blackboard.Board) []pipeline.ClassifiedFinding {
	t.Helper()
	fs, _ := b.Query(ctx, blackboard.Predicate{Types: []blackboard.FindingType{blackboard.TypeMisconfig}})
	var out []pipeline.ClassifiedFinding
	for _, f := range fs {
		var cf pipeline.ClassifiedFinding
		if json.Unmarshal(f.Data, &cf) == nil && cf.AttackCategory == "xxe" {
			out = append(out, cf)
		}
	}
	return out
}

func TestXXESweep_ProbesXMLEndpointsOnly(t *testing.T) {
	c := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, c, "http://t/api/notes/import/xml")
	writeEndpoint(ctx, board, c, "http://t/api/products") // not XML

	a := &ExploitAgent{campaignID: c, seenXXE: make(map[string]struct{})}
	var probed []string
	a.SetXXEVerifier(func(_ context.Context, target string, spec *verify.ProbeSpec) (bool, string) {
		probed = append(probed, target)
		return spec.Class == verify.ProbeOOB, "http interaction from 10.0.0.9"
	})
	a.xxeSweep(ctx, board, nil)

	for _, p := range probed {
		if p != "http://t/api/notes/import/xml" {
			t.Fatalf("only the XML endpoint should be probed, got %v", probed)
		}
	}
	found := xxeFindings(t, ctx, board)
	if len(found) != 1 || !found[0].Verified {
		t.Fatalf("expected 1 VERIFIED xxe finding, got %+v", found)
	}
}

func TestXXESweep_NilVerifierNoop(t *testing.T) {
	c := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, c, "http://t/import/xml")
	(&ExploitAgent{campaignID: c, seenXXE: make(map[string]struct{})}).xxeSweep(ctx, board, nil)
	if n := len(xxeFindings(t, ctx, board)); n != 0 {
		t.Errorf("no verifier -> no finding, got %d", n)
	}
}