package agents

import (
	"context"
	"encoding/json"
	"strings"
	"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 newCmdInjAgent(campaignID uuid.UUID) *ExploitAgent {
	return &ExploitAgent{campaignID: campaignID, seenCmdInj: make(map[string]struct{})}
}

func cmdInjFindings(t *testing.T, ctx context.Context, board blackboard.Board) []pipeline.ClassifiedFinding {
	t.Helper()
	fs, err := board.Query(ctx, blackboard.Predicate{Types: []blackboard.FindingType{blackboard.TypeMisconfig}})
	if err != nil {
		t.Fatal(err)
	}
	var out []pipeline.ClassifiedFinding
	for _, f := range fs {
		var cf pipeline.ClassifiedFinding
		if json.Unmarshal(f.Data, &cf) == nil && cf.AttackCategory == "command_injection" {
			out = append(out, cf)
		}
	}
	return out
}

func TestCmdInjectionSweep_PublishesCallingBackParam(t *testing.T) {
	campaignID := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, campaignID, "http://t/ping?host=seed&count=2")

	a := newCmdInjAgent(campaignID)
	a.SetCmdInjectionVerifier(func(_ context.Context, _ string, spec *verify.ProbeSpec) (bool, string) {
		// Only the host param is a shell sink.
		if strings.Contains(spec.BaseRequest, "host="+verify.DefaultMarker) {
			return true, "dns interaction from 127.0.0.1"
		}
		return false, ""
	})

	a.cmdInjectionSweep(ctx, board, nil)

	found := cmdInjFindings(t, ctx, board)
	if len(found) != 1 {
		t.Fatalf("expected 1 verified command-injection finding, got %d", len(found))
	}
	if !found[0].Verified || found[0].Severity != pipeline.SeverityCritical {
		t.Errorf("finding should be VERIFIED critical: %+v", found[0])
	}
}

func TestCmdInjectionSweep_SilencePublishesNothing(t *testing.T) {
	campaignID := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, campaignID, "http://t/ping?host=seed")

	a := newCmdInjAgent(campaignID)
	a.SetCmdInjectionVerifier(func(context.Context, string, *verify.ProbeSpec) (bool, string) { return false, "" })
	a.cmdInjectionSweep(ctx, board, nil)
	if n := len(cmdInjFindings(t, ctx, board)); n != 0 {
		t.Errorf("no callback → no finding, got %d", n)
	}
}

func TestCmdInjectionSweep_NilVerifierNoop(t *testing.T) {
	campaignID := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, campaignID, "http://t/ping?host=seed")

	a := newCmdInjAgent(campaignID)
	a.cmdInjectionSweep(ctx, board, nil)
	if n := len(cmdInjFindings(t, ctx, board)); n != 0 {
		t.Errorf("without --verify the sweep is a no-op, got %d findings", n)
	}
}

func TestCmdInjectionSweep_DedupesAcrossCalls(t *testing.T) {
	campaignID := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, campaignID, "http://t/ping?host=seed")

	a := newCmdInjAgent(campaignID)
	calls := 0
	a.SetCmdInjectionVerifier(func(context.Context, string, *verify.ProbeSpec) (bool, string) {
		calls++
		return false, ""
	})
	a.cmdInjectionSweep(ctx, board, nil)
	first := calls
	a.cmdInjectionSweep(ctx, board, nil)
	if calls != first {
		t.Errorf("sink re-probed after dedupe: %d then %d", first, calls)
	}
}