package exploit

import (
	"strings"
	"testing"

	"github.com/Armur-Ai/Pentest-Swarm-AI/internal/pipeline"
	"github.com/google/uuid"
)

func TestBuildKillChains_ComposesEnablerToTarget(t *testing.T) {
	leak := pipeline.ClassifiedFinding{ID: uuid.New(), Title: "Verbose error exposure", AttackCategory: "exposure", Severity: pipeline.SeverityLow, CVSSScore: 3}
	cred := pipeline.ClassifiedFinding{ID: uuid.New(), Title: "Leaked API credential", AttackCategory: "credential", Severity: pipeline.SeverityHigh, CVSSScore: 8}
	takeover := pipeline.ClassifiedFinding{ID: uuid.New(), Title: "Admin privilege escalation", AttackCategory: "privilege", Severity: pipeline.SeverityCritical, CVSSScore: 9.5}

	pb := NewPathBuilder()
	chains := pb.BuildKillChains([]pipeline.ClassifiedFinding{leak, cred, takeover})
	if len(chains) == 0 {
		t.Fatal("expected composed kill-chains")
	}
	// exposure→credential and credential→privilege should both exist.
	var haveLeakToCred, haveCredToPriv bool
	for _, c := range chains {
		if strings.Contains(c.Name, "Verbose error") && strings.Contains(c.Name, "credential") {
			haveLeakToCred = true
		}
		if strings.Contains(c.Name, "Leaked API credential") && strings.Contains(c.Name, "privilege") {
			haveCredToPriv = true
			if c.ExpectedImpact != "critical" {
				t.Errorf("cred→privilege impact should be critical (higher of two), got %q", c.ExpectedImpact)
			}
			if len(c.TargetFindingIDs) != 2 {
				t.Errorf("chain should tag both finding ids")
			}
		}
	}
	if !haveLeakToCred || !haveCredToPriv {
		t.Errorf("missing expected chains: leak→cred=%v cred→priv=%v", haveLeakToCred, haveCredToPriv)
	}
	// strongest-impact first (deterministic order)
	if severityRank(chains[0].ExpectedImpact) < severityRank(chains[len(chains)-1].ExpectedImpact) {
		t.Error("chains should be sorted strongest-impact first")
	}
}

func TestBuildKillChains_NoBridgeNoChain(t *testing.T) {
	a := pipeline.ClassifiedFinding{ID: uuid.New(), Title: "Missing security header", AttackCategory: "header"}
	b := pipeline.ClassifiedFinding{ID: uuid.New(), Title: "Clickjacking", AttackCategory: "ui"}
	if got := NewPathBuilder().BuildKillChains([]pipeline.ClassifiedFinding{a, b}); len(got) != 0 {
		t.Fatalf("unrelated findings should not chain, got %d", len(got))
	}
}