package exploit

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

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

// ExploitAgent constructs multi-step attack chains using LLM reasoning.
type ExploitAgent struct {
	provider    llm.Provider
	pathBuilder *PathBuilder
}

// NewExploitAgent creates a new exploit agent.
func NewExploitAgent(provider llm.Provider) *ExploitAgent {
	return &ExploitAgent{
		provider:    provider,
		pathBuilder: NewPathBuilder(),
	}
}

// BuildPlan constructs an attack plan from classified findings.
func (e *ExploitAgent) BuildPlan(ctx context.Context, findings pipeline.ClassifiedFindingSet, objective string) (*pipeline.AttackPlan, error) {
	// Step 1: Build rule-based chains
	ruleChains := e.pathBuilder.BuildChains(findings.Findings)

	// Step 2: Send to LLM for reasoning and enhancement
	findingsJSON, _ := json.Marshal(findings.Findings)
	chainsJSON, _ := json.Marshal(ruleChains)

	req := llm.CompletionRequest{
		SystemPrompt: exploitSystemPrompt,
		Messages: []llm.Message{
			{
				Role: "user",
				Content: fmt.Sprintf(
					"Objective: %s\n\nClassified Findings:\n%s\n\nRule-Based Attack Chains:\n%s\n\nAnalyze these findings. Think through each one in <think> tags, then produce your attack plan as JSON.",
					objective, string(findingsJSON), string(chainsJSON),
				),
			},
		},
		MaxTokens:   8192,
		Temperature: 0.2,
	}

	resp, err := e.provider.Complete(ctx, req)
	if err != nil {
		// Fall back to rule-based chains only
		return &pipeline.AttackPlan{
			ID:        uuid.New(),
			Paths:     ruleChains,
			Reasoning: "LLM unavailable — using rule-based chains only",
			CreatedAt: time.Now(),
		}, nil
	}

	// Parse <think> reasoning and JSON plan
	reasoning, planJSON := parseThinkAndJSON(resp.Content)

	var llmPaths []pipeline.AttackPath
	if planJSON != "" {
		_ = json.Unmarshal([]byte(planJSON), &llmPaths)
	}

	// Merge rule-based and LLM-generated paths
	allPaths := mergePaths(ruleChains, llmPaths)

	plan := &pipeline.AttackPlan{
		ID:        uuid.New(),
		Paths:     allPaths,
		Reasoning: reasoning,
		CreatedAt: time.Now(),
	}

	if len(allPaths) > 0 {
		plan.RecommendedPathID = allPaths[0].ID
	}

	return plan, nil
}

// AdaptPlan adjusts the attack plan based on execution results.
func (e *ExploitAgent) AdaptPlan(ctx context.Context, plan *pipeline.AttackPlan, result pipeline.ExecutionResult) (*pipeline.AttackPlan, error) {
	planJSON, _ := json.Marshal(plan)
	resultJSON, _ := json.Marshal(result)

	req := llm.CompletionRequest{
		SystemPrompt: "You are a penetration testing strategist. Given the current attack plan and the latest execution result, decide: continue on the current path, pivot to an alternative, or adjust the approach. Think in <think> tags, then output the updated plan as JSON.",
		Messages: []llm.Message{
			{
				Role:    "user",
				Content: fmt.Sprintf("Current Plan:\n%s\n\nLatest Result:\n%s", string(planJSON), string(resultJSON)),
			},
		},
		MaxTokens:   4096,
		Temperature: 0.2,
	}

	resp, err := e.provider.Complete(ctx, req)
	if err != nil {
		return plan, nil // keep existing plan on failure
	}

	reasoning, planJSONStr := parseThinkAndJSON(resp.Content)

	var updatedPaths []pipeline.AttackPath
	if planJSONStr != "" {
		_ = json.Unmarshal([]byte(planJSONStr), &updatedPaths)
	}

	if len(updatedPaths) > 0 {
		plan.Paths = updatedPaths
		plan.Reasoning = reasoning
	}

	return plan, nil
}

// parseThinkAndJSON splits LLM output into <think> reasoning and JSON content.
func parseThinkAndJSON(content string) (reasoning, jsonContent string) {
	// Extract <think> block
	thinkStart := strings.Index(content, "<think>")
	thinkEnd := strings.Index(content, "</think>")

	if thinkStart >= 0 && thinkEnd > thinkStart {
		reasoning = strings.TrimSpace(content[thinkStart+7 : thinkEnd])
		content = content[thinkEnd+8:]
	}

	// Extract JSON (find first [ or {)
	content = strings.TrimSpace(content)
	for i, c := range content {
		if c == '[' || c == '{' {
			jsonContent = content[i:]
			break
		}
	}

	// Strip trailing markdown
	if strings.HasSuffix(jsonContent, "```") {
		jsonContent = jsonContent[:len(jsonContent)-3]
	}
	jsonContent = strings.TrimSpace(jsonContent)

	return reasoning, jsonContent
}

// mergePaths combines rule-based and LLM-generated paths, deduplicating by name.
func mergePaths(ruleBased, llmGenerated []pipeline.AttackPath) []pipeline.AttackPath {
	seen := make(map[string]bool)
	var merged []pipeline.AttackPath

	// LLM paths take priority
	for _, p := range llmGenerated {
		seen[p.Name] = true
		merged = append(merged, p)
	}

	for _, p := range ruleBased {
		if !seen[p.Name] {
			merged = append(merged, p)
		}
	}

	return merged
}

const exploitSystemPrompt = `You are a specialized penetration testing strategist. Your role is to construct multi-step attack chains from classified security findings.

For each finding, think through:
1. How exploitable is this? What tools/techniques are needed?
2. Can this be chained with other findings for greater impact?
3. What is the most efficient path to the stated objective?

Output your thinking in <think> tags, then produce a JSON array of attack paths.

Each path should have:
- name: concise chain description (e.g., "SQLi → Data Extraction → Credential Access")
- description: detailed explanation
- steps: ordered array of {name, technique_id (MITRE ATT&CK), command, expected_output_pattern, cleanup_command}
- estimated_success_probability: 0.0-1.0
- expected_impact: low/medium/high/critical

Prioritize chains that lead to the objective with highest probability of success.`
