//! Prompt registry for MCP servers
//!
//! Provides registration and retrieval of MCP prompts

use std::sync::Arc;
use std::{future::Future, pin::Pin};

use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;

use crate::content::prompt::PromptMessage;
use crate::server::session::Session;
use crate::server::visibility::{ExecutionContext, VisibilityContext};

/// Prompt execution errors
#[derive(Debug, Error)]
pub enum PromptError {
    /// Prompt not found
    #[error("Prompt not found: {0}")]
    NotFound(String),

    /// Invalid arguments
    #[error("Invalid arguments: {0}")]
    InvalidArguments(String),

    /// Execution error
    #[error("Execution error: {0}")]
    Execution(String),

    /// Internal error
    #[error("Internal error: {0}")]
    Internal(String),
}

/// Prompt argument definition
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PromptArgument {
    /// Argument name
    pub name: String,

    /// Argument description
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,

    /// Whether the argument is required
    #[serde(skip_serializing_if = "Option::is_none")]
    pub required: Option<bool>,
}

/// Prompt metadata for listing
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PromptInfo {
    /// Prompt name
    pub name: String,

    /// Prompt description
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,

    /// Prompt arguments
    #[serde(skip_serializing_if = "Option::is_none")]
    pub arguments: Option<Vec<PromptArgument>>,
}

/// Result of getting a prompt
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GetPromptResult {
    /// Optional description of the prompt
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,

    /// Messages that make up the prompt
    pub messages: Vec<PromptMessage>,
}

/// Boxed future for prompt execution.
pub type PromptFuture<'a> =
    Pin<Box<dyn Future<Output = Result<GetPromptResult, PromptError>> + Send + 'a>>;

/// Prompt trait for implementing MCP prompts
pub trait Prompt: Send + Sync {
    /// Get prompt name
    fn name(&self) -> &str;

    /// Get prompt description
    fn description(&self) -> Option<&str> {
        None
    }

    /// Get prompt arguments
    fn arguments(&self) -> Option<Vec<PromptArgument>> {
        None
    }

    /// Check if this prompt should be visible in the given context
    ///
    /// Override this to implement contextual visibility. The default implementation
    /// always returns true (always visible).
    fn is_visible(&self, _ctx: &VisibilityContext) -> bool {
        true
    }

    /// Get the prompt with execution context
    ///
    /// The execution context provides access to:
    /// - `ctx.params`: Arguments passed to the prompt
    /// - `ctx.session`: Current session (roles, state, client info)
    /// - `ctx.environment`: Optional environment state
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// async fn get(&self, ctx: ExecutionContext<'_>) -> Result<GetPromptResult, PromptError> {
    ///     let name = ctx.params.get("name").and_then(|v| v.as_str()).unwrap_or("World");
    ///
    ///     // Customize based on session
    ///     let greeting = if ctx.is_vip() {
    ///         format!("Welcome, honored {}!", name)
    ///     } else {
    ///         format!("Hello, {}!", name)
    ///     };
    ///
    ///     Ok(GetPromptResult { ... })
    /// }
    /// ```
    fn get<'a>(&'a self, ctx: ExecutionContext<'a>) -> PromptFuture<'a>;

    /// Create a user message with text content
    fn user(&self, text: &str) -> PromptMessage {
        PromptMessage::user(text)
    }

    /// Create an assistant message with text content
    fn assistant(&self, text: &str) -> PromptMessage {
        PromptMessage::assistant(text)
    }

    /// Create a user message with image content
    fn user_image(&self, data: &str, mime_type: &str) -> PromptMessage {
        PromptMessage::user_image(data, mime_type)
    }

    /// Create an assistant message with image content
    fn assistant_image(&self, data: &str, mime_type: &str) -> PromptMessage {
        PromptMessage::assistant_image(data, mime_type)
    }
}

/// Prompt manager for managing available prompts
#[derive(Clone)]
pub struct PromptManager {
    prompts: Arc<DashMap<String, Arc<dyn Prompt>>>,
}

impl PromptManager {
    /// Create new prompt manager
    pub fn new() -> Self {
        Self {
            prompts: Arc::new(DashMap::new()),
        }
    }

    /// Register a prompt
    pub fn register<P: Prompt + 'static>(&self, prompt: P) {
        let name = prompt.name().to_string();
        self.prompts.insert(name, Arc::new(prompt));
    }

    /// Register a boxed prompt
    pub fn register_boxed(&self, prompt: Arc<dyn Prompt>) {
        let name = prompt.name().to_string();
        self.prompts.insert(name, prompt);
    }

    /// Get a prompt by name
    pub fn get(&self, name: &str) -> Option<Arc<dyn Prompt>> {
        self.prompts.get(name).map(|p| Arc::clone(&p))
    }

    /// List all registered prompts
    pub fn list(&self) -> Vec<PromptInfo> {
        self.prompts
            .iter()
            .map(|entry| {
                let prompt = entry.value();
                PromptInfo {
                    name: prompt.name().to_string(),
                    description: prompt.description().map(|s| s.to_string()),
                    arguments: prompt.arguments(),
                }
            })
            .collect()
    }

    /// Get a prompt by name with parameters
    pub async fn call(
        &self,
        name: &str,
        params: Value,
        session: &Session,
        logger: &crate::logging::McpLogger,
    ) -> Result<GetPromptResult, PromptError> {
        let prompt = self
            .get(name)
            .ok_or_else(|| PromptError::NotFound(name.to_string()))?;

        let ctx = ExecutionContext::new(params, session, logger);
        prompt.get(ctx).await
    }

    // ==================== Session-Aware Methods ====================

    /// List prompts visible to a specific session
    ///
    /// Resolution order:
    /// 1. Session overrides (replace global prompts with session-specific implementations)
    /// 2. Session extras (additional prompts added to session)
    /// 3. Global prompts (filtered by hidden list and visibility predicate)
    pub fn list_for_session(
        &self,
        session: &Session,
        ctx: &VisibilityContext<'_>,
    ) -> Vec<PromptInfo> {
        let mut prompts = std::collections::HashMap::new();

        // 1. Add global prompts (filtered by hidden and visibility)
        for entry in self.prompts.iter() {
            let name = entry.key().clone();
            if !session.is_prompt_hidden(&name) {
                let prompt = entry.value();
                if prompt.is_visible(ctx) {
                    prompts.insert(
                        name,
                        PromptInfo {
                            name: prompt.name().to_string(),
                            description: prompt.description().map(|s| s.to_string()),
                            arguments: prompt.arguments(),
                        },
                    );
                }
            }
        }

        // 2. Add session extras
        for entry in session.prompt_extras().iter() {
            let name = entry.key().clone();
            let prompt = entry.value();
            if prompt.is_visible(ctx) {
                prompts.insert(
                    name,
                    PromptInfo {
                        name: prompt.name().to_string(),
                        description: prompt.description().map(|s| s.to_string()),
                        arguments: prompt.arguments(),
                    },
                );
            }
        }

        // 3. Apply session overrides
        for entry in session.prompt_overrides().iter() {
            let name = entry.key().clone();
            let prompt = entry.value();
            if prompt.is_visible(ctx) {
                prompts.insert(
                    name,
                    PromptInfo {
                        name: prompt.name().to_string(),
                        description: prompt.description().map(|s| s.to_string()),
                        arguments: prompt.arguments(),
                    },
                );
            }
        }

        prompts.into_values().collect()
    }

    /// Get a prompt with session context
    ///
    /// Resolution order:
    /// 1. Check session overrides
    /// 2. Check session extras
    /// 3. Check session hidden
    /// 4. Check visibility predicate
    /// 5. Fall back to global registry
    pub async fn call_for_session(
        &self,
        name: &str,
        params: Value,
        session: &Session,
        logger: &crate::logging::McpLogger,
        visibility_ctx: &VisibilityContext<'_>,
    ) -> Result<GetPromptResult, PromptError> {
        // Create execution context (reuse environment from visibility context)
        let exec_ctx = match visibility_ctx.environment {
            Some(env) => ExecutionContext::with_environment(params.clone(), session, logger, env),
            None => ExecutionContext::new(params.clone(), session, logger),
        };

        // 1. Check session override first
        if let Some(prompt) = session.get_prompt_override(name) {
            if !prompt.is_visible(visibility_ctx) {
                return Err(PromptError::NotFound(name.to_string()));
            }
            return prompt.get(exec_ctx).await;
        }

        // 2. Check session extras
        if let Some(prompt) = session.get_prompt_extra(name) {
            if !prompt.is_visible(visibility_ctx) {
                return Err(PromptError::NotFound(name.to_string()));
            }
            return prompt.get(exec_ctx).await;
        }

        // 3. Check if hidden in session
        if session.is_prompt_hidden(name) {
            return Err(PromptError::NotFound(name.to_string()));
        }

        // 4. Check global registry with visibility check
        let prompt = self
            .get(name)
            .ok_or_else(|| PromptError::NotFound(name.to_string()))?;

        if !prompt.is_visible(visibility_ctx) {
            return Err(PromptError::NotFound(name.to_string()));
        }

        // 5. Execute
        self.call(name, params, session, logger).await
    }

    /// Get number of registered prompts
    pub fn len(&self) -> usize {
        self.prompts.len()
    }

    /// Check if manager is empty
    pub fn is_empty(&self) -> bool {
        self.prompts.is_empty()
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::server::session::Session;

    // Example prompt for testing
    struct GreetingPrompt;

    impl Prompt for GreetingPrompt {
        fn name(&self) -> &str {
            "greeting"
        }

        fn description(&self) -> Option<&str> {
            Some("Generates a greeting prompt")
        }

        fn arguments(&self) -> Option<Vec<PromptArgument>> {
            Some(vec![PromptArgument {
                name: "name".to_string(),
                description: Some("Name to greet".to_string()),
                required: Some(true),
            }])
        }

        fn get<'a>(&'a self, ctx: ExecutionContext<'a>) -> PromptFuture<'a> {
            Box::pin(async move {
                let name = ctx
                    .params
                    .get("name")
                    .and_then(|v| v.as_str())
                    .ok_or_else(|| {
                        PromptError::InvalidArguments("Missing 'name' field".to_string())
                    })?;

                Ok(GetPromptResult {
                    description: Some("A friendly greeting".to_string()),
                    messages: vec![self.user(&format!("Hello, {}!", name))],
                })
            })
        }
    }

    #[test]
    fn test_manager_creation() {
        let manager = PromptManager::new();
        assert!(manager.is_empty());
    }

    #[test]
    fn test_prompt_registration() {
        let manager = PromptManager::new();
        manager.register(GreetingPrompt);

        assert_eq!(manager.len(), 1);
        assert!(!manager.is_empty());
    }

    #[test]
    fn test_get_prompt() {
        let manager = PromptManager::new();
        manager.register(GreetingPrompt);

        let prompt = manager.get("greeting");
        assert!(prompt.is_some());
        assert_eq!(prompt.unwrap().name(), "greeting");

        let missing = manager.get("nonexistent");
        assert!(missing.is_none());
    }

    #[test]
    fn test_list_prompts() {
        let manager = PromptManager::new();
        manager.register(GreetingPrompt);

        let prompts = manager.list();
        assert_eq!(prompts.len(), 1);
        assert_eq!(prompts[0].name, "greeting");
        assert_eq!(
            prompts[0].description,
            Some("Generates a greeting prompt".to_string())
        );
        assert!(prompts[0].arguments.is_some());

        let args = prompts[0].arguments.as_ref().unwrap();
        assert_eq!(args.len(), 1);
        assert_eq!(args[0].name, "name");
        assert_eq!(args[0].required, Some(true));
    }

    #[tokio::test]
    async fn test_call_prompt() {
        let (_tx, _rx) = tokio::sync::mpsc::unbounded_channel();
        let logger = crate::logging::McpLogger::new(_tx, "test");
        let manager = PromptManager::new();
        manager.register(GreetingPrompt);
        let session = Session::new();

        let params = serde_json::json!({
            "name": "World"
        });

        let result = manager
            .call("greeting", params, &session, &logger)
            .await
            .unwrap();
        assert_eq!(result.messages.len(), 1);
        assert!(matches!(
            result.messages[0].role,
            crate::content::prompt::PromptRole::User
        ));
        assert_eq!(result.description, Some("A friendly greeting".to_string()));
    }

    #[tokio::test]
    async fn test_call_missing_prompt() {
        let (_tx, _rx) = tokio::sync::mpsc::unbounded_channel();
        let logger = crate::logging::McpLogger::new(_tx, "test");
        let manager = PromptManager::new();
        let session = Session::new();

        let params = serde_json::json!({});
        let result = manager.call("nonexistent", params, &session, &logger).await;

        assert!(matches!(result, Err(PromptError::NotFound(_))));
    }

    #[tokio::test]
    async fn test_prompt_invalid_arguments() {
        let (_tx, _rx) = tokio::sync::mpsc::unbounded_channel();
        let logger = crate::logging::McpLogger::new(_tx, "test");
        let manager = PromptManager::new();
        manager.register(GreetingPrompt);
        let session = Session::new();

        let params = serde_json::json!({}); // Missing required 'name' field

        let result = manager.call("greeting", params, &session, &logger).await;
        assert!(matches!(result, Err(PromptError::InvalidArguments(_))));
    }
}
