use crate::{AgentServer, AgentServerDelegate, load_proxy_env};
use acp_thread::AgentConnection;
use agent_client_protocol as acp;
use anyhow::{Context as _, Result};
use collections::HashSet;
use fs::Fs;
use gpui::{App, AppContext as _, SharedString, Task};
use project::agent_server_store::{AllAgentServersSettings, ExternalAgentServerName};
use settings::{SettingsStore, update_settings_file};
use std::{path::Path, rc::Rc, sync::Arc};
use ui::IconName;

/// A generic agent server implementation for custom user-defined agents
pub struct CustomAgentServer {
    name: SharedString,
}

impl CustomAgentServer {
    pub fn new(name: SharedString) -> Self {
        Self { name }
    }
}

impl AgentServer for CustomAgentServer {
    fn name(&self) -> SharedString {
        self.name.clone()
    }

    fn logo(&self) -> IconName {
        IconName::Terminal
    }

    fn default_mode(&self, cx: &App) -> Option<acp::SessionModeId> {
        let settings = cx.read_global(|settings: &SettingsStore, _| {
            settings
                .get::<AllAgentServersSettings>(None)
                .custom
                .get(self.name().as_ref())
                .cloned()
        });

        settings
            .as_ref()
            .and_then(|s| s.default_mode().map(acp::SessionModeId::new))
    }

    fn favorite_config_option_value_ids(
        &self,
        config_id: &acp::SessionConfigId,
        cx: &mut App,
    ) -> HashSet<acp::SessionConfigValueId> {
        let settings = cx.read_global(|settings: &SettingsStore, _| {
            settings
                .get::<AllAgentServersSettings>(None)
                .custom
                .get(self.name().as_ref())
                .cloned()
        });

        settings
            .as_ref()
            .and_then(|s| s.favorite_config_option_values(config_id.0.as_ref()))
            .map(|values| {
                values
                    .iter()
                    .cloned()
                    .map(acp::SessionConfigValueId::new)
                    .collect()
            })
            .unwrap_or_default()
    }

    fn toggle_favorite_config_option_value(
        &self,
        config_id: acp::SessionConfigId,
        value_id: acp::SessionConfigValueId,
        should_be_favorite: bool,
        fs: Arc<dyn Fs>,
        cx: &App,
    ) {
        let name = self.name();
        let config_id = config_id.to_string();
        let value_id = value_id.to_string();

        update_settings_file(fs, cx, move |settings, _| {
            let settings = settings
                .agent_servers
                .get_or_insert_default()
                .custom
                .entry(name.to_string())
                .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
                    default_model: None,
                    default_mode: None,
                    env: Default::default(),
                    favorite_models: Vec::new(),
                    default_config_options: Default::default(),
                    favorite_config_option_values: Default::default(),
                });

            match settings {
                settings::CustomAgentServerSettings::Custom {
                    favorite_config_option_values,
                    ..
                }
                | settings::CustomAgentServerSettings::Extension {
                    favorite_config_option_values,
                    ..
                }
                | settings::CustomAgentServerSettings::Registry {
                    favorite_config_option_values,
                    ..
                } => {
                    let entry = favorite_config_option_values
                        .entry(config_id.clone())
                        .or_insert_with(Vec::new);

                    if should_be_favorite {
                        if !entry.iter().any(|v| v == &value_id) {
                            entry.push(value_id.clone());
                        }
                    } else {
                        entry.retain(|v| v != &value_id);
                        if entry.is_empty() {
                            favorite_config_option_values.remove(&config_id);
                        }
                    }
                }
            }
        });
    }

    fn set_default_mode(&self, mode_id: Option<acp::SessionModeId>, fs: Arc<dyn Fs>, cx: &mut App) {
        let name = self.name();
        update_settings_file(fs, cx, move |settings, _| {
            let settings = settings
                .agent_servers
                .get_or_insert_default()
                .custom
                .entry(name.to_string())
                .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
                    default_model: None,
                    default_mode: None,
                    env: Default::default(),
                    favorite_models: Vec::new(),
                    default_config_options: Default::default(),
                    favorite_config_option_values: Default::default(),
                });

            match settings {
                settings::CustomAgentServerSettings::Custom { default_mode, .. }
                | settings::CustomAgentServerSettings::Extension { default_mode, .. }
                | settings::CustomAgentServerSettings::Registry { default_mode, .. } => {
                    *default_mode = mode_id.map(|m| m.to_string());
                }
            }
        });
    }

    fn default_model(&self, cx: &App) -> Option<acp::ModelId> {
        let settings = cx.read_global(|settings: &SettingsStore, _| {
            settings
                .get::<AllAgentServersSettings>(None)
                .custom
                .get(self.name().as_ref())
                .cloned()
        });

        settings
            .as_ref()
            .and_then(|s| s.default_model().map(acp::ModelId::new))
    }

    fn set_default_model(&self, model_id: Option<acp::ModelId>, fs: Arc<dyn Fs>, cx: &mut App) {
        let name = self.name();
        update_settings_file(fs, cx, move |settings, _| {
            let settings = settings
                .agent_servers
                .get_or_insert_default()
                .custom
                .entry(name.to_string())
                .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
                    default_model: None,
                    default_mode: None,
                    env: Default::default(),
                    favorite_models: Vec::new(),
                    default_config_options: Default::default(),
                    favorite_config_option_values: Default::default(),
                });

            match settings {
                settings::CustomAgentServerSettings::Custom { default_model, .. }
                | settings::CustomAgentServerSettings::Extension { default_model, .. }
                | settings::CustomAgentServerSettings::Registry { default_model, .. } => {
                    *default_model = model_id.map(|m| m.to_string());
                }
            }
        });
    }

    fn favorite_model_ids(&self, cx: &mut App) -> HashSet<acp::ModelId> {
        let settings = cx.read_global(|settings: &SettingsStore, _| {
            settings
                .get::<AllAgentServersSettings>(None)
                .custom
                .get(self.name().as_ref())
                .cloned()
        });

        settings
            .as_ref()
            .map(|s| {
                s.favorite_models()
                    .iter()
                    .map(|id| acp::ModelId::new(id.clone()))
                    .collect()
            })
            .unwrap_or_default()
    }

    fn toggle_favorite_model(
        &self,
        model_id: acp::ModelId,
        should_be_favorite: bool,
        fs: Arc<dyn Fs>,
        cx: &App,
    ) {
        let name = self.name();
        update_settings_file(fs, cx, move |settings, _| {
            let settings = settings
                .agent_servers
                .get_or_insert_default()
                .custom
                .entry(name.to_string())
                .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
                    default_model: None,
                    default_mode: None,
                    env: Default::default(),
                    favorite_models: Vec::new(),
                    default_config_options: Default::default(),
                    favorite_config_option_values: Default::default(),
                });

            let favorite_models = match settings {
                settings::CustomAgentServerSettings::Custom {
                    favorite_models, ..
                }
                | settings::CustomAgentServerSettings::Extension {
                    favorite_models, ..
                }
                | settings::CustomAgentServerSettings::Registry {
                    favorite_models, ..
                } => favorite_models,
            };

            let model_id_str = model_id.to_string();
            if should_be_favorite {
                if !favorite_models.contains(&model_id_str) {
                    favorite_models.push(model_id_str);
                }
            } else {
                favorite_models.retain(|id| id != &model_id_str);
            }
        });
    }

    fn default_config_option(&self, config_id: &str, cx: &App) -> Option<String> {
        let settings = cx.read_global(|settings: &SettingsStore, _| {
            settings
                .get::<AllAgentServersSettings>(None)
                .custom
                .get(self.name().as_ref())
                .cloned()
        });

        settings
            .as_ref()
            .and_then(|s| s.default_config_option(config_id).map(|s| s.to_string()))
    }

    fn set_default_config_option(
        &self,
        config_id: &str,
        value_id: Option<&str>,
        fs: Arc<dyn Fs>,
        cx: &mut App,
    ) {
        let name = self.name();
        let config_id = config_id.to_string();
        let value_id = value_id.map(|s| s.to_string());
        update_settings_file(fs, cx, move |settings, _| {
            let settings = settings
                .agent_servers
                .get_or_insert_default()
                .custom
                .entry(name.to_string())
                .or_insert_with(|| settings::CustomAgentServerSettings::Extension {
                    default_model: None,
                    default_mode: None,
                    env: Default::default(),
                    favorite_models: Vec::new(),
                    default_config_options: Default::default(),
                    favorite_config_option_values: Default::default(),
                });

            match settings {
                settings::CustomAgentServerSettings::Custom {
                    default_config_options,
                    ..
                }
                | settings::CustomAgentServerSettings::Extension {
                    default_config_options,
                    ..
                }
                | settings::CustomAgentServerSettings::Registry {
                    default_config_options,
                    ..
                } => {
                    if let Some(value) = value_id.clone() {
                        default_config_options.insert(config_id.clone(), value);
                    } else {
                        default_config_options.remove(&config_id);
                    }
                }
            }
        });
    }

    fn connect(
        &self,
        root_dir: Option<&Path>,
        delegate: AgentServerDelegate,
        cx: &mut App,
    ) -> Task<Result<(Rc<dyn AgentConnection>, Option<task::SpawnInTerminal>)>> {
        let name = self.name();
        let root_dir = root_dir.map(|root_dir| root_dir.to_string_lossy().into_owned());
        let is_remote = delegate.project.read(cx).is_via_remote_server();
        let default_mode = self.default_mode(cx);
        let default_model = self.default_model(cx);
        let (default_config_options, is_registry_agent) =
            cx.read_global(|settings: &SettingsStore, _| {
                let agent_settings = settings
                    .get::<AllAgentServersSettings>(None)
                    .custom
                    .get(self.name().as_ref());

                let is_registry = agent_settings
                    .map(|s| {
                        matches!(
                            s,
                            project::agent_server_store::CustomAgentServerSettings::Registry { .. }
                        )
                    })
                    .unwrap_or(false);

                let config_options = agent_settings
                    .map(|s| match s {
                        project::agent_server_store::CustomAgentServerSettings::Custom {
                            default_config_options,
                            ..
                        }
                        | project::agent_server_store::CustomAgentServerSettings::Extension {
                            default_config_options,
                            ..
                        }
                        | project::agent_server_store::CustomAgentServerSettings::Registry {
                            default_config_options,
                            ..
                        } => default_config_options.clone(),
                    })
                    .unwrap_or_default();

                (config_options, is_registry)
            });

        if is_registry_agent {
            if let Some(registry_store) = project::AgentRegistryStore::try_global(cx) {
                registry_store.update(cx, |store, cx| store.refresh_if_stale(cx));
            }
        }

        let store = delegate.store.downgrade();
        let extra_env = load_proxy_env(cx);
        cx.spawn(async move |cx| {
            let (command, root_dir, login) = store
                .update(cx, |store, cx| {
                    let agent = store
                        .get_external_agent(&ExternalAgentServerName(name.clone()))
                        .with_context(|| {
                            format!("Custom agent server `{}` is not registered", name)
                        })?;
                    anyhow::Ok(agent.get_command(
                        root_dir.as_deref(),
                        extra_env,
                        delegate.status_tx,
                        delegate.new_version_available,
                        &mut cx.to_async(),
                    ))
                })??
                .await?;
            let connection = crate::acp::connect(
                name,
                command,
                root_dir.as_ref(),
                default_mode,
                default_model,
                default_config_options,
                is_remote,
                cx,
            )
            .await?;
            Ok((connection, login))
        })
    }

    fn into_any(self: Rc<Self>) -> Rc<dyn std::any::Any> {
        self
    }
}