//! Deep Mode Fleet Executor
//!
//! Implements the iterative planning and execution loop (agentic pattern).
//! Unlike other modes that execute once and return, deep mode:
//! 1. Plans investigation steps using the LLM
//! 2. Executes steps iteratively
//! 3. Re-plans based on findings
//! 4. Synthesizes final answer with evidence

use aof_core::{AgentFleet, AofError, AofResult, DeepConfig};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, info, warn};

use crate::Runtime;

/// Investigation step generated by the planner
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InvestigationStep {
    /// Step number
    pub step: u32,
    /// Description of what this step does
    pub description: String,
    /// The action/query to execute
    pub action: String,
    /// Optional: specific tool to use
    pub tool: Option<String>,
    /// Whether this step is complete
    pub completed: bool,
    /// Result from executing this step
    pub result: Option<String>,
}

/// Investigation plan generated by the planner
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InvestigationPlan {
    /// Goal of the investigation
    pub goal: String,
    /// Steps to execute
    pub steps: Vec<InvestigationStep>,
    /// Current step index
    pub current_step: usize,
    /// Whether plan needs revision
    pub needs_revision: bool,
}

/// Finding from an investigation step
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Finding {
    /// Step that produced this finding
    pub step: u32,
    /// Description of the finding
    pub description: String,
    /// Raw data/evidence
    pub evidence: String,
    /// Significance (low, medium, high)
    pub significance: String,
}

/// Deep mode execution result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DeepResult {
    /// Final answer/conclusion
    pub conclusion: String,
    /// Evidence supporting the conclusion
    pub evidence: Vec<Finding>,
    /// Recommendations
    pub recommendations: Vec<String>,
    /// Number of iterations taken
    pub iterations: u32,
    /// Whether goal was achieved
    pub goal_achieved: bool,
}

/// Deep Fleet Executor - implements the agentic loop pattern
pub struct DeepFleetExecutor {
    /// Fleet configuration
    fleet: AgentFleet,
    /// Deep mode configuration
    config: DeepConfig,
    /// Runtime for agent execution
    runtime: Arc<RwLock<Runtime>>,
    /// Collected findings across iterations
    findings: Vec<Finding>,
    /// Memory of context across iterations
    memory: Vec<String>,
}

impl DeepFleetExecutor {
    /// Create a new deep fleet executor
    pub fn new(fleet: AgentFleet, runtime: Arc<RwLock<Runtime>>) -> Self {
        let config = fleet
            .spec
            .coordination
            .deep
            .clone()
            .unwrap_or_default();

        Self {
            fleet,
            config,
            runtime,
            findings: Vec::new(),
            memory: Vec::new(),
        }
    }

    /// Execute the deep mode investigation loop
    pub async fn execute(&mut self, query: &str) -> AofResult<DeepResult> {
        info!("Starting deep mode investigation: {}", query);

        let mut iteration = 0;
        let max_iterations = self.config.max_iterations;

        // Step 1: Generate initial plan
        let mut plan = if self.config.planning {
            self.generate_plan(query).await?
        } else {
            // Simple plan - just execute the query directly
            InvestigationPlan {
                goal: query.to_string(),
                steps: vec![InvestigationStep {
                    step: 1,
                    description: "Execute query".to_string(),
                    action: query.to_string(),
                    tool: None,
                    completed: false,
                    result: None,
                }],
                current_step: 0,
                needs_revision: false,
            }
        };

        info!("Generated plan with {} steps", plan.steps.len());

        // Step 2: Execute steps iteratively
        while iteration < max_iterations && !self.is_goal_achieved(&plan) {
            iteration += 1;
            debug!("Deep mode iteration {}/{}", iteration, max_iterations);

            // Execute current step
            if plan.current_step < plan.steps.len() {
                let step_result = self.execute_step(&mut plan.steps[plan.current_step]).await?;

                // Store in memory if enabled
                if self.config.memory {
                    self.memory.push(format!(
                        "Step {}: {} -> {}",
                        plan.steps[plan.current_step].step,
                        plan.steps[plan.current_step].description,
                        step_result
                    ));
                }

                // Analyze result and potentially add findings
                self.analyze_step_result(&plan.steps[plan.current_step], &step_result)
                    .await?;

                // Mark step as complete
                plan.steps[plan.current_step].completed = true;
                plan.steps[plan.current_step].result = Some(step_result.clone());
                plan.current_step += 1;
            }

            // Step 3: Re-plan if needed
            if self.should_replan(&plan) {
                info!("Re-planning based on findings...");
                plan = self.adjust_plan(&plan, query).await?;
            }
        }

        // Step 4: Synthesize final result
        let result = self.synthesize(query, &plan).await?;

        info!(
            "Deep mode completed in {} iterations, goal achieved: {}",
            iteration, result.goal_achieved
        );

        Ok(result)
    }

    /// Generate an investigation plan from the query
    async fn generate_plan(&self, query: &str) -> AofResult<InvestigationPlan> {
        // Get the first agent to use for planning (or use planner_model if specified)
        let planner_agent = self.fleet.spec.agents.first().ok_or_else(|| {
            AofError::runtime("No agents available for planning".to_string())
        })?;

        let planner_prompt = self.config.planner_prompt.clone().unwrap_or_else(|| {
            r#"You are a planning agent for root cause analysis investigations.
Given a query, generate a structured investigation plan.

Output your plan as JSON with this structure:
{
  "goal": "the investigation goal",
  "steps": [
    {
      "step": 1,
      "description": "what this step does",
      "action": "the specific query or command to execute",
      "tool": "optional tool name"
    }
  ]
}

Generate 3-7 logical investigation steps that will systematically find the root cause."#
                .to_string()
        });

        let prompt = format!(
            "{}\n\nQuery: {}\n\nGenerate your investigation plan:",
            planner_prompt, query
        );

        let response = {
            let runtime = self.runtime.read().await;
            runtime
                .execute(&planner_agent.name, &prompt)
                .await
                .map_err(|e| AofError::runtime(format!("Planning failed: {}", e)))?
        };

        // Parse the plan from the response
        self.parse_plan(&response, query)
    }

    /// Parse a plan from the LLM response
    fn parse_plan(&self, response: &str, query: &str) -> AofResult<InvestigationPlan> {
        // Try to extract JSON from the response
        let json_start = response.find('{');
        let json_end = response.rfind('}');

        if let (Some(start), Some(end)) = (json_start, json_end) {
            let json_str = &response[start..=end];
            if let Ok(plan) = serde_json::from_str::<InvestigationPlan>(json_str) {
                return Ok(plan);
            }
        }

        // Fallback: create a simple plan from the response
        warn!("Could not parse plan JSON, creating fallback plan");
        Ok(InvestigationPlan {
            goal: query.to_string(),
            steps: vec![
                InvestigationStep {
                    step: 1,
                    description: "Gather initial data".to_string(),
                    action: format!("Investigate: {}", query),
                    tool: None,
                    completed: false,
                    result: None,
                },
                InvestigationStep {
                    step: 2,
                    description: "Analyze findings".to_string(),
                    action: "Analyze the data gathered in previous step".to_string(),
                    tool: None,
                    completed: false,
                    result: None,
                },
                InvestigationStep {
                    step: 3,
                    description: "Identify root cause".to_string(),
                    action: "Based on analysis, identify the root cause".to_string(),
                    tool: None,
                    completed: false,
                    result: None,
                },
            ],
            current_step: 0,
            needs_revision: false,
        })
    }

    /// Execute a single investigation step
    async fn execute_step(&self, step: &mut InvestigationStep) -> AofResult<String> {
        debug!("Executing step {}: {}", step.step, step.description);

        // Get an appropriate agent for this step
        let agent = self.fleet.spec.agents.first().ok_or_else(|| {
            AofError::runtime("No agents available for step execution".to_string())
        })?;

        // Build context from memory
        let context = if self.config.memory && !self.memory.is_empty() {
            format!(
                "Previous findings:\n{}\n\n",
                self.memory.join("\n")
            )
        } else {
            String::new()
        };

        let prompt = format!(
            "{}Step: {}\nDescription: {}\nAction: {}\n\nExecute this step and provide your findings:",
            context, step.step, step.description, step.action
        );

        let response = {
            let runtime = self.runtime.read().await;
            runtime
                .execute(&agent.name, &prompt)
                .await
                .map_err(|e| AofError::runtime(format!("Step execution failed: {}", e)))?
        };

        Ok(response)
    }

    /// Analyze a step result and potentially add findings
    async fn analyze_step_result(
        &mut self,
        step: &InvestigationStep,
        result: &str,
    ) -> AofResult<()> {
        // Extract key findings from the result
        // For now, we'll add the entire result as a finding
        // In a more sophisticated implementation, we'd use the LLM to extract findings

        let significance = if result.to_lowercase().contains("error")
            || result.to_lowercase().contains("critical")
            || result.to_lowercase().contains("failed")
        {
            "high".to_string()
        } else if result.to_lowercase().contains("warning")
            || result.to_lowercase().contains("unusual")
        {
            "medium".to_string()
        } else {
            "low".to_string()
        };

        self.findings.push(Finding {
            step: step.step,
            description: step.description.clone(),
            evidence: result.to_string(),
            significance,
        });

        Ok(())
    }

    /// Check if we should re-plan based on current findings
    fn should_replan(&self, plan: &InvestigationPlan) -> bool {
        // Re-plan if we found high-significance findings
        let high_significance_count = self
            .findings
            .iter()
            .filter(|f| f.significance == "high")
            .count();

        // Re-plan if we have multiple high-significance findings and haven't revised yet
        high_significance_count >= 2 && !plan.needs_revision
    }

    /// Adjust the plan based on findings
    async fn adjust_plan(
        &self,
        current_plan: &InvestigationPlan,
        original_query: &str,
    ) -> AofResult<InvestigationPlan> {
        let agent = self.fleet.spec.agents.first().ok_or_else(|| {
            AofError::runtime("No agents available for re-planning".to_string())
        })?;

        let findings_summary: String = self
            .findings
            .iter()
            .map(|f| format!("- Step {}: {} ({})", f.step, f.description, f.significance))
            .collect::<Vec<_>>()
            .join("\n");

        let prompt = format!(
            r#"Based on new findings, revise the investigation plan.

Original query: {}

Completed steps:
{}

Key findings:
{}

Current remaining steps:
{}

Generate a revised plan (JSON format) that incorporates these findings and focuses on the most promising leads:"#,
            original_query,
            current_plan
                .steps
                .iter()
                .filter(|s| s.completed)
                .map(|s| format!("- {}: {}", s.step, s.description))
                .collect::<Vec<_>>()
                .join("\n"),
            findings_summary,
            current_plan
                .steps
                .iter()
                .filter(|s| !s.completed)
                .map(|s| format!("- {}: {}", s.step, s.description))
                .collect::<Vec<_>>()
                .join("\n")
        );

        let response = {
            let runtime = self.runtime.read().await;
            runtime
                .execute(&agent.name, &prompt)
                .await
                .map_err(|e| AofError::runtime(format!("Re-planning failed: {}", e)))?
        };

        let mut new_plan = self.parse_plan(&response, original_query)?;
        new_plan.needs_revision = true; // Mark as revised so we don't revise again

        // Preserve completed steps
        let completed_steps: Vec<_> = current_plan
            .steps
            .iter()
            .filter(|s| s.completed)
            .cloned()
            .collect();

        // Renumber new steps
        let next_step = completed_steps.len() as u32 + 1;
        for (i, step) in new_plan.steps.iter_mut().enumerate() {
            step.step = next_step + i as u32;
        }

        // Prepend completed steps
        let mut final_steps = completed_steps;
        final_steps.extend(new_plan.steps);
        new_plan.steps = final_steps;
        new_plan.current_step = current_plan.current_step;

        Ok(new_plan)
    }

    /// Check if the investigation goal has been achieved
    fn is_goal_achieved(&self, plan: &InvestigationPlan) -> bool {
        // Goal is achieved when all steps are complete
        plan.steps.iter().all(|s| s.completed)
    }

    /// Synthesize the final result from all findings
    async fn synthesize(
        &self,
        original_query: &str,
        plan: &InvestigationPlan,
    ) -> AofResult<DeepResult> {
        let agent = self.fleet.spec.agents.first().ok_or_else(|| {
            AofError::runtime("No agents available for synthesis".to_string())
        })?;

        let synthesis_prompt = self.config.synthesizer_prompt.clone().unwrap_or_else(|| {
            r#"Synthesize the investigation findings into a comprehensive report.

Provide:
1. A clear conclusion (root cause)
2. Supporting evidence
3. Actionable recommendations

Format your response as:
CONCLUSION: [your conclusion]
EVIDENCE: [list key evidence points]
RECOMMENDATIONS: [numbered list of recommendations]"#
                .to_string()
        });

        let findings_detail: String = self
            .findings
            .iter()
            .map(|f| {
                format!(
                    "Step {}: {}\nEvidence: {}\nSignificance: {}\n",
                    f.step, f.description, f.evidence, f.significance
                )
            })
            .collect::<Vec<_>>()
            .join("\n---\n");

        let prompt = format!(
            "{}\n\nOriginal Query: {}\n\nInvestigation Findings:\n{}\n\nSynthesize your report:",
            synthesis_prompt, original_query, findings_detail
        );

        let response = {
            let runtime = self.runtime.read().await;
            runtime
                .execute(&agent.name, &prompt)
                .await
                .map_err(|e| AofError::runtime(format!("Synthesis failed: {}", e)))?
        };

        // Parse the synthesis response
        let conclusion = self.extract_section(&response, "CONCLUSION")
            .unwrap_or_else(|| response.clone());
        let recommendations = self.extract_recommendations(&response);

        Ok(DeepResult {
            conclusion,
            evidence: self.findings.clone(),
            recommendations,
            iterations: plan.current_step as u32,
            goal_achieved: self.is_goal_achieved(plan),
        })
    }

    /// Extract a section from the response
    fn extract_section(&self, response: &str, section: &str) -> Option<String> {
        let section_start = format!("{}:", section);
        if let Some(start) = response.find(&section_start) {
            let after_label = &response[start + section_start.len()..];
            // Find the next section or end
            let end = after_label
                .find("\nEVIDENCE:")
                .or_else(|| after_label.find("\nRECOMMENDATIONS:"))
                .unwrap_or(after_label.len());
            Some(after_label[..end].trim().to_string())
        } else {
            None
        }
    }

    /// Extract recommendations from the response
    fn extract_recommendations(&self, response: &str) -> Vec<String> {
        if let Some(section) = self.extract_section(response, "RECOMMENDATIONS") {
            section
                .lines()
                .filter(|line| !line.trim().is_empty())
                .map(|line| {
                    // Remove numbering like "1." or "-"
                    line.trim()
                        .trim_start_matches(|c: char| c.is_numeric() || c == '.' || c == '-' || c == ' ')
                        .to_string()
                })
                .filter(|s| !s.is_empty())
                .collect()
        } else {
            Vec::new()
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_plan() {
        let config = DeepConfig::default();
        let fleet_yaml = r#"
apiVersion: aof.dev/v1
kind: AgentFleet
metadata:
  name: test-fleet
spec:
  agents:
    - name: test-agent
      spec:
        model: openai:gpt-4
        instructions: "Test agent"
        tools: []
  coordination:
    mode: deep
    deep:
      max_iterations: 5
"#;
        let fleet = aof_core::AgentFleet::from_yaml(fleet_yaml).unwrap();
        let runtime = Arc::new(RwLock::new(Runtime::new()));
        let executor = DeepFleetExecutor::new(fleet, runtime);

        let json_response = r#"
Based on your query, here's the plan:
{
  "goal": "Find root cause of API errors",
  "steps": [
    {"step": 1, "description": "Check logs", "action": "Fetch error logs", "tool": null, "completed": false, "result": null},
    {"step": 2, "description": "Analyze metrics", "action": "Query metrics", "tool": null, "completed": false, "result": null}
  ],
  "current_step": 0,
  "needs_revision": false
}
"#;

        let plan = executor.parse_plan(json_response, "Why are we seeing 500 errors?").unwrap();
        assert_eq!(plan.steps.len(), 2);
        assert_eq!(plan.goal, "Find root cause of API errors");
    }

    #[test]
    fn test_extract_section() {
        let config = DeepConfig::default();
        let fleet_yaml = r#"
apiVersion: aof.dev/v1
kind: AgentFleet
metadata:
  name: test-fleet
spec:
  agents:
    - name: test-agent
      spec:
        model: openai:gpt-4
        instructions: "Test agent"
        tools: []
  coordination:
    mode: deep
"#;
        let fleet = aof_core::AgentFleet::from_yaml(fleet_yaml).unwrap();
        let runtime = Arc::new(RwLock::new(Runtime::new()));
        let executor = DeepFleetExecutor::new(fleet, runtime);

        let response = r#"CONCLUSION: Memory leak in API service
EVIDENCE: Pods crashing with OOMKilled
RECOMMENDATIONS:
1. Increase memory limits
2. Profile for memory leaks"#;

        let conclusion = executor.extract_section(response, "CONCLUSION");
        assert_eq!(conclusion, Some("Memory leak in API service".to_string()));

        let recommendations = executor.extract_recommendations(response);
        assert_eq!(recommendations.len(), 2);
        assert_eq!(recommendations[0], "Increase memory limits");
    }
}
