package exploit

import (
	"context"
	"fmt"
	"sort"
	"strings"

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

// ScoredPath pairs a candidate attack path with its adaptive score and the
// source of that score ("jev" when Jev decided, "heuristic" on fallback).
type ScoredPath struct {
	Path   pipeline.AttackPath
	Score  float64
	Source string
}

// scoreboardMaxRows caps how many ranked strategies we render in the live
// scoreboard so a large candidate set doesn't flood the event stream.
const scoreboardMaxRows = 6

// FormatScoreboard renders the ranked candidate strategies as human-readable
// lines for the live TUI / dashboard stream: a header plus one row per strategy
// (rank, score, source, name), the top row marked as the one being pursued.
// Returns nil for an empty set.
func FormatScoreboard(scored []ScoredPath) []string {
	if len(scored) == 0 {
		return nil
	}
	lines := []string{fmt.Sprintf("adaptive scoring — %d candidate strategies ranked:", len(scored))}
	for i, sp := range scored {
		if i >= scoreboardMaxRows {
			lines = append(lines, fmt.Sprintf("  …and %d more", len(scored)-scoreboardMaxRows))
			break
		}
		marker := " "
		if i == 0 {
			marker = "▶" // pursued first
		}
		name := sp.Path.Name
		if name == "" {
			name = "attack path"
		}
		lines = append(lines, fmt.Sprintf("  %s #%d  %.2f  [%s]  %s", marker, i+1, sp.Score, sp.Source, name))
	}
	return lines
}

// strategyScorer is the minimal Jev surface the scorer needs. Declared as an
// interface so tests can inject a fake without a live API.
type strategyScorer interface {
	ScoreStrategies(ctx context.Context, state string, strategies map[string]string) (map[string]float64, error)
}

// JevScorer ranks candidate attack paths using Jev's real-time strategy scoring
// — the adaptive brain that decides which strategies the swarm should pursue
// most aggressively. It fails OPEN: if Jev is unavailable, disabled, or errors,
// each path keeps its heuristic EstimatedSuccessProbability so the swarm keeps
// moving; it just loses the adaptive edge rather than stalling.
type JevScorer struct {
	jev strategyScorer
}

// NewJevScorer builds a scorer over a Jev client. A nil client yields a scorer
// that always falls back to the heuristic ranking (safe no-op).
func NewJevScorer(client *jev.Client) *JevScorer {
	if client == nil {
		return &JevScorer{}
	}
	return &JevScorer{jev: client}
}

// describePath renders an attack path as a compact strategy description for Jev
// to reason over: its name, expected impact, and the technique of each step.
func describePath(p pipeline.AttackPath) string {
	var b strings.Builder
	if p.Name != "" {
		b.WriteString(p.Name)
	} else {
		b.WriteString("attack path")
	}
	if p.ExpectedImpact != "" {
		fmt.Fprintf(&b, " — impact: %s", p.ExpectedImpact)
	}
	b.WriteString("\nSteps:")
	for i, s := range p.Steps {
		tech := s.TechniqueID
		if tech == "" {
			tech = "action"
		}
		fmt.Fprintf(&b, "\n  %d. %s [%s]", i+1, s.Name, tech)
	}
	return b.String()
}

// RankPaths scores every candidate path against the live campaign state and
// returns them sorted best-first (highest score wins). state summarises what the
// swarm knows right now — target, discoveries, and which steps have already
// succeeded or failed — so Jev scores strategies against reality, not a static
// prior. RankPaths never errors: on any Jev failure it falls back to each path's
// heuristic probability. Safe to call on a nil *JevScorer.
func (s *JevScorer) RankPaths(ctx context.Context, state string, paths []pipeline.AttackPath) []ScoredPath {
	ranked := make([]ScoredPath, 0, len(paths))
	if len(paths) == 0 {
		return ranked
	}

	var scores map[string]float64
	if s != nil && s.jev != nil {
		strategies := make(map[string]string, len(paths))
		for _, p := range paths {
			strategies[p.ID.String()] = describePath(p)
		}
		if got, err := s.jev.ScoreStrategies(ctx, state, strategies); err == nil {
			scores = got
		}
	}

	for _, p := range paths {
		sp := ScoredPath{Path: p, Score: p.EstimatedSuccessProbability, Source: "heuristic"}
		if scores != nil {
			if v, ok := scores[p.ID.String()]; ok {
				sp.Score = v
				sp.Source = "jev"
			}
		}
		ranked = append(ranked, sp)
	}
	// Stable so equal scores keep planner order; highest score first.
	sort.SliceStable(ranked, func(i, j int) bool { return ranked[i].Score > ranked[j].Score })
	return ranked
}