use anyhow::{Context, Result, anyhow};
use arc_swap::ArcSwap;
use hashbrown::HashMap;
use rmcp::model::{
    CallToolRequestParams, CallToolResult, GetPromptRequestParams, InitializeRequestParams,
    InitializeResult, Prompt, ReadResourceRequestParams, Resource, Tool,
};
use serde_json::{Map, Value};
use std::ffi::OsString;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{Mutex, Semaphore};
use tracing::{Instrument, Span, warn};
use url::Url;

use super::{LATEST_PROTOCOL_VERSION, SUPPORTED_PROTOCOL_VERSIONS};

use crate::config::mcp::{McpAllowListConfig, McpProviderConfig, McpTransportConfig};
use vtcode_config::auth::McpOAuthService;
use vtcode_utility_tool_specs::parse_mcp_tool;

use super::{McpClient, RmcpClient};
use super::{
    McpElicitationHandler, McpPromptDetail, McpPromptInfo, McpResourceData, McpResourceInfo,
    McpToolInfo, TIMEZONE_ARGUMENT, build_headers, ensure_timezone_argument, schema_requires_field,
};

pub struct McpProvider {
    pub(super) name: String,
    #[expect(dead_code)]
    pub(super) protocol_version: String,
    client: ArcSwap<RmcpClient>,
    /// Stored config so we can reconnect after disconnection.
    config: McpProviderConfig,
    /// Stored elicitation handler for reconnection.
    elicitation_handler: Option<Arc<dyn McpElicitationHandler>>,
    pub(crate) semaphore: Arc<Semaphore>,
    caches: Mutex<ProviderCaches>,
    initialize_result: Mutex<Option<InitializeResult>>,
}

#[derive(Default)]
struct ProviderCaches {
    tools: Option<Arc<Vec<McpToolInfo>>>,
    resources: Option<Arc<Vec<McpResourceInfo>>>,
    prompts: Option<Arc<Vec<McpPromptInfo>>>,
}

impl McpProvider {
    pub(super) async fn connect(
        config: McpProviderConfig,
        elicitation_handler: Option<Arc<dyn McpElicitationHandler>>,
    ) -> Result<Self> {
        if config.name.trim().is_empty() {
            return Err(anyhow!("MCP provider name cannot be empty"));
        }

        let max_requests = std::cmp::max(1, config.max_concurrent_requests);

        let (client, protocol_version) = match &config.transport {
            McpTransportConfig::Stdio(stdio) => {
                let program = OsString::from(&stdio.command);
                let args: Vec<OsString> = stdio.args.iter().map(OsString::from).collect();
                let working_dir = stdio.working_directory.as_ref().map(PathBuf::from);
                let env: HashMap<OsString, OsString> = config
                    .env
                    .iter()
                    .map(|(key, value)| (OsString::from(key), OsString::from(value)))
                    .collect();
                let client = RmcpClient::new_stdio_client(
                    config.name.clone(),
                    program,
                    args,
                    working_dir,
                    Some(env),
                    elicitation_handler.clone(),
                )
                .await?;
                (client, LATEST_PROTOCOL_VERSION.to_string())
            }
            McpTransportConfig::Http(http) => {
                if !SUPPORTED_PROTOCOL_VERSIONS
                    .iter()
                    .any(|supported| supported == &http.protocol_version)
                {
                    return Err(anyhow!(
                        "MCP HTTP provider '{}' requested unsupported protocol version '{}'",
                        config.name,
                        http.protocol_version
                    ));
                }

                let bearer_token = if let Some(oauth) = http.oauth.as_ref() {
                    McpOAuthService::new()
                        .resolve_access_token(&config.name, oauth)
                        .await?
                        .ok_or_else(|| {
                            anyhow!(
                                "MCP HTTP provider '{}' requires OAuth login. Run `vtcode mcp login {}`.",
                                config.name,
                                config.name
                            )
                        })
                        .map(Some)?
                } else {
                    match http.api_key_env.as_ref() {
                        Some(var) => Some(std::env::var(var).with_context(|| {
                            format!("Missing MCP API key environment variable: {var}")
                        })?),
                        None => None,
                    }
                };

                let headers = build_headers(&http.http_headers, &http.env_http_headers);
                let client = RmcpClient::new_streamable_http_client(
                    config.name.clone(),
                    &http.endpoint,
                    bearer_token,
                    headers,
                    elicitation_handler.clone(),
                )
                .await?;
                (client, http.protocol_version.clone())
            }
        };

        Ok(Self {
            name: config.name.clone(),
            protocol_version,
            client: ArcSwap::from_pointee(client),
            config,
            elicitation_handler,
            semaphore: Arc::new(Semaphore::new(max_requests)),
            caches: Mutex::new(ProviderCaches::default()),
            initialize_result: Mutex::new(None),
        })
    }

    pub(super) fn invalidate_caches(&self) {
        if let Ok(mut caches) = self.caches.try_lock() {
            caches.tools = None;
            caches.resources = None;
            caches.prompts = None;
        }
    }

    pub(super) async fn initialize(
        &self,
        params: InitializeRequestParams,
        startup_timeout: Option<Duration>,
        tool_timeout: Option<Duration>,
        allowlist: &McpAllowListConfig,
    ) -> Result<()> {
        let client = self.client.load_full();
        let result = client.initialize(params, startup_timeout).await?;

        let protocol_version_str = result.protocol_version.to_string();
        if !SUPPORTED_PROTOCOL_VERSIONS
            .iter()
            .any(|supported| *supported == protocol_version_str)
        {
            return Err(anyhow!(
                "MCP server for '{}' negotiated unsupported protocol version '{}'",
                self.name,
                protocol_version_str
            ));
        }

        *self.initialize_result.lock().await = Some(result);
        let _ = self.refresh_tools(allowlist, tool_timeout).await;
        Ok(())
    }

    pub(super) async fn list_tools(
        &self,
        allowlist: &McpAllowListConfig,
        timeout: Option<Duration>,
    ) -> Result<Vec<McpToolInfo>> {
        Ok(self
            .list_tools_shared(allowlist, timeout)
            .await?
            .as_ref()
            .clone())
    }

    async fn list_tools_shared(
        &self,
        allowlist: &McpAllowListConfig,
        timeout: Option<Duration>,
    ) -> Result<Arc<Vec<McpToolInfo>>> {
        let mut caches = self.caches.lock().await;
        if self.client.load_full().take_tool_list_changed() {
            caches.tools = None;
        }

        if let Some(cache) = &caches.tools {
            return Ok(Arc::clone(cache));
        }
        drop(caches);

        self.refresh_tools_shared(allowlist, timeout).await
    }

    pub(super) async fn refresh_tools(
        &self,
        allowlist: &McpAllowListConfig,
        timeout: Option<Duration>,
    ) -> Result<Vec<McpToolInfo>> {
        Ok(self
            .refresh_tools_shared(allowlist, timeout)
            .await?
            .as_ref()
            .clone())
    }

    async fn refresh_tools_shared(
        &self,
        allowlist: &McpAllowListConfig,
        timeout: Option<Duration>,
    ) -> Result<Arc<Vec<McpToolInfo>>> {
        let client = self.client.load_full();
        let tools = client.list_all_tools(timeout).await?;
        let filtered = Arc::new(self.filter_tools(tools, allowlist));
        self.caches.lock().await.tools = Some(Arc::clone(&filtered));
        Ok(filtered)
    }

    pub(super) async fn has_tool(
        &self,
        tool_name: &str,
        allowlist: &McpAllowListConfig,
        timeout: Option<Duration>,
    ) -> Result<bool> {
        let tools = self.list_tools_shared(allowlist, timeout).await?;
        Ok(tools.iter().any(|tool| tool.name == tool_name))
    }

    pub(super) async fn call_tool(
        &self,
        tool_name: &str,
        args: &Value,
        timeout: Option<Duration>,
        allowlist: &McpAllowListConfig,
    ) -> Result<CallToolResult> {
        if !allowlist.is_tool_allowed(&self.name, tool_name) {
            return Err(anyhow!(
                "Tool '{}' is blocked by the MCP allow list for provider '{}'",
                tool_name,
                self.name
            ));
        }

        let _permit = self
            .semaphore
            .clone()
            .acquire_owned()
            .await
            .context("Failed to acquire MCP request slot")?;
        let mut arguments = McpClient::normalize_arguments(args);
        self.add_argument_defaults(tool_name, &mut arguments, allowlist, timeout)
            .await
            .with_context(|| {
                format!(
                    "failed to prepare arguments for MCP tool '{}' on provider '{}'",
                    tool_name, self.name
                )
            })?;
        let params = CallToolRequestParams::new(tool_name.to_string()).with_arguments(arguments);
        let client = self.client.load_full();
        async move { client.call_tool(params, timeout).await }
            .instrument(mcp_tool_call_span(
                &self.name,
                tool_name,
                &self.config.transport,
            ))
            .await
    }

    async fn add_argument_defaults(
        &self,
        tool_name: &str,
        arguments: &mut Map<String, Value>,
        allowlist: &McpAllowListConfig,
        timeout: Option<Duration>,
    ) -> Result<()> {
        let requires_timezone = self
            .tool_requires_field(tool_name, TIMEZONE_ARGUMENT, allowlist, timeout)
            .await?;
        ensure_timezone_argument(arguments, requires_timezone)?;
        Ok(())
    }

    async fn tool_requires_field(
        &self,
        tool_name: &str,
        field: &str,
        allowlist: &McpAllowListConfig,
        timeout: Option<Duration>,
    ) -> Result<bool> {
        if let Some(tools) = &self.caches.lock().await.tools
            && let Some(tool) = tools.iter().find(|tool| tool.name == tool_name)
        {
            return Ok(schema_requires_field(&tool.input_schema, field));
        }

        match self.refresh_tools_shared(allowlist, timeout).await {
            Ok(tools) => Ok(tools
                .iter()
                .find(|tool| tool.name == tool_name)
                .map(|tool| schema_requires_field(&tool.input_schema, field))
                .unwrap_or(false)),
            Err(err) => {
                warn!(
                    "Failed to refresh tools while inspecting schema for '{}' on provider '{}': {err}",
                    tool_name, self.name
                );
                Ok(false)
            }
        }
    }

    pub(super) async fn list_resources(
        &self,
        allowlist: &McpAllowListConfig,
        timeout: Option<Duration>,
    ) -> Result<Vec<McpResourceInfo>> {
        Ok(self
            .list_resources_shared(allowlist, timeout)
            .await?
            .as_ref()
            .clone())
    }

    async fn list_resources_shared(
        &self,
        allowlist: &McpAllowListConfig,
        timeout: Option<Duration>,
    ) -> Result<Arc<Vec<McpResourceInfo>>> {
        let mut caches = self.caches.lock().await;
        if self.client.load_full().take_resource_list_changed() {
            caches.resources = None;
        }

        if let Some(cache) = &caches.resources {
            return Ok(Arc::clone(cache));
        }
        drop(caches);

        self.refresh_resources_shared(allowlist, timeout).await
    }

    pub(super) async fn refresh_resources(
        &self,
        allowlist: &McpAllowListConfig,
        timeout: Option<Duration>,
    ) -> Result<Vec<McpResourceInfo>> {
        Ok(self
            .refresh_resources_shared(allowlist, timeout)
            .await?
            .as_ref()
            .clone())
    }

    async fn refresh_resources_shared(
        &self,
        allowlist: &McpAllowListConfig,
        timeout: Option<Duration>,
    ) -> Result<Arc<Vec<McpResourceInfo>>> {
        let client = self.client.load_full();
        let resources = client.list_all_resources(timeout).await?;
        let filtered = Arc::new(self.filter_resources(resources, allowlist));
        self.caches.lock().await.resources = Some(Arc::clone(&filtered));
        Ok(filtered)
    }

    pub(super) async fn has_resource(
        &self,
        uri: &str,
        allowlist: &McpAllowListConfig,
        timeout: Option<Duration>,
    ) -> Result<bool> {
        let resources = self.list_resources_shared(allowlist, timeout).await?;
        Ok(resources.iter().any(|resource| resource.uri == uri))
    }

    pub(super) async fn read_resource(
        &self,
        uri: &str,
        timeout: Option<Duration>,
        allowlist: &McpAllowListConfig,
    ) -> Result<McpResourceData> {
        if !allowlist.is_resource_allowed(&self.name, uri) {
            return Err(anyhow!(
                "Resource '{}' is blocked by the MCP allow list for provider '{}'",
                uri,
                self.name
            ));
        }

        let _permit = self
            .semaphore
            .clone()
            .acquire_owned()
            .await
            .context("Failed to acquire MCP request slot")?;
        let params = ReadResourceRequestParams::new(uri.to_string());
        let client = self.client.load_full();
        let result = client.read_resource(params, timeout).await?;
        Ok(McpResourceData {
            provider: self.name.clone(),
            uri: uri.to_string(),
            contents: result.contents,
            meta: Map::new(),
        })
    }

    pub(super) async fn list_prompts(
        &self,
        allowlist: &McpAllowListConfig,
        timeout: Option<Duration>,
    ) -> Result<Vec<McpPromptInfo>> {
        Ok(self
            .list_prompts_shared(allowlist, timeout)
            .await?
            .as_ref()
            .clone())
    }

    async fn list_prompts_shared(
        &self,
        allowlist: &McpAllowListConfig,
        timeout: Option<Duration>,
    ) -> Result<Arc<Vec<McpPromptInfo>>> {
        let mut caches = self.caches.lock().await;
        if self.client.load_full().take_prompt_list_changed() {
            caches.prompts = None;
        }

        if let Some(cache) = &caches.prompts {
            return Ok(Arc::clone(cache));
        }
        drop(caches);

        self.refresh_prompts_shared(allowlist, timeout).await
    }

    pub(super) async fn refresh_prompts(
        &self,
        allowlist: &McpAllowListConfig,
        timeout: Option<Duration>,
    ) -> Result<Vec<McpPromptInfo>> {
        Ok(self
            .refresh_prompts_shared(allowlist, timeout)
            .await?
            .as_ref()
            .clone())
    }

    async fn refresh_prompts_shared(
        &self,
        allowlist: &McpAllowListConfig,
        timeout: Option<Duration>,
    ) -> Result<Arc<Vec<McpPromptInfo>>> {
        let client = self.client.load_full();
        let prompts = client.list_all_prompts(timeout).await?;
        let filtered = Arc::new(self.filter_prompts(prompts, allowlist));
        self.caches.lock().await.prompts = Some(Arc::clone(&filtered));
        Ok(filtered)
    }

    pub(super) async fn has_prompt(
        &self,
        prompt_name: &str,
        allowlist: &McpAllowListConfig,
        timeout: Option<Duration>,
    ) -> Result<bool> {
        let prompts = self.list_prompts_shared(allowlist, timeout).await?;
        Ok(prompts.iter().any(|prompt| prompt.name == prompt_name))
    }

    pub(super) async fn get_prompt(
        &self,
        prompt_name: &str,
        arguments: HashMap<String, String>,
        timeout: Option<Duration>,
        allowlist: &McpAllowListConfig,
    ) -> Result<McpPromptDetail> {
        if !allowlist.is_prompt_allowed(&self.name, prompt_name) {
            return Err(anyhow!(
                "Prompt '{}' is blocked by the MCP allow list for provider '{}'",
                prompt_name,
                self.name
            ));
        }

        let _permit = self
            .semaphore
            .clone()
            .acquire_owned()
            .await
            .context("Failed to acquire MCP request slot")?;
        // Convert HashMap<String, String> to JsonObject (BTreeMap<String, Value>)
        let args_json: Map<String, Value> = arguments
            .into_iter()
            .map(|(k, v)| (k, Value::String(v)))
            .collect();

        let params = GetPromptRequestParams::new(prompt_name.to_string()).with_arguments(args_json);
        let client = self.client.load_full();
        let result = client.get_prompt(params, timeout).await?;
        Ok(McpPromptDetail {
            provider: self.name.clone(),
            name: prompt_name.to_string(),
            description: result.description,
            messages: result.messages,
            meta: Map::new(),
        })
    }

    pub(super) async fn cached_tools(&self) -> Option<Vec<McpToolInfo>> {
        self.caches
            .lock()
            .await
            .tools
            .as_ref()
            .map(|tools| tools.as_ref().clone())
    }

    pub(super) async fn cached_tools_or_refresh(
        &self,
        allowlist: &McpAllowListConfig,
        timeout: Option<Duration>,
    ) -> Result<Vec<McpToolInfo>> {
        if let Some(tools) = self.cached_tools().await {
            return Ok(tools);
        }

        self.refresh_tools(allowlist, timeout).await
    }

    pub(super) async fn shutdown(&self) -> Result<()> {
        let client = self.client.load_full();
        client.shutdown().await
    }

    /// Returns `true` when the underlying transport is still connected and responsive.
    pub(super) async fn is_healthy(&self) -> bool {
        let client = self.client.load_full();
        client.is_healthy().await
    }

    /// Attempt to re-establish the MCP connection using the stored configuration.
    ///
    /// This replaces the inner [`RmcpClient`] with a freshly connected one, then
    /// re-initialises the provider (tools/resources/prompts caches are invalidated).
    pub(super) async fn reconnect(
        &self,
        startup_timeout: Option<Duration>,
        tool_timeout: Option<Duration>,
        allowlist: &McpAllowListConfig,
    ) -> Result<()> {
        tracing::info!(provider = self.name.as_str(), "Attempting MCP reconnection");

        // Shut down the old client (best-effort).
        {
            let old = self.client.load_full();
            let _ = old.shutdown().await;
        }

        // Create a fresh client from the stored config.
        let new_provider =
            McpProvider::connect(self.config.clone(), self.elicitation_handler.clone())
                .await
                .with_context(|| format!("MCP reconnect failed for provider '{}'", self.name))?;

        // Swap inner client.
        {
            let new_client = new_provider.client.load_full();
            self.client.store(new_client);
        }

        // Invalidate all caches before re-initialisation.
        self.invalidate_caches();

        // Re-initialise (handshake + tool refresh).
        let init_params = InitializeRequestParams::new(
            rmcp::model::ClientCapabilities::default(),
            super::utils::build_client_implementation(),
        )
        .with_protocol_version(rmcp::model::ProtocolVersion::V_2024_11_05);
        self.initialize(init_params, startup_timeout, tool_timeout, allowlist)
            .await
            .with_context(|| {
                format!("MCP re-initialization failed for provider '{}'", self.name)
            })?;

        tracing::info!(provider = self.name.as_str(), "MCP reconnection successful");
        Ok(())
    }

    fn filter_tools(&self, tools: Vec<Tool>, allowlist: &McpAllowListConfig) -> Vec<McpToolInfo> {
        tools
            .into_iter()
            .filter(|tool| allowlist.is_tool_allowed(&self.name, &tool.name))
            .map(|tool| {
                let parsed = parse_mcp_tool(&tool);
                McpToolInfo {
                    description: parsed.description,
                    input_schema: parsed.input_schema,
                    provider: self.name.clone(),
                    name: parsed.name,
                }
            })
            .collect()
    }

    fn filter_resources(
        &self,
        resources: Vec<Resource>,
        allowlist: &McpAllowListConfig,
    ) -> Vec<McpResourceInfo> {
        resources
            .into_iter()
            .filter(|resource| allowlist.is_resource_allowed(&self.name, &resource.uri))
            .map(|resource| McpResourceInfo {
                provider: self.name.clone(),
                uri: resource.uri.clone(),
                name: resource.name.clone(),
                description: resource.description.clone(),
                mime_type: resource.mime_type.clone(),
                size: resource.size.map(|s| s as i64),
            })
            .collect()
    }

    fn filter_prompts(
        &self,
        prompts: Vec<Prompt>,
        allowlist: &McpAllowListConfig,
    ) -> Vec<McpPromptInfo> {
        prompts
            .into_iter()
            .filter(|prompt| allowlist.is_prompt_allowed(&self.name, &prompt.name))
            .map(|prompt| McpPromptInfo {
                provider: self.name.clone(),
                name: prompt.name.clone(),
                description: prompt.description.clone(),
                arguments: prompt.arguments.clone().unwrap_or_default(),
            })
            .collect()
    }
}

fn mcp_tool_call_span(
    provider_name: &str,
    tool_name: &str,
    transport: &McpTransportConfig,
) -> Span {
    let (transport_label, server_address, server_port) = match transport {
        McpTransportConfig::Stdio(_) => ("stdio", String::new(), 0_u16),
        McpTransportConfig::Http(http) => {
            let (server_address, server_port) = Url::parse(&http.endpoint)
                .ok()
                .and_then(|url| {
                    url.host_str().map(|host| {
                        (
                            host.to_string(),
                            url.port_or_known_default().unwrap_or_default(),
                        )
                    })
                })
                .unwrap_or_default();
            ("streamable_http", server_address, server_port)
        }
    };

    tracing::info_span!(
        "mcp.tools.call",
        provider = provider_name,
        tool = tool_name,
        rpc_system = "jsonrpc",
        rpc_method = "tools/call",
        transport = transport_label,
        server_address = server_address.as_str(),
        server_port = server_port,
    )
}

#[cfg(test)]
mod tests {
    use super::mcp_tool_call_span;
    use crate::config::mcp::{McpHttpServerConfig, McpStdioServerConfig, McpTransportConfig};
    use crate::utils::trace_writer::FlushableWriter;
    use std::fs;
    use tempfile::tempdir;
    use tracing_subscriber::{fmt::format::FmtSpan, prelude::*};

    #[test]
    fn mcp_tool_call_span_records_http_metadata() {
        let tempdir = tempdir().expect("tempdir");
        let log_file = tempdir.path().join("trace.log");
        let writer = FlushableWriter::open(&log_file).expect("trace writer");
        let writer_for_layer = writer.clone();
        let subscriber = tracing_subscriber::registry().with(
            tracing_subscriber::fmt::layer()
                .with_writer(move || writer_for_layer.clone())
                .with_span_events(FmtSpan::FULL)
                .with_ansi(false),
        );
        let _guard = tracing::subscriber::set_default(subscriber);

        {
            let span = mcp_tool_call_span(
                "calendar",
                "get_events",
                &McpTransportConfig::Http(McpHttpServerConfig {
                    endpoint: "https://example.com:8443/mcp".to_string(),
                    api_key_env: None,
                    oauth: None,
                    protocol_version: "2024-11-05".to_string(),
                    http_headers: Default::default(),
                    env_http_headers: Default::default(),
                }),
            );
            let _entered = span.enter();
        }

        writer.flush();
        let logs = fs::read_to_string(&log_file).expect("trace log");
        assert!(logs.contains("mcp.tools.call"));
        assert!(logs.contains("provider=\"calendar\""));
        assert!(logs.contains("tool=\"get_events\""));
        assert!(logs.contains("rpc_system=\"jsonrpc\""));
        assert!(logs.contains("rpc_method=\"tools/call\""));
        assert!(logs.contains("transport=\"streamable_http\""));
        assert!(logs.contains("server_address=\"example.com\""));
        assert!(logs.contains("server_port=8443"));
    }

    #[test]
    fn mcp_tool_call_span_defaults_stdio_transport() {
        let span = mcp_tool_call_span(
            "filesystem",
            "read_file",
            &McpTransportConfig::Stdio(McpStdioServerConfig {
                command: "rmcp-server".to_string(),
                args: Vec::new(),
                working_directory: None,
            }),
        );

        assert_eq!(span.metadata().expect("metadata").name(), "mcp.tools.call");
    }
}
