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 newNoSQLAgent(c uuid.UUID) *ExploitAgent {
	return &ExploitAgent{campaignID: c, seenNoSQL: make(map[string]struct{})}
}

func nosqlFindings(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 == "nosql_injection" {
			out = append(out, cf)
		}
	}
	return out
}

func TestNoSQLSweep_ProbesLoginEndpointsOnly(t *testing.T) {
	c := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, c, "http://t/api/user/login")
	writeEndpoint(ctx, board, c, "http://t/api/products") // not a login endpoint

	a := newNoSQLAgent(c)
	var probed []string
	a.SetNoSQLVerifier(func(_ context.Context, target string, spec *verify.ProbeSpec) (bool, string) {
		probed = append(probed, target)
		return true, "attack 200 vs control 401"
	})
	a.nosqlSweep(ctx, board, nil)

	for _, p := range probed {
		if p != "http://t/api/user/login" {
			t.Fatalf("only the login endpoint should be probed, got %v", probed)
		}
	}
	found := nosqlFindings(t, ctx, board)
	if len(found) != 1 || !found[0].Verified || found[0].Severity != pipeline.SeverityCritical {
		t.Fatalf("expected 1 VERIFIED critical nosql finding, got %+v", found)
	}
}

func TestNoSQLSweep_NoBypassPublishesNothing(t *testing.T) {
	c := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, c, "http://t/login")
	a := newNoSQLAgent(c)
	a.SetNoSQLVerifier(func(context.Context, string, *verify.ProbeSpec) (bool, string) { return false, "" })
	a.nosqlSweep(ctx, board, nil)
	if n := len(nosqlFindings(t, ctx, board)); n != 0 {
		t.Errorf("no bypass -> no finding, got %d", n)
	}
}

func TestNoSQLSweep_NilVerifierNoop(t *testing.T) {
	c := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, c, "http://t/login")
	newNoSQLAgent(c).nosqlSweep(ctx, board, nil)
	if n := len(nosqlFindings(t, ctx, board)); n != 0 {
		t.Errorf("no verifier -> no finding, got %d", n)
	}
}