//! Agent execution layer
//!
//! Executes agents with context injection and output validation

use crate::error::{Error, Result};
use crate::nickel::AgentDefinition;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::process::Command;
use tera::{Context, Tera};

/// Agent execution result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionResult {
    /// Generated output from agent
    pub output: String,
    /// Whether validation passed
    pub validation_passed: bool,
    /// Validation errors if any
    #[serde(default)]
    pub validation_errors: Vec<String>,
    /// Execution metadata
    #[serde(default)]
    pub metadata: ExecutionMetadata,
}

/// Execution metadata
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ExecutionMetadata {
    /// Time taken in milliseconds
    pub duration_ms: Option<u64>,
    /// Token count if available
    pub tokens: Option<usize>,
    /// Model used
    pub model: Option<String>,
}

/// Context data assembled from imports and shell commands
#[derive(Debug, Clone)]
pub struct ContextData {
    /// File imports content (alias → content)
    pub imports: HashMap<String, String>,
    /// Shell command outputs (alias → output)
    pub shell_outputs: HashMap<String, String>,
}

/// Agent executor
pub struct AgentExecutor;

impl AgentExecutor {
    /// Create a new agent executor
    pub fn new() -> Self {
        Self
    }

    /// Execute an agent with given inputs
    pub async fn execute(
        &self,
        agent: &AgentDefinition,
        inputs: HashMap<String, serde_json::Value>,
    ) -> Result<ExecutionResult> {
        let start = std::time::Instant::now();

        // Load context data
        let context_data = self.load_context(agent).await?;

        // Render template with inputs and context
        let rendered_prompt = self.render_template(agent, &inputs, &context_data)?;

        // Execute LLM with rendered prompt
        let (output, tokens) = self.execute_llm(agent, &rendered_prompt).await?;

        // Validate output
        let validation_errors = self.validate_output(&output, agent)?;
        let validation_passed = validation_errors.is_empty();

        let duration_ms = start.elapsed().as_millis() as u64;

        Ok(ExecutionResult {
            output,
            validation_passed,
            validation_errors,
            metadata: ExecutionMetadata {
                duration_ms: Some(duration_ms),
                tokens,
                model: Some(agent.config.llm.clone()),
            },
        })
    }

    /// Execute an agent with streaming output
    ///
    /// The callback is called for each chunk of output as it arrives.
    /// Returns the final execution result when the stream completes.
    pub async fn execute_streaming<F>(
        &self,
        agent: &AgentDefinition,
        inputs: HashMap<String, serde_json::Value>,
        mut on_chunk: F,
    ) -> Result<ExecutionResult>
    where
        F: FnMut(&str),
    {
        use crate::llm::{create_provider, LlmMessage, LlmRequest, MessageRole, StreamChunk};
        use futures::stream::StreamExt;

        let start = std::time::Instant::now();

        // Load context data
        let context_data = self.load_context(agent).await?;

        // Render template with inputs and context
        let rendered_prompt = self.render_template(agent, &inputs, &context_data)?;

        // Create provider and build request
        let provider = create_provider(&agent.config.llm)?;
        let request = LlmRequest {
            model: agent.config.llm.clone(),
            messages: vec![LlmMessage {
                role: MessageRole::User,
                content: rendered_prompt.clone(),
            }],
            max_tokens: Some(agent.config.max_tokens),
            temperature: Some(agent.config.temperature),
            system: Some(format!("You are a {}.", agent.config.role)),
        };

        // Stream the response
        let mut stream = provider.stream(request).await?;
        let mut full_output = String::new();
        let mut tokens = None;

        while let Some(chunk_result) = stream.next().await {
            match chunk_result? {
                StreamChunk::Content(text) => {
                    full_output.push_str(&text);
                    on_chunk(&text);
                }
                StreamChunk::Done(metadata) => {
                    tokens = metadata.usage.map(|u| u.total_tokens);
                }
                StreamChunk::Error(err) => {
                    return Err(Error::execution("Streaming error", err));
                }
            }
        }

        // Validate output
        let validation_errors = self.validate_output(&full_output, agent)?;
        let validation_passed = validation_errors.is_empty();

        let duration_ms = start.elapsed().as_millis() as u64;

        Ok(ExecutionResult {
            output: full_output,
            validation_passed,
            validation_errors,
            metadata: ExecutionMetadata {
                duration_ms: Some(duration_ms),
                tokens,
                model: Some(agent.config.llm.clone()),
            },
        })
    }

    /// Execute LLM with the rendered prompt
    async fn execute_llm(
        &self,
        agent: &AgentDefinition,
        prompt: &str,
    ) -> Result<(String, Option<usize>)> {
        use crate::llm::{create_provider, LlmMessage, LlmRequest, MessageRole};

        // Create provider based on model name
        let provider = create_provider(&agent.config.llm)?;

        // Build request
        let request = LlmRequest {
            model: agent.config.llm.clone(),
            messages: vec![LlmMessage {
                role: MessageRole::User,
                content: prompt.to_string(),
            }],
            max_tokens: Some(agent.config.max_tokens),
            temperature: Some(agent.config.temperature),
            system: Some(format!("You are a {}.", agent.config.role)),
        };

        // Execute request
        let response = provider.complete(request).await?;

        // Extract token count
        let tokens = response.usage.as_ref().map(|u| u.total_tokens);

        Ok((response.content, tokens))
    }

    /// Load context from imports and shell commands
    async fn load_context(&self, agent: &AgentDefinition) -> Result<ContextData> {
        let mut imports = HashMap::new();
        let mut shell_outputs = HashMap::new();

        if let Some(context) = &agent.context {
            // Load file imports
            for import_path in &context.imports {
                let content = self.load_import(import_path).await?;
                imports.insert(import_path.clone(), content);
            }

            // Execute shell commands
            for shell_cmd in &context.shell_commands {
                let output = self.execute_shell_command(shell_cmd)?;
                shell_outputs.insert(shell_cmd.clone(), output);
            }
        }

        Ok(ContextData {
            imports,
            shell_outputs,
        })
    }

    /// Load content from import path (file, glob, or URL)
    async fn load_import(&self, path: &str) -> Result<String> {
        // Check if it's a URL
        if path.starts_with("http://") || path.starts_with("https://") {
            return self.load_url(path).await;
        }

        // Check if it's a glob pattern
        if path.contains('*') {
            return self.load_glob(path);
        }

        // Regular file path
        std::fs::read_to_string(path).map_err(|e| {
            Error::execution(
                format!("Failed to read import file: {}", path),
                e.to_string(),
            )
        })
    }

    /// Load content from URL
    async fn load_url(&self, url: &str) -> Result<String> {
        let response = reqwest::get(url).await.map_err(|e| {
            Error::execution(format!("Failed to fetch URL: {}", url), e.to_string())
        })?;

        response.text().await.map_err(|e| {
            Error::execution(
                format!("Failed to read response from URL: {}", url),
                e.to_string(),
            )
        })
    }

    /// Load files matching glob pattern
    fn load_glob(&self, pattern: &str) -> Result<String> {
        use globset::GlobBuilder;
        use ignore::WalkBuilder;

        let glob = GlobBuilder::new(pattern)
            .build()
            .map_err(|e| {
                Error::execution(format!("Invalid glob pattern: {}", pattern), e.to_string())
            })?
            .compile_matcher();

        let mut contents = Vec::new();

        // Extract base path from pattern
        let base_path = pattern.split("**").next().unwrap_or(".");

        for entry in WalkBuilder::new(base_path).build() {
            let entry = entry.map_err(|e| {
                Error::execution(
                    format!("Failed to walk directory: {}", base_path),
                    e.to_string(),
                )
            })?;

            if entry.file_type().map(|ft| ft.is_file()).unwrap_or(false)
                && glob.is_match(entry.path())
            {
                let content = std::fs::read_to_string(entry.path()).map_err(|e| {
                    Error::execution(
                        format!("Failed to read file: {:?}", entry.path()),
                        e.to_string(),
                    )
                })?;

                contents.push(format!("// File: {}\n{}", entry.path().display(), content));
            }
        }

        if contents.is_empty() {
            return Err(Error::execution(
                format!("No files matched glob pattern: {}", pattern),
                "glob returned empty",
            ));
        }

        Ok(contents.join("\n\n"))
    }

    /// Execute shell command and capture output
    fn execute_shell_command(&self, command: &str) -> Result<String> {
        let output = Command::new("sh")
            .arg("-c")
            .arg(command)
            .output()
            .map_err(|e| {
                Error::execution(
                    format!("Failed to execute shell command: {}", command),
                    e.to_string(),
                )
            })?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(Error::execution(
                format!("Shell command failed: {}", command),
                stderr.to_string(),
            ));
        }

        Ok(String::from_utf8_lossy(&output.stdout).to_string())
    }

    /// Render template with inputs and context
    fn render_template(
        &self,
        agent: &AgentDefinition,
        inputs: &HashMap<String, serde_json::Value>,
        context_data: &ContextData,
    ) -> Result<String> {
        let mut tera = Tera::default();

        // Add template
        tera.add_raw_template("agent", &agent.template)
            .map_err(|e| Error::execution("Failed to parse template", e.to_string()))?;

        // Build context
        let mut ctx = Context::new();

        // Add inputs
        for (key, value) in inputs {
            ctx.insert(key, value);
        }

        // Add agent config
        ctx.insert("role", &agent.config.role);
        ctx.insert("llm", &agent.config.llm);
        ctx.insert("tools", &agent.config.tools);

        // Add imports
        for (alias, content) in &context_data.imports {
            ctx.insert(alias, content);
        }

        // Add shell outputs
        for (alias, output) in &context_data.shell_outputs {
            ctx.insert(alias, output);
        }

        // Render
        tera.render("agent", &ctx)
            .map_err(|e| Error::execution("Failed to render template", e.to_string()))
    }

    /// Validate output against agent validation rules
    pub fn validate_output(&self, output: &str, agent: &AgentDefinition) -> Result<Vec<String>> {
        let mut errors = Vec::new();

        if let Some(validation) = &agent.validation {
            // Check must_contain
            for pattern in &validation.must_contain {
                if !output.contains(pattern) {
                    errors.push(format!("Output must contain: {}", pattern));
                }
            }

            // Check must_not_contain
            for pattern in &validation.must_not_contain {
                if output.contains(pattern) {
                    errors.push(format!("Output must not contain: {}", pattern));
                }
            }

            // Check min_length
            if let Some(min_len) = validation.min_length {
                if output.len() < min_len {
                    errors.push(format!(
                        "Output too short: {} chars (minimum: {})",
                        output.len(),
                        min_len
                    ));
                }
            }

            // Check max_length
            if let Some(max_len) = validation.max_length {
                if output.len() > max_len {
                    errors.push(format!(
                        "Output too long: {} chars (maximum: {})",
                        output.len(),
                        max_len
                    ));
                }
            }

            // Check format
            match validation.format.as_str() {
                "json" => {
                    if serde_json::from_str::<serde_json::Value>(output).is_err() {
                        errors.push("Output is not valid JSON".to_string());
                    }
                }
                "yaml" => {
                    if serde_yaml::from_str::<serde_yaml::Value>(output).is_err() {
                        errors.push("Output is not valid YAML".to_string());
                    }
                }
                _ => {
                    // Markdown, text, and other formats are permissive - any text is valid
                }
            }
        }

        Ok(errors)
    }
}

impl Default for AgentExecutor {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::nickel::{AgentConfig, ValidationRules};

    #[test]
    fn test_validate_output_must_contain() {
        let executor = AgentExecutor::new();
        let agent = AgentDefinition {
            config: AgentConfig {
                role: "tester".to_string(),
                llm: "claude".to_string(),
                tools: vec![],
                max_tokens: 4096,
                temperature: 0.7,
            },
            inputs: HashMap::new(),
            context: None,
            validation: Some(ValidationRules {
                must_contain: vec!["test".to_string(), "pass".to_string()],
                must_not_contain: vec![],
                format: "markdown".to_string(),
                min_length: None,
                max_length: None,
            }),
            template: "".to_string(),
        };

        let output = "This is a test and it should pass";
        let errors = executor.validate_output(output, &agent).unwrap();
        assert!(errors.is_empty());

        let output_bad = "This is missing something";
        let errors = executor.validate_output(output_bad, &agent).unwrap();
        assert!(!errors.is_empty());
        assert!(errors[0].contains("must contain"));
    }

    #[test]
    fn test_validate_output_length() {
        let executor = AgentExecutor::new();
        let agent = AgentDefinition {
            config: AgentConfig {
                role: "tester".to_string(),
                llm: "claude".to_string(),
                tools: vec![],
                max_tokens: 4096,
                temperature: 0.7,
            },
            inputs: HashMap::new(),
            context: None,
            validation: Some(ValidationRules {
                must_contain: vec![],
                must_not_contain: vec![],
                format: "markdown".to_string(),
                min_length: Some(10),
                max_length: Some(50),
            }),
            template: "".to_string(),
        };

        let output = "Just right length";
        let errors = executor.validate_output(output, &agent).unwrap();
        assert!(errors.is_empty());

        let output_short = "Short";
        let errors = executor.validate_output(output_short, &agent).unwrap();
        assert!(!errors.is_empty());
        assert!(errors[0].contains("too short"));

        let output_long = "This is way too long for the maximum length constraint that we've set";
        let errors = executor.validate_output(output_long, &agent).unwrap();
        assert!(!errors.is_empty());
        assert!(errors[0].contains("too long"));
    }

    #[test]
    fn test_validate_output_json_format() {
        let executor = AgentExecutor::new();
        let agent = AgentDefinition {
            config: AgentConfig {
                role: "tester".to_string(),
                llm: "claude".to_string(),
                tools: vec![],
                max_tokens: 4096,
                temperature: 0.7,
            },
            inputs: HashMap::new(),
            context: None,
            validation: Some(ValidationRules {
                must_contain: vec![],
                must_not_contain: vec![],
                format: "json".to_string(),
                min_length: None,
                max_length: None,
            }),
            template: "".to_string(),
        };

        let output_valid = r#"{"key": "value"}"#;
        let errors = executor.validate_output(output_valid, &agent).unwrap();
        assert!(errors.is_empty());

        let output_invalid = "not json at all";
        let errors = executor.validate_output(output_invalid, &agent).unwrap();
        assert!(!errors.is_empty());
        assert!(errors[0].contains("not valid JSON"));
    }

    #[test]
    fn test_render_template_basic() {
        let executor = AgentExecutor::new();
        let agent = AgentDefinition {
            config: AgentConfig {
                role: "architect".to_string(),
                llm: "claude-opus-4".to_string(),
                tools: vec![],
                max_tokens: 4096,
                temperature: 0.7,
            },
            inputs: HashMap::new(),
            context: None,
            validation: None,
            template: "Design feature: {{ feature_name }}".to_string(),
        };

        let mut inputs = HashMap::new();
        inputs.insert("feature_name".to_string(), serde_json::json!("auth"));

        let context_data = ContextData {
            imports: HashMap::new(),
            shell_outputs: HashMap::new(),
        };

        let result = executor.render_template(&agent, &inputs, &context_data);
        assert!(result.is_ok());
        let rendered = result.unwrap();
        assert_eq!(rendered, "Design feature: auth");
    }

    #[tokio::test]
    async fn test_execute_shell_command() {
        let executor = AgentExecutor::new();
        let result = executor.execute_shell_command("echo hello");
        assert!(result.is_ok());
        assert_eq!(result.unwrap().trim(), "hello");
    }

    #[tokio::test]
    #[ignore] // Requires ANTHROPIC_API_KEY
    async fn test_execute_with_real_llm() {
        let executor = AgentExecutor::new();
        let agent = AgentDefinition {
            config: AgentConfig {
                role: "assistant".to_string(),
                llm: "claude-3-5-haiku-20241022".to_string(),
                tools: vec![],
                max_tokens: 100,
                temperature: 0.7,
            },
            inputs: HashMap::new(),
            context: None,
            validation: None,
            template: "Say hello to {{ name }} in exactly 3 words.".to_string(),
        };

        let mut inputs = HashMap::new();
        inputs.insert("name".to_string(), serde_json::json!("Alice"));

        let result = executor.execute(&agent, inputs).await;
        assert!(result.is_ok());

        let exec_result = result.unwrap();
        assert!(!exec_result.output.is_empty());
        assert!(exec_result.metadata.tokens.is_some());

        println!("LLM Response: {}", exec_result.output);
        println!("Tokens used: {:?}", exec_result.metadata.tokens);
    }

    // Mock provider for testing streaming without API calls
    mod mock {
        use super::*;
        use crate::llm::{
            LlmProvider, LlmRequest, LlmResponse, LlmStream, StreamChunk, StreamMetadata,
            TokenUsage,
        };
        use async_trait::async_trait;
        use futures::stream;

        pub struct MockStreamingProvider {
            chunks: Vec<String>,
            should_error: bool,
        }

        impl MockStreamingProvider {
            pub fn new(chunks: Vec<String>) -> Self {
                Self {
                    chunks,
                    should_error: false,
                }
            }

            pub fn with_error() -> Self {
                Self {
                    chunks: vec![],
                    should_error: true,
                }
            }
        }

        #[async_trait]
        impl LlmProvider for MockStreamingProvider {
            async fn complete(&self, _request: LlmRequest) -> Result<LlmResponse> {
                Ok(LlmResponse {
                    content: self.chunks.join(""),
                    model: "mock-model".to_string(),
                    usage: Some(TokenUsage {
                        input_tokens: 10,
                        output_tokens: 20,
                        total_tokens: 30,
                    }),
                })
            }

            async fn stream(&self, _request: LlmRequest) -> Result<LlmStream> {
                if self.should_error {
                    let error_stream =
                        stream::iter(vec![Ok(StreamChunk::Error("Mock error".to_string()))]);
                    return Ok(Box::pin(error_stream));
                }

                let mut events: Vec<Result<StreamChunk>> = self
                    .chunks
                    .iter()
                    .map(|chunk| Ok(StreamChunk::Content(chunk.clone())))
                    .collect();

                // Add final Done event with usage
                events.push(Ok(StreamChunk::Done(StreamMetadata {
                    model: "mock-model".to_string(),
                    usage: Some(TokenUsage {
                        input_tokens: 10,
                        output_tokens: 20,
                        total_tokens: 30,
                    }),
                })));

                Ok(Box::pin(stream::iter(events)))
            }

            fn name(&self) -> &str {
                "MockStreaming"
            }
        }
    }

    #[tokio::test]
    async fn test_streaming_chunks_accumulate() {
        use crate::llm::LlmProvider;
        use std::sync::{Arc, Mutex};

        // Track chunks as they arrive
        let received_chunks = Arc::new(Mutex::new(Vec::new()));
        let chunks_clone = Arc::clone(&received_chunks);

        // Mock provider that streams "Hello", " ", "world"
        let mock_provider = mock::MockStreamingProvider::new(vec![
            "Hello".to_string(),
            " ".to_string(),
            "world".to_string(),
        ]);

        // Manually test streaming with mock provider
        use crate::llm::{LlmMessage, LlmRequest, MessageRole, StreamChunk};
        use futures::stream::StreamExt;

        let request = LlmRequest {
            model: "mock-model".to_string(),
            messages: vec![LlmMessage {
                role: MessageRole::User,
                content: "Test".to_string(),
            }],
            max_tokens: Some(100),
            temperature: Some(0.7),
            system: Some("Test system".to_string()),
        };

        let mut stream = mock_provider.stream(request).await.unwrap();
        let mut full_output = String::new();
        let mut tokens = None;

        while let Some(chunk_result) = stream.next().await {
            match chunk_result.unwrap() {
                StreamChunk::Content(text) => {
                    chunks_clone.lock().unwrap().push(text.clone());
                    full_output.push_str(&text);
                }
                StreamChunk::Done(metadata) => {
                    tokens = metadata.usage.map(|u| u.total_tokens);
                }
                StreamChunk::Error(_) => panic!("Unexpected error"),
            }
        }

        // Verify chunks were received in order
        let chunks = received_chunks.lock().unwrap();
        assert_eq!(chunks.len(), 3);
        assert_eq!(chunks[0], "Hello");
        assert_eq!(chunks[1], " ");
        assert_eq!(chunks[2], "world");

        // Verify full output
        assert_eq!(full_output, "Hello world");

        // Verify tokens
        assert_eq!(tokens, Some(30));
    }

    #[tokio::test]
    async fn test_streaming_error_handling() {
        use crate::llm::{LlmMessage, LlmProvider, LlmRequest, MessageRole, StreamChunk};
        use futures::stream::StreamExt;

        let mock_provider = mock::MockStreamingProvider::with_error();

        let request = LlmRequest {
            model: "mock-model".to_string(),
            messages: vec![LlmMessage {
                role: MessageRole::User,
                content: "Test".to_string(),
            }],
            max_tokens: Some(100),
            temperature: Some(0.7),
            system: Some("Test system".to_string()),
        };

        let mut stream = mock_provider.stream(request).await.unwrap();
        let mut error_received = false;

        while let Some(chunk_result) = stream.next().await {
            if let StreamChunk::Error(msg) = chunk_result.unwrap() {
                assert_eq!(msg, "Mock error");
                error_received = true;
            }
        }

        assert!(error_received);
    }

    #[tokio::test]
    async fn test_streaming_empty_chunks() {
        use crate::llm::{LlmMessage, LlmProvider, LlmRequest, MessageRole, StreamChunk};
        use futures::stream::StreamExt;

        // Provider with no content chunks, just Done
        let mock_provider = mock::MockStreamingProvider::new(vec![]);

        let request = LlmRequest {
            model: "mock-model".to_string(),
            messages: vec![LlmMessage {
                role: MessageRole::User,
                content: "Test".to_string(),
            }],
            max_tokens: Some(100),
            temperature: Some(0.7),
            system: Some("Test system".to_string()),
        };

        let mut stream = mock_provider.stream(request).await.unwrap();
        let mut full_output = String::new();
        let mut done_received = false;

        while let Some(chunk_result) = stream.next().await {
            match chunk_result.unwrap() {
                StreamChunk::Content(text) => {
                    full_output.push_str(&text);
                }
                StreamChunk::Done(_) => {
                    done_received = true;
                }
                StreamChunk::Error(_) => panic!("Unexpected error"),
            }
        }

        assert_eq!(full_output, "");
        assert!(done_received);
    }
}
