use std::collections::HashMap;
use std::sync::Arc;

use anyhow::{Context, Result};
use rmcp::model::{CallToolRequestParam, RawContent};
use rmcp::ServiceExt;
use rmcp::transport::child_process::TokioChildProcess;
use serde_json::Value;
use tokio::sync::Mutex;
use tracing::warn;

use super::types::{KBServerConfig, KBSource, KnowledgeEntry};

/// Manages connections to external knowledge base MCP servers.
///
/// Mirrors `McpConnectionManager` from the sync module: maintains a registry
/// of `KBServerConfig` entries and spawns a fresh child process per call.
pub struct KBConnector {
    configs: Arc<Mutex<HashMap<String, KBServerConfig>>>,
}

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

impl KBConnector {
    pub fn new() -> Self {
        Self {
            configs: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Register a knowledge base MCP server configuration.
    pub async fn register(&self, config: KBServerConfig) {
        let mut configs = self.configs.lock().await;
        configs.insert(config.name.clone(), config);
    }

    /// Check if a source has a registered MCP server.
    pub async fn is_registered(&self, name: &str) -> bool {
        let configs = self.configs.lock().await;
        configs.contains_key(name)
    }

    /// Return names of all registered KB sources.
    pub async fn registered_sources(&self) -> Vec<String> {
        let configs = self.configs.lock().await;
        configs.keys().cloned().collect()
    }

    /// List available collections on a KB server.
    pub async fn list_collections(&self, name: &str) -> Result<Vec<String>> {
        let config = self.get_config(name).await?;
        let tool_name = Self::list_collections_tool(&config.source);

        let result = self.call_kb_tool(&config, &tool_name, serde_json::json!({})).await?;

        // Parse collection names from the response.
        // ChromaDB returns an array of collection objects; others may vary.
        let collections = if let Some(arr) = result.as_array() {
            arr.iter()
                .filter_map(|v| {
                    v.as_str()
                        .map(String::from)
                        .or_else(|| v.get("name").and_then(|n| n.as_str()).map(String::from))
                })
                .collect()
        } else if let Some(obj) = result.as_object() {
            // Handle { "collections": [...] } format
            obj.get("collections")
                .and_then(|c| c.as_array())
                .map(|arr| {
                    arr.iter()
                        .filter_map(|v| {
                            v.as_str()
                                .map(String::from)
                                .or_else(|| v.get("name").and_then(|n| n.as_str()).map(String::from))
                        })
                        .collect()
                })
                .unwrap_or_default()
        } else {
            Vec::new()
        };

        Ok(collections)
    }

    /// Query a knowledge base for documents matching a query.
    pub async fn query(
        &self,
        name: &str,
        query: &str,
        max_results: usize,
        collection: &str,
    ) -> Result<Vec<KnowledgeEntry>> {
        let config = self.get_config(name).await?;
        let tool_name = Self::query_tool(&config.source);

        let args = Self::build_query_args(&config.source, query, max_results, collection);
        let result = self.call_kb_tool(&config, &tool_name, args).await?;

        Self::parse_query_result(&result, name, collection)
    }

    /// Ingest a document into a knowledge base collection.
    pub async fn ingest(
        &self,
        name: &str,
        title: &str,
        content: &str,
        metadata: &HashMap<String, String>,
        collection: &str,
    ) -> Result<Value> {
        let config = self.get_config(name).await?;
        let tool_name = Self::ingest_tool(&config.source);

        let id = format!("furnace-{}", chrono::Utc::now().timestamp_millis());
        let args = Self::build_ingest_args(&config.source, &id, title, content, metadata, collection);

        self.call_kb_tool(&config, &tool_name, args).await
    }

    /// Resolve the effective collection(s) based on config and params.
    pub async fn resolve_collections(
        &self,
        name: &str,
        collection: Option<&str>,
        collection_pattern: Option<&str>,
    ) -> Result<Vec<String>> {
        if let Some(c) = collection {
            return Ok(vec![c.to_string()]);
        }

        if let Some(pattern) = collection_pattern {
            let all_collections = self.list_collections(name).await?;
            let matched = Self::glob_filter(&all_collections, pattern);
            if matched.is_empty() {
                warn!(pattern = %pattern, "No collections matched pattern");
            }
            return Ok(matched);
        }

        // Fall back to config default.
        let config = self.get_config(name).await?;
        if let Some(ref pat) = config.collection_pattern {
            let all_collections = self.list_collections(name).await?;
            return Ok(Self::glob_filter(&all_collections, pat));
        }
        if let Some(ref c) = config.collection {
            return Ok(vec![c.clone()]);
        }

        Ok(vec!["default".to_string()])
    }

    // -----------------------------------------------------------------------
    // Internal helpers
    // -----------------------------------------------------------------------

    async fn get_config(&self, name: &str) -> Result<KBServerConfig> {
        let configs = self.configs.lock().await;
        configs
            .get(name)
            .cloned()
            .context(format!("No KB server configured for '{}'", name))
    }

    async fn call_kb_tool(
        &self,
        config: &KBServerConfig,
        tool_name: &str,
        arguments: Value,
    ) -> Result<Value> {
        let mut cmd = tokio::process::Command::new(&config.command);
        for arg in &config.args {
            cmd.arg(arg);
        }
        for (key, val) in &config.env_vars {
            cmd.env(key, val);
        }

        let transport = TokioChildProcess::new(&mut cmd)
            .context(format!("Failed to spawn KB MCP server for {}", config.name))?;
        let client = ().serve(transport)
            .await
            .map_err(|e| anyhow::anyhow!("Failed to connect to {} KB server: {}", config.name, e))?;

        let result = client
            .call_tool(CallToolRequestParam {
                name: tool_name.to_string().into(),
                arguments: Some(arguments.as_object().cloned().unwrap_or_default()),
            })
            .await
            .map_err(|e| {
                anyhow::anyhow!("KB tool '{}' failed on {}: {}", tool_name, config.name, e)
            })?;

        let content_text = result
            .content
            .iter()
            .filter_map(|c| {
                if let RawContent::Text(t) = &c.raw {
                    Some(t.text.clone())
                } else {
                    None
                }
            })
            .collect::<Vec<_>>()
            .join("\n");

        let value = serde_json::from_str(&content_text)
            .unwrap_or(Value::String(content_text));

        let _ = client.cancel().await;

        Ok(value)
    }

    /// Map KB source to the appropriate tool name for listing collections.
    fn list_collections_tool(source: &KBSource) -> String {
        match source {
            KBSource::ChromaDB => "chroma_list_collections".to_string(),
            _ => "list_collections".to_string(),
        }
    }

    /// Map KB source to the appropriate tool name for querying.
    fn query_tool(source: &KBSource) -> String {
        match source {
            KBSource::ChromaDB => "chroma_query_documents".to_string(),
            _ => "search".to_string(),
        }
    }

    /// Map KB source to the appropriate tool name for ingestion.
    fn ingest_tool(source: &KBSource) -> String {
        match source {
            KBSource::ChromaDB => "chroma_add_documents".to_string(),
            _ => "add".to_string(),
        }
    }

    /// Build query arguments for the specific KB source.
    fn build_query_args(source: &KBSource, query: &str, max_results: usize, collection: &str) -> Value {
        match source {
            KBSource::ChromaDB => serde_json::json!({
                "collection_name": collection,
                "query_texts": [query],
                "n_results": max_results,
            }),
            _ => serde_json::json!({
                "collection": collection,
                "query": query,
                "limit": max_results,
            }),
        }
    }

    /// Build ingest arguments for the specific KB source.
    fn build_ingest_args(
        source: &KBSource,
        id: &str,
        title: &str,
        content: &str,
        metadata: &HashMap<String, String>,
        collection: &str,
    ) -> Value {
        let mut meta = metadata.clone();
        meta.insert("title".into(), title.into());

        match source {
            KBSource::ChromaDB => serde_json::json!({
                "collection_name": collection,
                "documents": [content],
                "ids": [id],
                "metadatas": [meta],
            }),
            _ => serde_json::json!({
                "collection": collection,
                "id": id,
                "content": content,
                "metadata": meta,
            }),
        }
    }

    /// Parse a KB query response into KnowledgeEntry items.
    fn parse_query_result(result: &Value, source: &str, collection: &str) -> Result<Vec<KnowledgeEntry>> {
        let now = chrono::Utc::now().timestamp();
        let mut entries = Vec::new();

        // ChromaDB format: { documents: [[...]], ids: [[...]], distances: [[...]], metadatas: [[...]] }
        if let (Some(docs), Some(ids)) = (
            result.get("documents").and_then(|d| d.get(0)).and_then(|d| d.as_array()),
            result.get("ids").and_then(|d| d.get(0)).and_then(|d| d.as_array()),
        ) {
            let distances = result
                .get("distances")
                .and_then(|d| d.get(0))
                .and_then(|d| d.as_array());
            let metadatas = result
                .get("metadatas")
                .and_then(|d| d.get(0))
                .and_then(|d| d.as_array());

            for (i, (doc, id)) in docs.iter().zip(ids.iter()).enumerate() {
                let content = doc.as_str().unwrap_or("").to_string();
                let entry_id = id.as_str().unwrap_or("").to_string();

                // ChromaDB returns distances (lower = better); convert to similarity score.
                let score = distances
                    .and_then(|d| d.get(i))
                    .and_then(|v| v.as_f64())
                    .map(|dist| 1.0 / (1.0 + dist))
                    .unwrap_or(0.5);

                let metadata: HashMap<String, String> = metadatas
                    .and_then(|m| m.get(i))
                    .and_then(|m| m.as_object())
                    .map(|obj| {
                        obj.iter()
                            .map(|(k, v)| (k.clone(), v.as_str().unwrap_or("").to_string()))
                            .collect()
                    })
                    .unwrap_or_default();

                let title = metadata
                    .get("title")
                    .cloned()
                    .unwrap_or_else(|| format!("Entry {}", entry_id));

                entries.push(KnowledgeEntry {
                    id: entry_id,
                    source: source.to_string(),
                    collection: collection.to_string(),
                    title,
                    content,
                    metadata,
                    relevance_score: score,
                    retrieved_at: now,
                });
            }
        } else if let Some(results_arr) = result.get("results").and_then(|r| r.as_array()) {
            // Generic format: { results: [{ id, content, score, metadata }] }
            for item in results_arr {
                let entry_id = item.get("id").and_then(|v| v.as_str()).unwrap_or("").to_string();
                let content = item.get("content").and_then(|v| v.as_str()).unwrap_or("").to_string();
                let score = item.get("score").and_then(|v| v.as_f64()).unwrap_or(0.5);
                let title = item.get("title").and_then(|v| v.as_str()).unwrap_or("").to_string();
                let metadata: HashMap<String, String> = item
                    .get("metadata")
                    .and_then(|m| m.as_object())
                    .map(|obj| {
                        obj.iter()
                            .map(|(k, v)| (k.clone(), v.as_str().unwrap_or("").to_string()))
                            .collect()
                    })
                    .unwrap_or_default();

                entries.push(KnowledgeEntry {
                    id: entry_id,
                    source: source.to_string(),
                    collection: collection.to_string(),
                    title,
                    content,
                    metadata,
                    relevance_score: score,
                    retrieved_at: now,
                });
            }
        }

        Ok(entries)
    }

    /// Simple glob-style pattern matching for collection names.
    /// Supports `*` as a wildcard.
    fn glob_filter(names: &[String], pattern: &str) -> Vec<String> {
        let regex_pattern = format!(
            "^{}$",
            regex::escape(pattern).replace(r"\*", ".*")
        );
        let re = match regex::Regex::new(&regex_pattern) {
            Ok(r) => r,
            Err(_) => return Vec::new(),
        };

        names
            .iter()
            .filter(|name| re.is_match(name))
            .cloned()
            .collect()
    }
}

// ===========================================================================
// Tests
// ===========================================================================

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_register_and_check() {
        let connector = KBConnector::new();
        assert!(!connector.is_registered("test-kb").await);

        connector
            .register(KBServerConfig {
                source: KBSource::ChromaDB,
                name: "test-kb".into(),
                command: "echo".into(),
                args: vec![],
                env_vars: vec![],
                collection: Some("default".into()),
                collection_pattern: None,
            })
            .await;

        assert!(connector.is_registered("test-kb").await);
    }

    #[tokio::test]
    async fn test_registered_sources() {
        let connector = KBConnector::new();
        connector
            .register(KBServerConfig {
                source: KBSource::ChromaDB,
                name: "kb-a".into(),
                command: "echo".into(),
                args: vec![],
                env_vars: vec![],
                collection: None,
                collection_pattern: None,
            })
            .await;
        connector
            .register(KBServerConfig {
                source: KBSource::Pinecone,
                name: "kb-b".into(),
                command: "echo".into(),
                args: vec![],
                env_vars: vec![],
                collection: None,
                collection_pattern: None,
            })
            .await;

        let sources = connector.registered_sources().await;
        assert_eq!(sources.len(), 2);
        assert!(sources.contains(&"kb-a".to_string()));
        assert!(sources.contains(&"kb-b".to_string()));
    }

    #[test]
    fn test_glob_filter() {
        let names = vec![
            "requirements-auth".to_string(),
            "requirements-api".to_string(),
            "design-docs".to_string(),
            "requirements-ui".to_string(),
        ];

        let matched = KBConnector::glob_filter(&names, "requirements-*");
        assert_eq!(matched.len(), 3);
        assert!(matched.contains(&"requirements-auth".to_string()));
        assert!(matched.contains(&"requirements-api".to_string()));
        assert!(matched.contains(&"requirements-ui".to_string()));
        assert!(!matched.contains(&"design-docs".to_string()));
    }

    #[test]
    fn test_glob_filter_exact_match() {
        let names = vec!["design-docs".to_string(), "other".to_string()];
        let matched = KBConnector::glob_filter(&names, "design-docs");
        assert_eq!(matched, vec!["design-docs".to_string()]);
    }

    #[test]
    fn test_glob_filter_no_match() {
        let names = vec!["foo".to_string(), "bar".to_string()];
        let matched = KBConnector::glob_filter(&names, "baz-*");
        assert!(matched.is_empty());
    }

    #[test]
    fn test_parse_chromadb_result() {
        let result = serde_json::json!({
            "documents": [["doc content 1", "doc content 2"]],
            "ids": [["id-1", "id-2"]],
            "distances": [[0.1, 0.5]],
            "metadatas": [[{"title": "Auth Design"}, {"title": "API Spec"}]],
        });

        let entries = KBConnector::parse_query_result(&result, "ChromaDB", "design-docs").unwrap();
        assert_eq!(entries.len(), 2);
        assert_eq!(entries[0].id, "id-1");
        assert_eq!(entries[0].content, "doc content 1");
        assert_eq!(entries[0].title, "Auth Design");
        assert_eq!(entries[0].source, "ChromaDB");
        assert_eq!(entries[0].collection, "design-docs");
        // distance 0.1 → score = 1/(1+0.1) ≈ 0.909
        assert!(entries[0].relevance_score > 0.9);
        // distance 0.5 → score = 1/(1+0.5) ≈ 0.667
        assert!(entries[1].relevance_score > 0.6 && entries[1].relevance_score < 0.7);
    }

    #[test]
    fn test_parse_generic_result() {
        let result = serde_json::json!({
            "results": [
                {"id": "g-1", "content": "Generic content", "score": 0.8, "title": "Test", "metadata": {"key": "val"}},
            ]
        });

        let entries = KBConnector::parse_query_result(&result, "Custom", "my-collection").unwrap();
        assert_eq!(entries.len(), 1);
        assert_eq!(entries[0].id, "g-1");
        assert_eq!(entries[0].relevance_score, 0.8);
        assert_eq!(entries[0].metadata.get("key").unwrap(), "val");
    }

    #[test]
    fn test_parse_empty_result() {
        let result = serde_json::json!({});
        let entries = KBConnector::parse_query_result(&result, "KB", "col").unwrap();
        assert!(entries.is_empty());
    }

    #[tokio::test]
    async fn test_resolve_collections_explicit() {
        let connector = KBConnector::new();
        connector
            .register(KBServerConfig {
                source: KBSource::ChromaDB,
                name: "test".into(),
                command: "echo".into(),
                args: vec![],
                env_vars: vec![],
                collection: Some("default-col".into()),
                collection_pattern: None,
            })
            .await;

        // Explicit collection overrides config default.
        let cols = connector
            .resolve_collections("test", Some("explicit"), None)
            .await
            .unwrap();
        assert_eq!(cols, vec!["explicit"]);
    }

    #[tokio::test]
    async fn test_resolve_collections_config_default() {
        let connector = KBConnector::new();
        connector
            .register(KBServerConfig {
                source: KBSource::ChromaDB,
                name: "test".into(),
                command: "echo".into(),
                args: vec![],
                env_vars: vec![],
                collection: Some("my-default".into()),
                collection_pattern: None,
            })
            .await;

        let cols = connector
            .resolve_collections("test", None, None)
            .await
            .unwrap();
        assert_eq!(cols, vec!["my-default"]);
    }

    #[test]
    fn test_build_chromadb_query_args() {
        let args = KBConnector::build_query_args(&KBSource::ChromaDB, "test query", 5, "my-col");
        assert_eq!(args["collection_name"], "my-col");
        assert_eq!(args["n_results"], 5);
        assert_eq!(args["query_texts"][0], "test query");
    }

    #[test]
    fn test_build_generic_query_args() {
        let args = KBConnector::build_query_args(&KBSource::Pinecone, "test query", 10, "idx");
        assert_eq!(args["collection"], "idx");
        assert_eq!(args["limit"], 10);
        assert_eq!(args["query"], "test query");
    }
}
