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

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

func TestSSTISweep_PublishesEvaluatingParam(t *testing.T) {
	campaignID := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, campaignID, "http://t/greet?name=bob&lang=en")

	a := newSSTIAgent(campaignID)
	a.SetSSTIVerifier(func(_ context.Context, _ string, spec *verify.ProbeSpec) (bool, string) {
		// Only the `name` parameter templates; `lang` is inert.
		if !strings.Contains(spec.BaseRequest, "name="+verify.DefaultMarker) {
			return false, ""
		}
		// Polyglot gate proves; the Jinja per-engine spec also proves.
		if strings.Contains(spec.AttackPayload, "%3C%25") { // polyglot carries encoded "<%"
			return true, "polyglot"
		}
		if strings.HasPrefix(spec.AttackPayload, "%7B%7B") { // Jinja "{{"
			return true, "indicator 36863 present in attack"
		}
		return false, ""
	})

	a.sstiSweep(ctx, board, nil)

	found := sstiFindings(t, ctx, board)
	if len(found) != 1 {
		t.Fatalf("expected 1 verified SSTI finding, got %d", len(found))
	}
	if !found[0].Verified || found[0].Severity != pipeline.SeverityCritical {
		t.Errorf("SSTI finding should be VERIFIED critical: %+v", found[0])
	}
	if !strings.Contains(found[0].Title, "Jinja") {
		t.Errorf("title should name the engine: %q", found[0].Title)
	}
}

func TestSSTISweep_PlainReflectionPublishesNothing(t *testing.T) {
	campaignID := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, campaignID, "http://t/greet?name=bob")

	a := newSSTIAgent(campaignID)
	a.SetSSTIVerifier(func(context.Context, string, *verify.ProbeSpec) (bool, string) { return false, "" })
	a.sstiSweep(ctx, board, nil)
	if n := len(sstiFindings(t, ctx, board)); n != 0 {
		t.Errorf("a non-templating reflection must publish nothing, got %d", n)
	}
}

func TestSSTISweep_NilVerifierNoop(t *testing.T) {
	campaignID := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeEndpoint(ctx, board, campaignID, "http://t/greet?name=bob")

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

func TestQueryParamNames(t *testing.T) {
	got := queryParamNames("http://t/p?a=1&b=2&a=3&c")
	want := []string{"a", "b", "c"}
	if strings.Join(got, ",") != strings.Join(want, ",") {
		t.Errorf("queryParamNames = %v, want %v", got, want)
	}
	if queryParamNames("http://t/nopath") != nil {
		t.Error("no query → nil")
	}
}