package agents

import (
	"context"
	"encoding/json"
	"fmt"

	exploitpkg "github.com/Armur-Ai/Pentest-Swarm-AI/internal/agent/exploit"
	"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/tuning"
	"github.com/google/uuid"
)

// ExploitAgent wakes on high-pheromone CVE_MATCH findings, builds an attack
// chain via the existing exploit agent, and publishes EXPLOIT_CHAIN to the
// board. The scheduler will never hand this agent a stale finding — the
// MinPheromone predicate takes care of that at the DB/memory layer.
type ExploitAgent struct {
	exploit    *exploitpkg.ExploitAgent
	executor   *exploitpkg.Executor
	objective  string
	campaignID uuid.UUID
	parallel   int
	dryRun     bool
	tun        *tuning.Settings
}

// NewExploitAgent wires the existing exploit agent + executor into the swarm.
// Pass nil for tun to use the baked-in default pheromone tuning.
func NewExploitAgent(inner *exploitpkg.ExploitAgent, exec *exploitpkg.Executor, objective string, campaignID uuid.UUID, parallel int, dryRun bool, tun *tuning.Settings) *ExploitAgent {
	if parallel <= 0 {
		parallel = 1
	}
	if tun == nil {
		tun = tuning.Default()
	}
	return &ExploitAgent{
		exploit:    inner,
		executor:   exec,
		objective:  objective,
		campaignID: campaignID,
		parallel:   parallel,
		dryRun:     dryRun,
		tun:        tun,
	}
}

// Name implements swarm.Agent.
func (a *ExploitAgent) Name() string { return "exploit" }

// Trigger implements swarm.Agent.
func (a *ExploitAgent) Trigger() blackboard.Predicate {
	return blackboard.Predicate{
		Types:        []blackboard.FindingType{blackboard.TypeCVEMatch},
		MinPheromone: 0.5, // only act on credible, recent findings
	}
}

// MaxConcurrency implements swarm.Agent.
func (a *ExploitAgent) MaxConcurrency() int { return a.parallel }

// Handle builds an attack chain for a single CVE_MATCH and (optionally)
// executes the top path.
func (a *ExploitAgent) Handle(ctx context.Context, f blackboard.Finding, board blackboard.Board) error {
	var cf pipeline.ClassifiedFinding
	if err := json.Unmarshal(f.Data, &cf); err != nil {
		return fmt.Errorf("decode classified finding: %w", err)
	}

	set := pipeline.ClassifiedFindingSet{
		CampaignID: a.campaignID,
		Findings:   []pipeline.ClassifiedFinding{cf},
	}
	plan, err := a.exploit.BuildPlan(ctx, set, a.objective)
	if err != nil {
		return fmt.Errorf("build plan: %w", err)
	}
	if plan == nil || len(plan.Paths) == 0 {
		return nil
	}

	// Always publish the chain so other agents (and the final report) can see it.
	chainData, _ := json.Marshal(plan.Paths[0])
	chainBase, chainHalf := a.tun.Lookup(blackboard.TypeExploitChain)
	chainID, _ := board.Write(ctx, blackboard.Finding{
		CampaignID:    a.campaignID,
		AgentName:     a.Name(),
		Type:          blackboard.TypeExploitChain,
		Target:        cf.Target,
		Data:          chainData,
		PheromoneBase: chainBase,
		HalfLifeSec:   chainHalf,
	})

	if a.dryRun || a.executor == nil {
		return nil
	}

	// Execute steps of the top path. Each result becomes its own finding.
	for _, step := range plan.Paths[0].Steps {
		if step.Command == "" {
			continue
		}
		res, execErr := a.executor.Execute(ctx, step, a.campaignID)
		if execErr != nil {
			continue
		}
		resData, _ := json.Marshal(map[string]any{
			"chain_id": chainID, "step": step.Name, "result": res,
		})
		pheromone, halfLife := a.tun.Lookup(blackboard.TypeExploitResult)
		if res.Success {
			// Successful exploits are maximally interesting; override the
			// tuned base so downstream agents definitely notice.
			pheromone = 1.0
		}
		_, _ = board.Write(ctx, blackboard.Finding{
			CampaignID:    a.campaignID,
			AgentName:     a.Name(),
			Type:          blackboard.TypeExploitResult,
			Target:        cf.Target,
			Data:          resData,
			PheromoneBase: pheromone,
			HalfLifeSec:   halfLife,
		})
	}
	return nil
}
