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

// writeProtectedEndpoint stores an endpoint recon saw return a given status.
func writeProtectedEndpoint(ctx context.Context, board blackboard.Board, campaignID uuid.UUID, urlStr string, status int) {
	data, _ := json.Marshal(pipeline.EndpointRecord{URL: urlStr, Method: "GET", StatusCode: status})
	_, _ = board.Write(ctx, blackboard.Finding{
		CampaignID: campaignID, AgentName: "recon", Type: blackboard.TypeHTTPEndpoint,
		Target: urlStr, Data: data, PheromoneBase: 0.6, HalfLifeSec: 3600,
	})
}

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

func TestAuthBypassSweep_PublishesProvenBypass(t *testing.T) {
	campaignID := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeProtectedEndpoint(ctx, board, campaignID, "http://t/admin", 403)
	writeProtectedEndpoint(ctx, board, campaignID, "http://t/public", 200) // not protected → skipped

	a := newAuthBypassAgent(campaignID)
	var probed []string
	a.SetAuthBypassVerifier(func(_ context.Context, target string, spec *verify.ProbeSpec) (bool, string) {
		probed = append(probed, target)
		// Prove the trailing-slash variant only.
		return spec.AttackPayload == "/admin/", "attack 200 vs control 403 — authorization boundary crossed"
	})

	a.authBypassSweep(ctx, board, nil)

	for _, p := range probed {
		if strings.Contains(p, "/public") {
			t.Fatalf("a 200 endpoint must not be bypass-probed: %v", probed)
		}
	}
	found := authBypassFindings(t, ctx, board)
	if len(found) != 1 {
		t.Fatalf("expected 1 verified bypass finding, got %d", len(found))
	}
	if !found[0].Verified || !strings.Contains(found[0].Title, "trailing-slash") {
		t.Errorf("finding should be VERIFIED and name the technique: %+v", found[0])
	}
}

func TestAuthBypassSweep_NoBypassPublishesNothing(t *testing.T) {
	campaignID := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeProtectedEndpoint(ctx, board, campaignID, "http://t/admin", 403)

	a := newAuthBypassAgent(campaignID)
	a.SetAuthBypassVerifier(func(context.Context, string, *verify.ProbeSpec) (bool, string) { return false, "" })
	a.authBypassSweep(ctx, board, nil)
	if n := len(authBypassFindings(t, ctx, board)); n != 0 {
		t.Errorf("no working variant → no finding, got %d", n)
	}
}

func TestAuthBypassSweep_NilVerifierNoop(t *testing.T) {
	campaignID := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeProtectedEndpoint(ctx, board, campaignID, "http://t/admin", 403)

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

func TestAuthBypassSweep_DedupesAcrossCalls(t *testing.T) {
	campaignID := uuid.New()
	ctx := context.Background()
	board := blackboard.NewMemoryBoard(time.Now)
	writeProtectedEndpoint(ctx, board, campaignID, "http://t/admin", 403)

	a := newAuthBypassAgent(campaignID)
	endpoints := map[string]int{}
	a.SetAuthBypassVerifier(func(_ context.Context, target string, _ *verify.ProbeSpec) (bool, string) {
		endpoints[target]++
		return false, ""
	})
	a.authBypassSweep(ctx, board, nil)
	a.authBypassSweep(ctx, board, nil) // same endpoint — must not re-probe
	if endpoints["http://t/admin"] == 0 {
		t.Fatal("endpoint was never probed")
	}
	// Second sweep must add no further probes for the same endpoint.
	first := endpoints["http://t/admin"]
	a.authBypassSweep(ctx, board, nil)
	if endpoints["http://t/admin"] != first {
		t.Errorf("endpoint re-probed after dedupe: %d then %d", first, endpoints["http://t/admin"])
	}
}