use super::*;

use crate::{
    client::Model,
    function::{run_llm_function, Functions},
};

use anyhow::{Context, Result};
use inquire::{validator::Validation, Text};
use rust_embed::Embed;
use serde::{Deserialize, Serialize};
use std::{fs, fs::read_to_string, path::Path};
use std::ffi::OsStr;

const DEFAULT_AGENT_NAME: &str = "rag";

pub type AgentVariables = IndexMap<String, String>;

#[derive(Embed)]
#[folder = "assets/agents/"]
struct AgentAssets;

#[derive(Debug, Clone)]
pub struct Agent {
    name: String,
    config: AgentConfig,
    shared_variables: AgentVariables,
    session_variables: Option<AgentVariables>,
    shared_dynamic_instructions: Option<String>,
    session_dynamic_instructions: Option<String>,
    functions: Functions,
    rag: Option<Arc<Rag>>,
    model: Model,
}

impl Agent {
    pub fn install_builtin_agents() -> Result<()> {
        info!(
            "Installing built-in agents in {}",
            Config::agents_data_dir().display()
        );

        for file in AgentAssets::iter() {
            debug!("Processing agent file: {}", file.as_ref());

            let embedded_file = AgentAssets::get(&file)
                .ok_or_else(|| anyhow!("Failed to load embedded agent file: {}", file.as_ref()))?;
            let content = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) };
            let file_path = Config::agents_data_dir().join(file.as_ref());
					let file_extension = file_path
						.extension()
						.and_then(OsStr::to_str)
						.map(|s| s.to_lowercase());
					let is_script = matches!(file_extension.as_deref(), Some("sh") | Some("py"));

            if file_path.exists() {
                debug!(
                    "Agent file already exists, skipping: {}",
                    file_path.display()
                );
                continue;
            }

            ensure_parent_exists(&file_path)?;
            info!("Creating agent file: {}", file_path.display());
            let mut agent_file = File::create(&file_path)?;
            agent_file.write_all(content.as_bytes())?;

					#[cfg(unix)]
					if is_script {
						use std::os::unix::fs::PermissionsExt;
						fs::set_permissions(&file_path, fs::Permissions::from_mode(0o755))?;
					}
        }

        Ok(())
    }

    pub async fn init(
        config: &GlobalConfig,
        name: &str,
        abort_signal: AbortSignal,
    ) -> Result<Self> {
        let agent_data_dir = Config::agent_data_dir(name);
        let loaders = config.read().document_loaders.clone();
        let rag_path = Config::agent_rag_file(name, DEFAULT_AGENT_NAME);
        let config_path = Config::agent_config_file(name);
        let mut agent_config = if config_path.exists() {
            AgentConfig::load(&config_path)?
        } else {
            bail!("Agent config file not found at '{}'", config_path.display())
        };
        let mut functions = Functions::init_agent(name, &agent_config.global_tools)?;

        config.write().functions.clear_mcp_meta_functions();
        let mcp_servers =
            (!agent_config.mcp_servers.is_empty()).then(|| agent_config.mcp_servers.join(","));
        let registry = config
            .write()
            .mcp_registry
            .take()
            .expect("MCP registry should be initialized");
        let new_mcp_registry =
            McpRegistry::reinit(registry, mcp_servers, abort_signal.clone()).await?;

        if !new_mcp_registry.is_empty() {
            functions.append_mcp_meta_functions(new_mcp_registry.list_servers());
        }

        config.write().mcp_registry = Some(new_mcp_registry);
        agent_config.replace_tools_placeholder(&functions);

        agent_config.load_envs(&config.read());

        let model = {
            let config = config.read();
            match agent_config.model_id.as_ref() {
                Some(model_id) => Model::retrieve_model(&config, model_id, ModelType::Chat)?,
                None => {
                    if agent_config.temperature.is_none() {
                        agent_config.temperature = config.temperature;
                    }
                    if agent_config.top_p.is_none() {
                        agent_config.top_p = config.top_p;
                    }
                    config.current_model().clone()
                }
            }
        };

        let rag = if rag_path.exists() {
            Some(Arc::new(Rag::load(config, DEFAULT_AGENT_NAME, &rag_path)?))
        } else if !agent_config.documents.is_empty() && !config.read().info_flag {
            let mut ans = false;
            if *IS_STDOUT_TERMINAL {
                ans = Confirm::new("The agent has documents attached, init RAG?")
                    .with_default(true)
                    .prompt()?;
            }
            if ans {
                let mut document_paths = vec![];
                for path in &agent_config.documents {
                    if is_url(path) {
                        document_paths.push(path.to_string());
                    } else if is_loader_protocol(&loaders, path) {
                        let (protocol, document_path) = path
                            .split_once(':')
                            .with_context(|| "Invalid loader protocol path")?;
                        let resolved_path = resolve_home_dir(document_path);
                        let new_path = if Path::new(&resolved_path).is_relative() {
                            safe_join_path(&agent_data_dir, resolved_path)
                                .ok_or_else(|| anyhow!("Invalid document path: '{path}'"))?
                        } else {
                            PathBuf::from(&resolved_path)
                        };
                        document_paths.push(format!("{}:{}", protocol, new_path.display()));
                    } else if Path::new(&resolve_home_dir(path)).is_relative() {
                        let new_path = safe_join_path(&agent_data_dir, path)
                            .ok_or_else(|| anyhow!("Invalid document path: '{path}'"))?;
                        document_paths.push(new_path.display().to_string())
                    } else {
                        document_paths.push(path.to_string())
                    }
                }
                let rag =
                    Rag::init(config, "rag", &rag_path, &document_paths, abort_signal).await?;
                Some(Arc::new(rag))
            } else {
                None
            }
        } else {
            None
        };

        Ok(Self {
            name: name.to_string(),
            config: agent_config,
            shared_variables: Default::default(),
            session_variables: None,
            shared_dynamic_instructions: None,
            session_dynamic_instructions: None,
            functions,
            rag,
            model,
        })
    }

    pub fn init_agent_variables(
        agent_variables: &[AgentVariable],
        no_interaction: bool,
    ) -> Result<AgentVariables> {
        let mut output = IndexMap::new();
        if agent_variables.is_empty() {
            return Ok(output);
        }
        let mut printed = false;
        let mut unset_variables = vec![];
        for agent_variable in agent_variables {
            let key = agent_variable.name.clone();
            if let Some(value) = agent_variable.default.clone() {
                output.insert(key, value);
                continue;
            }
            if no_interaction {
                continue;
            }
            if *IS_STDOUT_TERMINAL {
                if !printed {
                    println!("⚙ Init agent variables...");
                    printed = true;
                }
                let value = Text::new(&format!(
                    "{} ({}):",
                    agent_variable.name, agent_variable.description
                ))
                .with_validator(|input: &str| {
                    if input.trim().is_empty() {
                        Ok(Validation::Invalid("This field is required".into()))
                    } else {
                        Ok(Validation::Valid)
                    }
                })
                .prompt()?;
                output.insert(key, value);
            } else {
                unset_variables.push(agent_variable)
            }
        }
        if !unset_variables.is_empty() {
            bail!(
                "The following agent variables are required:\n{}",
                unset_variables
                    .iter()
                    .map(|v| format!("  - {}: {}", v.name, v.description))
                    .collect::<Vec<_>>()
                    .join("\n")
            )
        }
        Ok(output)
    }

    pub fn export(&self) -> Result<String> {
        let mut value = json!({});
        value["name"] = json!(self.name());
        let variables = self.variables();
        if !variables.is_empty() {
            value["variables"] = serde_json::to_value(variables)?;
        }
        value["config"] = json!(self.config);
        let mut config = self.config.clone();
        config.instructions = self.interpolated_instructions();
        value["definition"] = json!(config);
        value["data_dir"] = Config::agent_data_dir(&self.name)
            .display()
            .to_string()
            .into();
        value["config_file"] = Config::agent_config_file(&self.name)
            .display()
            .to_string()
            .into();
        let data = serde_yaml::to_string(&value)?;
        Ok(data)
    }

    pub fn banner(&self) -> String {
        self.config.banner()
    }

    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn functions(&self) -> &Functions {
        &self.functions
    }

    pub fn rag(&self) -> Option<Arc<Rag>> {
        self.rag.clone()
    }

    pub fn conversation_starters(&self) -> &[String] {
        &self.config.conversation_starters
    }

    pub fn interpolated_instructions(&self) -> String {
        let mut output = self
            .session_dynamic_instructions
            .clone()
            .or_else(|| self.shared_dynamic_instructions.clone())
            .unwrap_or_else(|| self.config.instructions.clone());
        for (k, v) in self.variables() {
            output = output.replace(&format!("{{{{{k}}}}}"), v)
        }
        interpolate_variables(&mut output);
        output
    }

    pub fn agent_session(&self) -> Option<&str> {
        self.config.agent_session.as_deref()
    }

    pub fn variables(&self) -> &AgentVariables {
        match &self.session_variables {
            Some(variables) => variables,
            None => &self.shared_variables,
        }
    }

    pub fn variable_envs(&self) -> HashMap<String, String> {
        self.variables()
            .iter()
            .map(|(k, v)| {
                (
                    format!("LLM_AGENT_VAR_{}", normalize_env_name(k)),
                    v.clone(),
                )
            })
            .collect()
    }

    pub fn shared_variables(&self) -> &AgentVariables {
        &self.shared_variables
    }

    pub fn set_shared_variables(&mut self, shared_variables: AgentVariables) {
        self.shared_variables = shared_variables;
    }

    pub fn set_session_variables(&mut self, session_variables: AgentVariables) {
        self.session_variables = Some(session_variables);
    }

    pub fn defined_variables(&self) -> &[AgentVariable] {
        &self.config.variables
    }

    pub fn exit_session(&mut self) {
        self.session_variables = None;
        self.session_dynamic_instructions = None;
    }

    pub fn is_dynamic_instructions(&self) -> bool {
        self.config.dynamic_instructions
    }

    pub fn update_shared_dynamic_instructions(&mut self, force: bool) -> Result<()> {
        if self.is_dynamic_instructions() && (force || self.shared_dynamic_instructions.is_none()) {
            self.shared_dynamic_instructions = Some(self.run_instructions_fn()?);
        }
        Ok(())
    }

    pub fn update_session_dynamic_instructions(&mut self, value: Option<String>) -> Result<()> {
        if self.is_dynamic_instructions() {
            let value = match value {
                Some(v) => v,
                None => self.run_instructions_fn()?,
            };
            self.session_dynamic_instructions = Some(value);
        }
        Ok(())
    }

    fn run_instructions_fn(&self) -> Result<String> {
        let value = run_llm_function(
            self.name().to_string(),
            vec!["_instructions".into(), "{}".into()],
            self.variable_envs(),
        )?;
        match value {
            Some(v) => Ok(v),
            _ => bail!("No return value from '_instructions' function"),
        }
    }
}

impl RoleLike for Agent {
    fn to_role(&self) -> Role {
        let prompt = self.interpolated_instructions();
        let mut role = Role::new("", &prompt);
        role.sync(self);
        role
    }

    fn model(&self) -> &Model {
        &self.model
    }

    fn temperature(&self) -> Option<f64> {
        self.config.temperature
    }

    fn top_p(&self) -> Option<f64> {
        self.config.top_p
    }

    fn use_tools(&self) -> Option<String> {
        self.config.global_tools.clone().join(",").into()
    }

    fn use_mcp_servers(&self) -> Option<String> {
        self.config.mcp_servers.clone().join(",").into()
    }

    fn set_model(&mut self, model: Model) {
        self.config.model_id = Some(model.id());
        self.model = model;
    }

    fn set_temperature(&mut self, value: Option<f64>) {
        self.config.temperature = value;
    }

    fn set_top_p(&mut self, value: Option<f64>) {
        self.config.top_p = value;
    }

    fn set_use_tools(&mut self, value: Option<String>) {
        match value {
            Some(tools) => {
                let tools = tools
                    .split(',')
                    .map(|v| v.trim().to_string())
                    .filter(|v| !v.is_empty())
                    .collect::<Vec<_>>();
                self.config.global_tools = tools;
            }
            None => {
                self.config.global_tools.clear();
            }
        }
    }

    fn set_use_mcp_servers(&mut self, value: Option<String>) {
        match value {
            Some(servers) => {
                let servers = servers
                    .split(',')
                    .map(|v| v.trim().to_string())
                    .filter(|v| !v.is_empty())
                    .collect::<Vec<_>>();
                self.config.mcp_servers = servers;
            }
            None => {
                self.config.mcp_servers.clear();
            }
        }
    }
}

#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct AgentConfig {
    pub name: String,
    #[serde(rename(serialize = "model", deserialize = "model"))]
    pub model_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub top_p: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub agent_session: Option<String>,
    #[serde(default)]
    pub description: String,
    #[serde(default)]
    pub version: String,
    #[serde(default)]
    pub mcp_servers: Vec<String>,
    #[serde(default)]
    pub global_tools: Vec<String>,
    #[serde(default)]
    pub instructions: String,
    #[serde(default)]
    pub dynamic_instructions: bool,
    #[serde(default)]
    pub variables: Vec<AgentVariable>,
    #[serde(default)]
    pub conversation_starters: Vec<String>,
    #[serde(default)]
    pub documents: Vec<String>,
}

impl AgentConfig {
    pub fn load(path: &Path) -> Result<Self> {
        let contents = read_to_string(path)
            .with_context(|| format!("Failed to read agent config file at '{}'", path.display()))?;
        let agent_config: Self = serde_yaml::from_str(&contents)
            .with_context(|| format!("Failed to load agent config at '{}'", path.display()))?;

        Ok(agent_config)
    }

    fn load_envs(&mut self, config: &Config) {
        let name = &self.name;
        let with_prefix = |v: &str| normalize_env_name(&format!("{name}_{v}"));

        if self.agent_session.is_none() {
            self.agent_session = config.agent_session.clone();
        }

        if let Some(v) = read_env_value::<String>(&with_prefix("model")) {
            self.model_id = v;
        }
        if let Some(v) = read_env_value::<f64>(&with_prefix("temperature")) {
            self.temperature = v;
        }
        if let Some(v) = read_env_value::<f64>(&with_prefix("top_p")) {
            self.top_p = v;
        }
        if let Some(v) = read_env_value::<String>(&with_prefix("agent_session")) {
            self.agent_session = v;
        }
        if let Ok(v) = env::var(with_prefix("variables")) {
            if let Ok(v) = serde_json::from_str(&v) {
                self.variables = v;
            }
        }
    }

    fn banner(&self) -> String {
        let AgentConfig {
            name,
            description,
            version,
            conversation_starters,
            ..
        } = self;
        let starters = if conversation_starters.is_empty() {
            String::new()
        } else {
            let starters = conversation_starters
                .iter()
                .map(|v| format!("- {v}"))
                .collect::<Vec<_>>()
                .join("\n");
            format!(
                r#"

## Conversation Starters
{starters}"#
            )
        };
        format!(
            r#"# {name} {version}
{description}{starters}"#
        )
    }

    fn replace_tools_placeholder(&mut self, functions: &Functions) {
        let tools_placeholder: &str = "{{__tools__}}";
        if self.instructions.contains(tools_placeholder) {
            let tools = functions
                .declarations()
                .iter()
                .enumerate()
                .map(|(i, v)| {
                    let description = match v.description.split_once('\n') {
                        Some((v, _)) => v,
                        None => &v.description,
                    };
                    format!("{}. {}: {description}", i + 1, v.name)
                })
                .collect::<Vec<String>>()
                .join("\n");
            self.instructions = self.instructions.replace(tools_placeholder, &tools);
        }
    }
}

#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct AgentVariable {
    pub name: String,
    pub description: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub default: Option<String>,
    #[serde(skip_deserializing, default)]
    pub value: String,
}

pub fn list_agents() -> Vec<String> {
    let agents_data_dir = Config::agents_data_dir();
    if !agents_data_dir.exists() {
        return vec![];
    }

    let mut agents = Vec::new();
    if let Ok(entries) = read_dir(agents_data_dir) {
        for entry in entries.flatten() {
            if entry.path().is_dir() {
                if let Some(name) = entry.file_name().to_str() {
                    agents.push(name.to_string());
                }
            }
        }
    }

    agents
}

pub fn complete_agent_variables(agent_name: &str) -> Vec<(String, Option<String>)> {
    let config_path = Config::agent_config_file(agent_name);
    if !config_path.exists() {
        return vec![];
    }
    let Ok(config) = AgentConfig::load(&config_path) else {
        return vec![];
    };
    config
        .variables
        .iter()
        .map(|v| {
            let description = match &v.default {
                Some(default) => format!("{} [default: {default}]", v.description),
                None => v.description.clone(),
            };
            (format!("{}=", v.name), Some(description))
        })
        .collect()
}