//! Session management
//!
//! Tracks individual MCP client connections with state storage, lifecycle state machine,
//! and per-session tool/resource/prompt customization.

use std::collections::{HashMap, HashSet};
use std::sync::{Arc, RwLock};

use dashmap::DashMap;
use serde_json::Value;
use uuid::Uuid;

use crate::protocol::capabilities::ClientCapabilities;
use crate::protocol::types::Implementation;
use crate::registry::prompts::Prompt;
use crate::registry::resources::Resource;
use crate::registry::tools::Tool;
use crate::server::profile::SessionProfile;

/// Session state storage (thread-safe)
pub type SessionState = Arc<RwLock<HashMap<String, Value>>>;

/// Session lifecycle states
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SessionLifecycle {
    /// Session created but not initialized
    Created,
    /// Session initialized and ready for requests
    Ready,
    /// Session encountered errors but still operational
    Degraded,
    /// Session closed
    Closed,
}

impl SessionLifecycle {
    /// Check if session can accept requests
    pub fn can_accept_requests(&self) -> bool {
        matches!(self, Self::Ready | Self::Degraded)
    }

    /// Check if session is healthy
    pub fn is_healthy(&self) -> bool {
        matches!(self, Self::Ready)
    }
}

/// Session represents a single MCP client connection
#[derive(Clone)]
pub struct Session {
    /// Unique session ID
    pub id: String,

    /// Client information (set during initialization)
    pub client_info: Option<Implementation>,

    /// Client capabilities (set during initialization)
    pub capabilities: Option<ClientCapabilities>,

    /// Negotiated protocol version (set during initialization)
    pub protocol_version: Option<String>,

    /// Session lifecycle state (Arc-wrapped for shared state across clones)
    lifecycle: Arc<RwLock<SessionLifecycle>>,

    /// Error count for degraded state transitions
    error_count: Arc<RwLock<u32>>,

    /// Per-session state storage
    state: SessionState,

    // Tool customization
    /// Tool overrides (same name, different implementation)
    tool_overrides: Arc<DashMap<String, Arc<dyn Tool>>>,
    /// Extra tools added to this session
    tool_extras: Arc<DashMap<String, Arc<dyn Tool>>>,
    /// Tools hidden from this session
    tool_hidden: Arc<RwLock<HashSet<String>>>,
    /// Tool aliases (alias -> actual name)
    tool_aliases: Arc<RwLock<HashMap<String, String>>>,

    // Resource customization
    /// Resource overrides (same URI, different implementation)
    resource_overrides: Arc<DashMap<String, Arc<dyn Resource>>>,
    /// Extra resources added to this session
    resource_extras: Arc<DashMap<String, Arc<dyn Resource>>>,
    /// Resources hidden from this session
    resource_hidden: Arc<RwLock<HashSet<String>>>,

    // Prompt customization
    /// Prompt overrides (same name, different implementation)
    prompt_overrides: Arc<DashMap<String, Arc<dyn Prompt>>>,
    /// Extra prompts added to this session
    prompt_extras: Arc<DashMap<String, Arc<dyn Prompt>>>,
    /// Prompts hidden from this session
    prompt_hidden: Arc<RwLock<HashSet<String>>>,
}

impl Session {
    /// Create new session with random UUID
    pub fn new() -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            client_info: None,
            capabilities: None,
            protocol_version: None,
            lifecycle: Arc::new(RwLock::new(SessionLifecycle::Created)),
            error_count: Arc::new(RwLock::new(0)),
            state: Arc::new(RwLock::new(HashMap::new())),
            // Tool customization
            tool_overrides: Arc::new(DashMap::new()),
            tool_extras: Arc::new(DashMap::new()),
            tool_hidden: Arc::new(RwLock::new(HashSet::new())),
            tool_aliases: Arc::new(RwLock::new(HashMap::new())),
            // Resource customization
            resource_overrides: Arc::new(DashMap::new()),
            resource_extras: Arc::new(DashMap::new()),
            resource_hidden: Arc::new(RwLock::new(HashSet::new())),
            // Prompt customization
            prompt_overrides: Arc::new(DashMap::new()),
            prompt_extras: Arc::new(DashMap::new()),
            prompt_hidden: Arc::new(RwLock::new(HashSet::new())),
        }
    }

    /// Create session with specific ID
    pub fn with_id(id: impl Into<String>) -> Self {
        Self {
            id: id.into(),
            client_info: None,
            capabilities: None,
            protocol_version: None,
            lifecycle: Arc::new(RwLock::new(SessionLifecycle::Created)),
            error_count: Arc::new(RwLock::new(0)),
            state: Arc::new(RwLock::new(HashMap::new())),
            // Tool customization
            tool_overrides: Arc::new(DashMap::new()),
            tool_extras: Arc::new(DashMap::new()),
            tool_hidden: Arc::new(RwLock::new(HashSet::new())),
            tool_aliases: Arc::new(RwLock::new(HashMap::new())),
            // Resource customization
            resource_overrides: Arc::new(DashMap::new()),
            resource_extras: Arc::new(DashMap::new()),
            resource_hidden: Arc::new(RwLock::new(HashSet::new())),
            // Prompt customization
            prompt_overrides: Arc::new(DashMap::new()),
            prompt_extras: Arc::new(DashMap::new()),
            prompt_hidden: Arc::new(RwLock::new(HashSet::new())),
        }
    }

    /// Initialize session with client info and capabilities
    /// Transitions: Created -> Ready
    pub fn initialize(
        &mut self,
        client_info: Implementation,
        capabilities: ClientCapabilities,
        protocol_version: String,
    ) {
        self.client_info = Some(client_info);
        self.capabilities = Some(capabilities);
        self.protocol_version = Some(protocol_version);
        *self.lifecycle.write().unwrap() = SessionLifecycle::Ready;
    }

    /// Check if session is initialized (Ready or Degraded)
    pub fn is_initialized(&self) -> bool {
        self.lifecycle.read().unwrap().can_accept_requests()
    }

    /// Get the negotiated protocol version
    pub fn protocol_version(&self) -> Option<&str> {
        self.protocol_version.as_deref()
    }

    /// Record an error, potentially transitioning to Degraded state
    /// Transitions: Ready -> Degraded (after threshold errors)
    pub fn record_error(&mut self) {
        if let Ok(mut count) = self.error_count.write() {
            *count += 1;
            // Transition to Degraded after 3 consecutive errors
            if *count >= 3 && *self.lifecycle.read().unwrap() == SessionLifecycle::Ready {
                *self.lifecycle.write().unwrap() = SessionLifecycle::Degraded;
            }
        }
    }

    /// Record success, potentially recovering from Degraded state
    /// Transitions: Degraded -> Ready (resets error count)
    pub fn record_success(&mut self) {
        if let Ok(mut count) = self.error_count.write() {
            *count = 0;
            if *self.lifecycle.read().unwrap() == SessionLifecycle::Degraded {
                *self.lifecycle.write().unwrap() = SessionLifecycle::Ready;
            }
        }
    }

    /// Close the session
    /// Transitions: Any -> Closed
    pub fn close(&mut self) {
        *self.lifecycle.write().unwrap() = SessionLifecycle::Closed;
    }

    /// Get current lifecycle state
    pub fn lifecycle(&self) -> SessionLifecycle {
        *self.lifecycle.read().unwrap()
    }

    /// Get current error count
    pub fn error_count(&self) -> u32 {
        self.error_count.read().map(|c| *c).unwrap_or(0)
    }

    /// Get state value
    pub fn get_state(&self, key: &str) -> Option<Value> {
        self.state.read().ok()?.get(key).cloned()
    }

    /// Set state value
    pub fn set_state(&self, key: impl Into<String>, value: Value) {
        if let Ok(mut state) = self.state.write() {
            state.insert(key.into(), value);
        }
    }

    /// Remove state value
    pub fn remove_state(&self, key: &str) -> Option<Value> {
        self.state.write().ok()?.remove(key)
    }

    /// Clear all state
    pub fn clear_state(&self) {
        if let Ok(mut state) = self.state.write() {
            state.clear();
        }
    }

    /// Get all state keys
    pub fn state_keys(&self) -> Vec<String> {
        self.state
            .read()
            .ok()
            .map(|state| state.keys().cloned().collect())
            .unwrap_or_default()
    }

    // ==================== Tool Management ====================

    /// Add an extra tool to this session
    pub fn add_tool(&self, tool: Arc<dyn Tool>) {
        let name = tool.name().to_string();
        self.tool_extras.insert(name, tool);
    }

    /// Override a tool (same name, different implementation)
    pub fn override_tool(&self, name: impl Into<String>, tool: Arc<dyn Tool>) {
        self.tool_overrides.insert(name.into(), tool);
    }

    /// Hide a tool from this session
    pub fn hide_tool(&self, name: impl Into<String>) {
        if let Ok(mut hidden) = self.tool_hidden.write() {
            hidden.insert(name.into());
        }
    }

    /// Unhide a tool
    pub fn unhide_tool(&self, name: &str) {
        if let Ok(mut hidden) = self.tool_hidden.write() {
            hidden.remove(name);
        }
    }

    /// Add a tool alias (alias -> target)
    pub fn alias_tool(&self, alias: impl Into<String>, target: impl Into<String>) {
        if let Ok(mut aliases) = self.tool_aliases.write() {
            aliases.insert(alias.into(), target.into());
        }
    }

    /// Remove a tool alias
    pub fn remove_tool_alias(&self, alias: &str) {
        if let Ok(mut aliases) = self.tool_aliases.write() {
            aliases.remove(alias);
        }
    }

    /// Check if a tool is hidden
    pub fn is_tool_hidden(&self, name: &str) -> bool {
        self.tool_hidden
            .read()
            .map(|hidden| hidden.contains(name))
            .unwrap_or(false)
    }

    /// Resolve a tool alias to its target name
    pub fn resolve_tool_alias<'a>(&self, name: &'a str) -> std::borrow::Cow<'a, str> {
        self.tool_aliases
            .read()
            .ok()
            .and_then(|aliases| aliases.get(name).cloned())
            .map(std::borrow::Cow::Owned)
            .unwrap_or(std::borrow::Cow::Borrowed(name))
    }

    /// Get a tool override for this session
    pub fn get_tool_override(&self, name: &str) -> Option<Arc<dyn Tool>> {
        self.tool_overrides.get(name).map(|r| Arc::clone(&r))
    }

    /// Get an extra tool added to this session
    pub fn get_tool_extra(&self, name: &str) -> Option<Arc<dyn Tool>> {
        self.tool_extras.get(name).map(|r| Arc::clone(&r))
    }

    /// Get all tool overrides
    pub fn tool_overrides(&self) -> &Arc<DashMap<String, Arc<dyn Tool>>> {
        &self.tool_overrides
    }

    /// Get all extra tools
    pub fn tool_extras(&self) -> &Arc<DashMap<String, Arc<dyn Tool>>> {
        &self.tool_extras
    }

    // ==================== Resource Management ====================

    /// Add an extra resource to this session
    pub fn add_resource(&self, resource: Arc<dyn Resource>) {
        let uri = resource.uri().to_string();
        self.resource_extras.insert(uri, resource);
    }

    /// Override a resource (same URI, different implementation)
    pub fn override_resource(&self, uri: impl Into<String>, resource: Arc<dyn Resource>) {
        self.resource_overrides.insert(uri.into(), resource);
    }

    /// Hide a resource from this session
    pub fn hide_resource(&self, uri: impl Into<String>) {
        if let Ok(mut hidden) = self.resource_hidden.write() {
            hidden.insert(uri.into());
        }
    }

    /// Unhide a resource
    pub fn unhide_resource(&self, uri: &str) {
        if let Ok(mut hidden) = self.resource_hidden.write() {
            hidden.remove(uri);
        }
    }

    /// Check if a resource is hidden
    pub fn is_resource_hidden(&self, uri: &str) -> bool {
        self.resource_hidden
            .read()
            .map(|hidden| hidden.contains(uri))
            .unwrap_or(false)
    }

    /// Get a resource override for this session
    pub fn get_resource_override(&self, uri: &str) -> Option<Arc<dyn Resource>> {
        self.resource_overrides.get(uri).map(|r| Arc::clone(&r))
    }

    /// Get an extra resource added to this session
    pub fn get_resource_extra(&self, uri: &str) -> Option<Arc<dyn Resource>> {
        self.resource_extras.get(uri).map(|r| Arc::clone(&r))
    }

    /// Get all resource overrides
    pub fn resource_overrides(&self) -> &Arc<DashMap<String, Arc<dyn Resource>>> {
        &self.resource_overrides
    }

    /// Get all extra resources
    pub fn resource_extras(&self) -> &Arc<DashMap<String, Arc<dyn Resource>>> {
        &self.resource_extras
    }

    // ==================== Prompt Management ====================

    /// Add an extra prompt to this session
    pub fn add_prompt(&self, prompt: Arc<dyn Prompt>) {
        let name = prompt.name().to_string();
        self.prompt_extras.insert(name, prompt);
    }

    /// Override a prompt (same name, different implementation)
    pub fn override_prompt(&self, name: impl Into<String>, prompt: Arc<dyn Prompt>) {
        self.prompt_overrides.insert(name.into(), prompt);
    }

    /// Hide a prompt from this session
    pub fn hide_prompt(&self, name: impl Into<String>) {
        if let Ok(mut hidden) = self.prompt_hidden.write() {
            hidden.insert(name.into());
        }
    }

    /// Unhide a prompt
    pub fn unhide_prompt(&self, name: &str) {
        if let Ok(mut hidden) = self.prompt_hidden.write() {
            hidden.remove(name);
        }
    }

    /// Check if a prompt is hidden
    pub fn is_prompt_hidden(&self, name: &str) -> bool {
        self.prompt_hidden
            .read()
            .map(|hidden| hidden.contains(name))
            .unwrap_or(false)
    }

    /// Get a prompt override for this session
    pub fn get_prompt_override(&self, name: &str) -> Option<Arc<dyn Prompt>> {
        self.prompt_overrides.get(name).map(|r| Arc::clone(&r))
    }

    /// Get an extra prompt added to this session
    pub fn get_prompt_extra(&self, name: &str) -> Option<Arc<dyn Prompt>> {
        self.prompt_extras.get(name).map(|r| Arc::clone(&r))
    }

    /// Get all prompt overrides
    pub fn prompt_overrides(&self) -> &Arc<DashMap<String, Arc<dyn Prompt>>> {
        &self.prompt_overrides
    }

    /// Get all extra prompts
    pub fn prompt_extras(&self) -> &Arc<DashMap<String, Arc<dyn Prompt>>> {
        &self.prompt_extras
    }

    // ==================== Profile Management ====================

    /// Apply a session profile
    ///
    /// This adds/overrides/hides tools, resources, and prompts according to the profile
    pub fn apply_profile(&self, profile: &SessionProfile) {
        // Apply tool configuration
        for tool in &profile.tool_extras {
            self.add_tool(Arc::clone(tool));
        }
        for (name, tool) in &profile.tool_overrides {
            self.override_tool(name.clone(), Arc::clone(tool));
        }
        for name in &profile.tool_hidden {
            self.hide_tool(name.clone());
        }
        for (alias, target) in &profile.tool_aliases {
            self.alias_tool(alias.clone(), target.clone());
        }

        // Apply resource configuration
        for resource in &profile.resource_extras {
            self.add_resource(Arc::clone(resource));
        }
        for (uri, resource) in &profile.resource_overrides {
            self.override_resource(uri.clone(), Arc::clone(resource));
        }
        for uri in &profile.resource_hidden {
            self.hide_resource(uri.clone());
        }

        // Apply prompt configuration
        for prompt in &profile.prompt_extras {
            self.add_prompt(Arc::clone(prompt));
        }
        for (name, prompt) in &profile.prompt_overrides {
            self.override_prompt(name.clone(), Arc::clone(prompt));
        }
        for name in &profile.prompt_hidden {
            self.hide_prompt(name.clone());
        }
    }

    /// Clear all session customizations
    pub fn clear_customizations(&self) {
        // Clear tools
        self.tool_overrides.clear();
        self.tool_extras.clear();
        if let Ok(mut hidden) = self.tool_hidden.write() {
            hidden.clear();
        }
        if let Ok(mut aliases) = self.tool_aliases.write() {
            aliases.clear();
        }

        // Clear resources
        self.resource_overrides.clear();
        self.resource_extras.clear();
        if let Ok(mut hidden) = self.resource_hidden.write() {
            hidden.clear();
        }

        // Clear prompts
        self.prompt_overrides.clear();
        self.prompt_extras.clear();
        if let Ok(mut hidden) = self.prompt_hidden.write() {
            hidden.clear();
        }
    }
}

// Implement Debug manually since trait objects don't implement Debug
impl std::fmt::Debug for Session {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Session")
            .field("id", &self.id)
            .field("client_info", &self.client_info)
            .field("capabilities", &self.capabilities)
            .field("lifecycle", &self.lifecycle)
            .field("error_count", &self.error_count())
            .field("tool_overrides_count", &self.tool_overrides.len())
            .field("tool_extras_count", &self.tool_extras.len())
            .field("resource_overrides_count", &self.resource_overrides.len())
            .field("resource_extras_count", &self.resource_extras.len())
            .field("prompt_overrides_count", &self.prompt_overrides.len())
            .field("prompt_extras_count", &self.prompt_extras.len())
            .finish()
    }
}

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

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

    #[test]
    fn test_session_creation() {
        let session = Session::new();
        assert!(!session.id.is_empty());
        assert_eq!(session.lifecycle(), SessionLifecycle::Created);
        assert!(!session.is_initialized());
        assert!(session.client_info.is_none());
        assert!(session.capabilities.is_none());
    }

    #[test]
    fn test_session_with_id() {
        let session = Session::with_id("test-session");
        assert_eq!(session.id, "test-session");
        assert_eq!(session.lifecycle(), SessionLifecycle::Created);
    }

    #[test]
    fn test_session_initialization() {
        let mut session = Session::new();
        let client_info = Implementation {
            name: "test-client".to_string(),
            version: "1.0.0".to_string(),
        };
        let capabilities = ClientCapabilities::default();

        session.initialize(client_info.clone(), capabilities, "2025-06-18".to_string());

        assert!(session.is_initialized());
        assert_eq!(session.lifecycle(), SessionLifecycle::Ready);
        assert_eq!(session.client_info.unwrap().name, "test-client");
    }

    #[test]
    fn test_session_lifecycle_transitions() {
        let mut session = Session::new();
        assert_eq!(session.lifecycle(), SessionLifecycle::Created);
        assert!(!session.lifecycle().can_accept_requests());

        // Initialize -> Ready
        session.initialize(
            Implementation {
                name: "test".into(),
                version: "1.0".into(),
            },
            ClientCapabilities::default(),
            "2025-06-18".to_string(),
        );
        assert_eq!(session.lifecycle(), SessionLifecycle::Ready);
        assert!(session.lifecycle().can_accept_requests());
        assert!(session.lifecycle().is_healthy());

        // Record errors -> Degraded (after 3)
        session.record_error();
        assert_eq!(session.lifecycle(), SessionLifecycle::Ready);
        session.record_error();
        assert_eq!(session.lifecycle(), SessionLifecycle::Ready);
        session.record_error();
        assert_eq!(session.lifecycle(), SessionLifecycle::Degraded);
        assert!(session.lifecycle().can_accept_requests());
        assert!(!session.lifecycle().is_healthy());

        // Record success -> Ready (recovery)
        session.record_success();
        assert_eq!(session.lifecycle(), SessionLifecycle::Ready);
        assert_eq!(session.error_count(), 0);

        // Close -> Closed
        session.close();
        assert_eq!(session.lifecycle(), SessionLifecycle::Closed);
        assert!(!session.lifecycle().can_accept_requests());
    }

    #[test]
    fn test_session_state() {
        let session = Session::new();

        // Set state
        session.set_state("key1", Value::String("value1".to_string()));
        session.set_state("key2", Value::Number(42.into()));

        // Get state
        assert_eq!(
            session.get_state("key1"),
            Some(Value::String("value1".to_string()))
        );
        assert_eq!(session.get_state("key2"), Some(Value::Number(42.into())));
        assert_eq!(session.get_state("nonexistent"), None);

        // State keys
        let keys = session.state_keys();
        assert_eq!(keys.len(), 2);
        assert!(keys.contains(&"key1".to_string()));
        assert!(keys.contains(&"key2".to_string()));

        // Remove state
        let removed = session.remove_state("key1");
        assert_eq!(removed, Some(Value::String("value1".to_string())));
        assert_eq!(session.get_state("key1"), None);

        // Clear state
        session.clear_state();
        assert_eq!(session.state_keys().len(), 0);
    }

    #[test]
    fn test_session_clone() {
        let session1 = Session::with_id("test");
        session1.set_state("shared", Value::Bool(true));

        let session2 = session1.clone();

        // Both sessions share the same state storage
        assert_eq!(session1.id, session2.id);
        assert_eq!(session2.get_state("shared"), Some(Value::Bool(true)));

        // Modifying state in one affects the other
        session2.set_state("shared", Value::Bool(false));
        assert_eq!(session1.get_state("shared"), Some(Value::Bool(false)));
    }
}
