//! Dynamic Schema Tool Generation — Phase-83 T83.6
//!
//! Inspects available collections and federation stores at runtime,
//! producing a structured schema description for the LLM.
//!
//! This allows the LLM to understand:
//! - Which collections exist and their vector dimensions
//! - Which federation stores are registered and their types
//! - What tools are available and how to use them
//!
//! The generated schema can be:
//! 1. Injected into the system prompt as context
//! 2. Returned as a tool result (via `db_list_sources`)
//! 3. Used to dynamically construct `ToolDefinition`s with enum constraints

use std::sync::Arc;

use async_trait::async_trait;
use serde::{Deserialize, Serialize};

use crate::federation::Federation;
use crate::storage::collection::CollectionStore;
use crate::llm::error::{LlmError, LlmResult};
use crate::llm::tool_executor::{ToolExecutor, ToolMetadata, ToolResult, ToolRegistry};

// ============================================================================
// Schema types
// ============================================================================

/// Description of a single collection.
#[derive(Debug, Clone, Serialize)]
pub struct CollectionSchema {
    /// Collection name.
    pub name: String,
    /// Vector dimension.
    pub dimension: usize,
    /// Number of documents in the collection.
    pub document_count: usize,
    /// Number of vectors in the collection.
    pub vector_count: usize,
    /// Whether full-text search is available.
    pub fts_available: bool,
    /// Optional description from config.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
}

/// Description of a federation store.
#[derive(Debug, Clone, Serialize)]
pub struct StoreSchema {
    /// Store alias.
    pub alias: String,
    /// Store type (graph, document, vector).
    pub store_type: String,
}

/// Complete schema of all available data sources.
#[derive(Debug, Clone, Serialize)]
pub struct DataSourceSchema {
    /// Available collections.
    pub collections: Vec<CollectionSchema>,
    /// Available federation stores.
    pub stores: Vec<StoreSchema>,
}

// ============================================================================
// SchemaIntrospector
// ============================================================================

/// Introspects the database to discover available data sources and their schemas.
pub struct SchemaIntrospector {
    collections: Arc<CollectionStore>,
    federation: Arc<Federation>,
}

impl SchemaIntrospector {
    pub fn new(
        collections: Arc<CollectionStore>,
        federation: Arc<Federation>,
    ) -> Self {
        Self {
            collections,
            federation,
        }
    }

    /// Discover all available data sources and return their schemas.
    pub fn introspect(&self) -> DataSourceSchema {
        let collections = self.introspect_collections();
        let stores = self.introspect_stores();
        DataSourceSchema {
            collections,
            stores,
        }
    }

    /// Discover available collections and their schemas.
    pub fn introspect_collections(&self) -> Vec<CollectionSchema> {
        let names = self.collections.list_collections();
        names
            .into_iter()
            .filter_map(|name| {
                let collection = self.collections.get_collection(&name).ok()??;
                Some(CollectionSchema {
                    name: collection.name.clone(),
                    dimension: collection.config.dimension,
                    document_count: collection.document_count(),
                    vector_count: collection.vector_count(),
                    fts_available: collection.docs.supports_fts(),
                    description: collection.config.description.clone(),
                })
            })
            .collect()
    }

    /// Discover available federation stores.
    pub fn introspect_stores(&self) -> Vec<StoreSchema> {
        match self.federation.list_stores() {
            Ok(stores) => stores
                .into_iter()
                .map(|(alias, store_type)| StoreSchema {
                    alias,
                    store_type: store_type.to_string(),
                })
                .collect(),
            Err(_) => Vec::new(),
        }
    }

    /// Generate a system prompt snippet describing all available data sources.
    ///
    /// This text can be prepended to the LLM's system prompt to give it
    /// awareness of the database schema.
    pub fn generate_system_prompt_snippet(&self) -> String {
        let schema = self.introspect();
        let mut lines = Vec::new();

        lines.push("## Available Data Sources\n".to_string());

        if !schema.collections.is_empty() {
            lines.push("### Collections (for db_vector_search, db_fts_search, db_doc_crud)".to_string());
            for c in &schema.collections {
                let fts_tag = if c.fts_available { ", FTS ✓" } else { "" };
                let desc = c.description.as_deref().unwrap_or("");
                let desc_suffix = if desc.is_empty() {
                    String::new()
                } else {
                    format!(" — {}", desc)
                };
                lines.push(format!(
                    "- **{}**: {} docs, {} vectors, dim={}{}{}", 
                    c.name, c.document_count, c.vector_count, c.dimension, fts_tag, desc_suffix
                ));
            }
            lines.push(String::new());
        }

        if !schema.stores.is_empty() {
            lines.push("### Federation Stores (for db_sql_query)".to_string());
            for s in &schema.stores {
                lines.push(format!("- **{}**: type={}", s.alias, s.store_type));
            }
            lines.push(String::new());
        }

        if schema.collections.is_empty() && schema.stores.is_empty() {
            lines.push("No data sources are currently registered.".to_string());
        }

        lines.join("\n")
    }

    /// Update tool definitions in the registry with dynamic enum constraints.
    ///
    /// For example, the `db_vector_search` tool's `collection` parameter
    /// gets its `enum` values populated with the actual collection names.
    pub fn update_tool_definitions(&self, registry: &ToolRegistry) {
        let schema = self.introspect();

        // Store the schema as metadata on the registry
        // (tools can access this at execution time)
        if let Ok(schema_json) = serde_json::to_value(&schema) {
            registry.set_context("data_source_schema", schema_json);
        }
    }
}

// ============================================================================
// db_list_sources tool
// ============================================================================

/// Built-in tool: `db_list_sources`
///
/// Allows the LLM to discover available data sources at runtime.
/// Returns collection names, store aliases, and their capabilities.
pub struct ListSourcesTool {
    collections: Arc<CollectionStore>,
    federation: Arc<Federation>,
}

impl ListSourcesTool {
    pub fn new(
        collections: Arc<CollectionStore>,
        federation: Arc<Federation>,
    ) -> Self {
        Self {
            collections,
            federation,
        }
    }
}

/// Arguments for the list sources tool.
#[derive(Debug, Deserialize)]
struct ListSourcesArgs {
    /// Optional: filter by source type ("collections", "stores", or "all").
    #[serde(default = "default_source_type")]
    source_type: String,
}

fn default_source_type() -> String {
    "all".to_string()
}

#[async_trait]
impl ToolExecutor for ListSourcesTool {
    fn metadata(&self) -> ToolMetadata {
        ToolMetadata::new(
            "db_list_sources",
            "List all available data sources (collections and federation stores). \
             Use this tool first to discover what data is available before \
             querying with other tools. Returns collection names with their \
             capabilities (document count, vector dimension, FTS support) \
             and federation store aliases with their types.",
            serde_json::json!({
                "type": "object",
                "properties": {
                    "source_type": {
                        "type": "string",
                        "enum": ["all", "collections", "stores"],
                        "description": "Filter by source type (default: 'all')",
                        "default": "all"
                    }
                }
            }),
        )
        .with_category("discovery")
    }

    async fn execute(&self, arguments: &str) -> LlmResult<ToolResult> {
        let args: ListSourcesArgs = serde_json::from_str(arguments)
            .unwrap_or(ListSourcesArgs {
                source_type: "all".to_string(),
            });

        let introspector = SchemaIntrospector::new(
            self.collections.clone(),
            self.federation.clone(),
        );

        match args.source_type.as_str() {
            "collections" => {
                let collections = introspector.introspect_collections();
                Ok(ToolResult::ok_json(&serde_json::json!({
                    "source_type": "collections",
                    "count": collections.len(),
                    "collections": collections,
                })))
            }
            "stores" => {
                let stores = introspector.introspect_stores();
                Ok(ToolResult::ok_json(&serde_json::json!({
                    "source_type": "stores",
                    "count": stores.len(),
                    "stores": stores,
                })))
            }
            _ => {
                let schema = introspector.introspect();
                Ok(ToolResult::ok_json(&serde_json::json!({
                    "source_type": "all",
                    "collection_count": schema.collections.len(),
                    "store_count": schema.stores.len(),
                    "collections": schema.collections,
                    "stores": schema.stores,
                })))
            }
        }
    }
}

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

    #[test]
    fn test_list_sources_metadata() {
        let collections = Arc::new(
            CollectionStore::open_or_create(std::env::temp_dir().join("test_ls_meta")).unwrap(),
        );
        let federation = Arc::new(Federation::new(FederationConfig::default()));
        let tool = ListSourcesTool::new(collections, federation);
        let meta = tool.metadata();
        assert_eq!(meta.name, "db_list_sources");
        assert_eq!(meta.category, "discovery");
        assert!(meta.is_read_only);
    }

    #[test]
    fn test_schema_introspector_empty() {
        let collections = Arc::new(
            CollectionStore::open_or_create(std::env::temp_dir().join("test_si_empty")).unwrap(),
        );
        let federation = Arc::new(Federation::new(FederationConfig::default()));
        let introspector = SchemaIntrospector::new(collections, federation);
        let schema = introspector.introspect();
        // At minimum, a default collection exists
        assert!(!schema.collections.is_empty() || schema.stores.is_empty());
    }

    #[test]
    fn test_generate_system_prompt_snippet() {
        let collections = Arc::new(
            CollectionStore::open_or_create(std::env::temp_dir().join("test_prompt_snip")).unwrap(),
        );
        let federation = Arc::new(Federation::new(FederationConfig::default()));
        let introspector = SchemaIntrospector::new(collections, federation);
        let snippet = introspector.generate_system_prompt_snippet();
        assert!(snippet.contains("Available Data Sources"));
    }

    #[tokio::test]
    async fn test_list_sources_execute_all() {
        let collections = Arc::new(
            CollectionStore::open_or_create(std::env::temp_dir().join("test_ls_exec")).unwrap(),
        );
        let federation = Arc::new(Federation::new(FederationConfig::default()));
        let tool = ListSourcesTool::new(collections, federation);
        let result = tool.execute(r#"{"source_type": "all"}"#).await.unwrap();
        assert!(result.success);
        let output: serde_json::Value = serde_json::from_str(&result.output).unwrap();
        assert_eq!(output["source_type"], "all");
    }
}
