//! Lifecycle WASM Plugin Loader
//!
//! This module loads and interfaces with lifecycle extensions which handle
//! template management, task classification, and system information gathering.
//!
//! # Modes
//!
//! - **WASM Extension**: Full lifecycle with templates and task classification
//! - **Built-in Simple**: Minimal lifecycle without classification (when `lifecycle.enabled = false`)

use anyhow::{Context, Result};
use std::future::Future;
use std::path::PathBuf;
use std::pin::Pin;
use std::sync::Mutex;

use crate::extension::LifecycleExtensionInstance;

macro_rules! debug {
    ($($arg:tt)*) => {
        if std::env::var("RUST_LOG").map(|v| v.to_lowercase().contains("debug")).unwrap_or(false) {
            eprintln!("[DEBUG LIFECYCLE] {}", format!($($arg)*));
        }
    };
}

/// Lifecycle trait for polymorphic lifecycle handling
pub trait Lifecycle: Send + Sync {
    /// Load a template by name
    fn load_template(&self, name: &str) -> Pin<Box<dyn Future<Output = Result<String>> + Send + '_>>;
    
    /// Render a template with variable substitution
    fn render_template(&self, template: &str, variables: &[(String, String)]) -> Pin<Box<dyn Future<Output = Result<String>> + Send + '_>>;
    
    /// Classify a task (returns "query" for simple lifecycle)
    fn classify_task(&self, task_description: &str) -> Pin<Box<dyn Future<Output = Result<(String, f32)>> + Send + '_>>;
    
    /// Get system info variables
    fn get_system_info_variables(&self) -> Pin<Box<dyn Future<Output = Result<Vec<(String, String)>>> + Send + '_>>;
    
    /// Load useful commands
    fn load_useful_commands(&self) -> Pin<Box<dyn Future<Output = Result<String>> + Send + '_>>;
    
    /// Get metadata
    fn get_metadata(&self) -> Pin<Box<dyn Future<Output = Result<String>> + Send + '_>>;
}

/// WASM-based lifecycle extension wrapper
pub struct WasmLifecycle {
    instance: Mutex<LifecycleExtensionInstance>,
    #[allow(dead_code)]
    extension_path: PathBuf,
}

impl WasmLifecycle {
    pub async fn new(extension_dir: PathBuf) -> Result<Self> {
        debug!("Creating lifecycle plugin from: {}", extension_dir.display());

        let mut instance = create_standalone_instance(&extension_dir)?;
        instance.init().context("Failed to initialize lifecycle extension")?;

        let caps = instance.list_capabilities()
            .context("Failed to list capabilities")?;
        if !caps.iter().any(|c| c == "lifecycle") {
            anyhow::bail!("Extension at '{}' does not have lifecycle capability", extension_dir.display());
        }

        debug!("Lifecycle plugin compiled successfully");

        Ok(Self {
            instance: Mutex::new(instance),
            extension_path: extension_dir,
        })
    }
}

impl Lifecycle for WasmLifecycle {
    fn load_template(&self, template_name: &str) -> Pin<Box<dyn Future<Output = Result<String>> + Send + '_>> {
        let template_name = template_name.to_string();
        Box::pin(async move {
            debug!("Loading template: {}", template_name);
            let result = self.instance.lock().unwrap().load_template(&template_name)
                .context(format!("Failed to load template: {}", template_name))?;
            debug!("Template loaded successfully: {} ({} bytes)", template_name, result.len());
            Ok(result)
        })
    }

    fn render_template(&self, template_content: &str, variables: &[(String, String)]) -> Pin<Box<dyn Future<Output = Result<String>> + Send + '_>> {
        let template_content = template_content.to_string();
        let variables = variables.to_vec();
        Box::pin(async move {
            debug!("Rendering template with {} variables", variables.len());
            let result = self.instance.lock().unwrap().render_template(&template_content, &variables)
                .context("Failed to render template")?;
            debug!("Template rendered successfully");
            Ok(result)
        })
    }

    fn classify_task(&self, task_description: &str) -> Pin<Box<dyn Future<Output = Result<(String, f32)>> + Send + '_>> {
        let task_description = task_description.to_string();
        Box::pin(async move {
            debug!("Classifying task: {}", task_description);
            let (task_type, confidence) = self.instance.lock().unwrap().classify_task(&task_description)
                .context("Failed to classify task")?;
            debug!("Task classified as: {} (confidence: {})", task_type, confidence);
            Ok((task_type, confidence))
        })
    }

    fn get_system_info_variables(&self) -> Pin<Box<dyn Future<Output = Result<Vec<(String, String)>>> + Send + '_>> {
        Box::pin(async move {
            debug!("Getting system info variables");
            let vars = self.instance.lock().unwrap().get_system_info_variables()
                .context("Failed to get system info variables")?;
            debug!("Retrieved {} system info variables", vars.len());
            Ok(vars)
        })
    }

    fn load_useful_commands(&self) -> Pin<Box<dyn Future<Output = Result<String>> + Send + '_>> {
        Box::pin(async move {
            debug!("Loading useful commands");
            let result = self.instance.lock().unwrap().load_useful_commands()
                .context("Failed to load useful commands")?;
            debug!("Useful commands loaded successfully ({} bytes)", result.len());
            Ok(result)
        })
    }

    fn get_metadata(&self) -> Pin<Box<dyn Future<Output = Result<String>> + Send + '_>> {
        Box::pin(async move {
            debug!("Getting lifecycle plugin metadata");
            let metadata = self.instance.lock().unwrap().get_metadata()
                .context("Failed to get metadata")?;
            let json = serde_json::json!({
                "id": metadata.id,
                "name": metadata.name,
                "version": metadata.version,
                "api_version": metadata.api_version,
                "description": metadata.description,
            });
            debug!("Metadata retrieved");
            Ok(json.to_string())
        })
    }
}

/// Simple built-in lifecycle (no classification, no templates)
pub struct SimpleLifecycle {
    system_template: Option<String>,
}

impl SimpleLifecycle {
    pub fn new(system_template: Option<String>) -> Self {
        debug!("Using simple built-in lifecycle (no WASM extension)");
        Self { system_template }
    }
}

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

const SIMPLE_SYSTEM_TEMPLATE: &str = r#"You are a helpful AI assistant. Follow the user's instructions directly.

## Tools

You have access to tools. Use them when needed to accomplish tasks.

## Workflow

1. Understand what the user wants
2. Execute the necessary actions using available tools
3. Provide clear, helpful responses

When you have completed the task, respond with your final result."#;

const SIMPLE_TASK_TEMPLATE: &str = r#"Task: {task_description}

{task_context}

Complete this task using available tools. When done, provide your final result."#;

impl Lifecycle for SimpleLifecycle {
    fn load_template(&self, name: &str) -> Pin<Box<dyn Future<Output = Result<String>> + Send + '_>> {
        let name = name.to_string();
        let system_template = self.system_template.clone();
        Box::pin(async move {
            match name.as_str() {
                "system" => Ok(system_template.unwrap_or_else(|| SIMPLE_SYSTEM_TEMPLATE.to_string())),
                "task/query" | "task/feature" | "task/bug_fix" | "task/maintenance" | "task/fallback" => {
                    Ok(SIMPLE_TASK_TEMPLATE.to_string())
                }
                _ => Ok(String::new()),
            }
        })
    }

    fn render_template(&self, template: &str, variables: &[(String, String)]) -> Pin<Box<dyn Future<Output = Result<String>> + Send + '_>> {
        let template = template.to_string();
        let variables = variables.to_vec();
        Box::pin(async move {
            let mut result = template;
            for (key, value) in variables {
                result = result.replace(&format!("{{{}}}", key), &value);
            }
            Ok(result)
        })
    }

    fn classify_task(&self, _task_description: &str) -> Pin<Box<dyn Future<Output = Result<(String, f32)>> + Send + '_>> {
        Box::pin(async move {
            Ok(("query".to_string(), 1.0))
        })
    }

    fn get_system_info_variables(&self) -> Pin<Box<dyn Future<Output = Result<Vec<(String, String)>>> + Send + '_>> {
        Box::pin(async move {
            Ok(vec![
                ("working_directory".to_string(), std::env::current_dir()
                    .unwrap_or_default()
                    .to_string_lossy()
                    .to_string()),
            ])
        })
    }

    fn load_useful_commands(&self) -> Pin<Box<dyn Future<Output = Result<String>> + Send + '_>> {
        Box::pin(async move {
            Ok("Use available tools to complete tasks.".to_string())
        })
    }

    fn get_metadata(&self) -> Pin<Box<dyn Future<Output = Result<String>> + Send + '_>> {
        Box::pin(async move {
            Ok(r#"{"id": "simple-lifecycle", "name": "Simple Lifecycle", "version": "0.1.0", "description": "Built-in simple lifecycle without classification"}"#.to_string())
        })
    }
}

/// Legacy LifecyclePlugin type alias for backward compatibility
pub type LifecyclePlugin = WasmLifecycle;

fn create_standalone_instance(extension_dir: &PathBuf) -> Result<LifecycleExtensionInstance> {
    use crate::extension::ExtensionManifest;
    use std::sync::Arc;
    use wasmtime::component::Component;
    use wasmtime::{Config, Engine};

    let manifest_path = extension_dir.join("extension.toml");
    debug!("Loading manifest from: {}", manifest_path.display());
    let manifest = ExtensionManifest::from_file(&manifest_path)
        .with_context(|| format!("Failed to load extension manifest: {}", manifest_path.display()))?;
    debug!("Manifest loaded: {} v{}", manifest.extension.name, manifest.extension.version);

    let mut config = Config::new();
    config.wasm_component_model(true);
    debug!("Creating WASM engine (sync mode)...");
    let engine = Arc::new(Engine::new(&config)
        .context("Failed to create WASM engine for lifecycle")?);
    debug!("WASM engine created");

    let wasm_path = extension_dir.join(&manifest.lib.path);
    debug!("Loading WASM component from: {}", wasm_path.display());
    let wasm_bytes = std::fs::read(&wasm_path)
        .with_context(|| format!("Failed to read WASM file: {}", wasm_path.display()))?;
    debug!("Read {} bytes from WASM file", wasm_bytes.len());
    
    let component = Component::from_binary(&engine, &wasm_bytes)
        .with_context(|| format!("Failed to parse WASM component: {}", wasm_path.display()))?;
    debug!("WASM component parsed successfully");

    debug!("Creating LifecycleExtensionInstance...");
    let instance = LifecycleExtensionInstance::new(&engine, &component)
        .with_context(|| "Failed to create lifecycle extension instance from WASM component")?;
    debug!("LifecycleExtensionInstance created successfully");

    Ok(instance)
}

trait ExpandHome {
    fn expand_home(&self) -> Result<PathBuf>;
}

impl ExpandHome for PathBuf {
    fn expand_home(&self) -> Result<PathBuf> {
        if let Some(path_str) = self.to_str() {
            if path_str.starts_with("~") {
                if let Ok(home) = crate::get_home_dir() {
                    let expanded = path_str.replacen("~", &home, 1);
                    return Ok(PathBuf::from(expanded));
                }
            }
        }
        Ok(self.clone())
    }
}

/// Find and load the lifecycle extension (WASM or simple built-in)
/// 
/// If `lifecycle_enabled` is false, returns the simple built-in lifecycle.
/// Otherwise searches for WASM extension.
pub async fn find_lifecycle_plugin_with_config(
    lifecycle_enabled: bool,
    system_template: Option<String>,
) -> Result<Box<dyn Lifecycle>> {
    if !lifecycle_enabled {
        return Ok(Box::new(SimpleLifecycle::new(system_template)));
    }
    
    let agent_name = std::env::var("ABK_AGENT_NAME").unwrap_or_else(|_| "NO_AGENT_NAME".to_string());
    
    let extension_paths = vec![
        PathBuf::from("extensions/coder-lifecycle"),
        PathBuf::from(format!("~/.{}/extensions/coder-lifecycle", agent_name)).expand_home()?,
    ];

    for path in &extension_paths {
        let manifest_path = path.join("extension.toml");
        if manifest_path.exists() {
            debug!("Found lifecycle extension at: {}", path.display());
            return Ok(Box::new(WasmLifecycle::new(path.clone()).await?));
        }
    }

    // Legacy paths
    let legacy_paths = vec![
        PathBuf::from("providers/lifecycle"),
        PathBuf::from(format!("~/.{}/providers/lifecycle", agent_name)).expand_home()?,
    ];

    for path in &legacy_paths {
        let manifest_path = path.join("extension.toml");
        if manifest_path.exists() {
            debug!("Found legacy lifecycle extension at: {}", path.display());
            eprintln!("[WARN] Using deprecated lifecycle plugin location. Please migrate to ~/.{}/extensions/coder-lifecycle/", agent_name);
            return Ok(Box::new(WasmLifecycle::new(path.clone()).await?));
        }
        
        let wasm_path = path.join("lifecycle.wasm");
        if wasm_path.exists() {
            anyhow::bail!(
                "Found old-format lifecycle plugin at {}. \n\
                The new extension system requires an extension.toml manifest.\n\
                Please update to the new coder-lifecycle-wasm extension format.\n\
                See: https://github.com/podtan/coder-lifecycle-wasm",
                wasm_path.display()
            );
        }
    }

    // Not found, fall back to simple lifecycle
    debug!("Lifecycle extension not found, using simple built-in lifecycle");
    Ok(Box::new(SimpleLifecycle::new(system_template)))
}

/// Legacy function for backward compatibility (always tries to load WASM)
pub async fn find_lifecycle_plugin() -> Result<LifecyclePlugin> {
    let agent_name = std::env::var("ABK_AGENT_NAME").unwrap_or_else(|_| "NO_AGENT_NAME".to_string());
    
    let extension_paths = vec![
        PathBuf::from("extensions/coder-lifecycle"),
        PathBuf::from(format!("~/.{}/extensions/coder-lifecycle", agent_name)).expand_home()?,
    ];

    for path in &extension_paths {
        let manifest_path = path.join("extension.toml");
        if manifest_path.exists() {
            debug!("Found lifecycle extension at: {}", path.display());
            return LifecyclePlugin::new(path.clone()).await;
        }
    }

    anyhow::bail!(
        "Lifecycle extension not found. Expected at one of:\n  \
        - extensions/coder-lifecycle/extension.toml\n  \
        - ~/.{}/extensions/coder-lifecycle/extension.toml\n\n\
        Install with: trustee extension install <path-to-coder-lifecycle-wasm>\n\
        Or set lifecycle.enabled = false in config to use simple built-in lifecycle.",
        agent_name
    )
}
