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 TestLDAPSweep_ProbesLDAPEndpointsOnly(t *testing.T) {
	c := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, c, "http://t/api/users/ldap-search")
	writeEndpoint(ctx, board, c, "http://t/api/products") // not ldap

	a := &ExploitAgent{campaignID: c, seenLDAP: make(map[string]struct{})}
	var probed []string
	a.SetLDAPVerifier(func(_ context.Context, target string, _ *verify.ProbeSpec) (bool, string) {
		probed = append(probed, target)
		return true, "attack diverges from control"
	})
	a.ldapSweep(ctx, board, nil)

	for _, p := range probed {
		if p != "http://t/api/users/ldap-search" {
			t.Fatalf("only the ldap endpoint should be 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 == "ldap injection" && cf.Verified {
			n++
		}
	}
	if n != 1 {
		t.Fatalf("expected 1 VERIFIED ldap finding, got %d", n)
	}
}