//! LLM Gateway — Conversation Manager (Phase-79: T79.3 ~ T79.9)
//!
//! Provides:
//! - Conversation CRUD with ownership-based permission checks
//! - Message persistence + history query + context window management
//! - `send_message()` / `send_with_rag()` chat modes
//! - Token usage tracking + cost calculation per ModelPricing
//! - CancellationToken stop-generation support
//! - Conversation model switch / share / favorite / archive
//!
//! # Architecture
//!
//! ```text
//! ConversationManager
//! ├── conversations: DashMap<conv_id, LlmConversation>
//! ├── messages: DashMap<conv_id, Vec<StoredMessage>>
//! ├── active_streams: DashMap<conv_id, CancellationToken>
//! ├── store_dir: PathBuf   (JSON persistence)
//! └── refs ─┬─ LlmStore (provider/model registry)
//!           └─ ProviderRouter (route + failover)
//! ```

use std::path::{Path, PathBuf};
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Instant;

use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use tokio_util::sync::CancellationToken;
use tracing::{debug, info, warn};

use super::client::{LlmMessage, MessageRole};
use super::conversation::{
    ConversationConfig, ConversationQuery, LlmConversation,
    MessageQuery, MessageStatus, StoredMessage, UsageStats,
};
use super::error::{LlmError, LlmResult};
use super::memory::MemoryManager;
use super::router::ProviderRouter;
use super::store::LlmStore;
use super::types::{
    ChatConfig, StreamEvent, TokenUsage,
};

// ============================================================================
// Persistence structure
// ============================================================================

#[derive(Serialize, Deserialize, Default)]
struct ConversationData {
    conversations: Vec<LlmConversation>,
    messages: Vec<StoredMessage>,
}

// ============================================================================
// ConversationManager
// ============================================================================

/// Manages conversations, messages, and chat interactions with LLMs.
///
/// Thread-safe via `DashMap` — safe for concurrent use from multiple
/// async tasks. Persistence is JSON file based (auto-save on write).
///
/// # Ownership Model
///
/// Each conversation has an `owner_id`. CRUD operations check:
/// - `can_read()`: Owner always, Public everyone, Shared specific users
/// - `can_edit()`: Owner only (unless protection blocks)
/// - `can_delete()`: Owner only (unless protection blocks)
pub struct ConversationManager {
    /// All conversations indexed by ID.
    conversations: DashMap<String, LlmConversation>,

    /// Messages indexed by conversation_id → Vec<StoredMessage>.
    messages: DashMap<String, Vec<StoredMessage>>,

    /// Active streaming sessions: conversation_id → CancellationToken.
    active_streams: DashMap<String, CancellationToken>,

    /// Provider & model registry (for model lookup + pricing).
    store: Arc<LlmStore>,

    /// Provider router (for routing chat requests).
    router: Arc<ProviderRouter>,

    /// Persistence file path.
    file_path: PathBuf,

    /// Dirty flag for save optimization.
    dirty: AtomicBool,

    /// Long-term conversational memory manager (Phase-84).
    ///
    /// When set, `send_message()` will automatically:
    /// - Recall relevant memories before sending
    /// - Store QA pairs for semantic recall
    /// - Trigger summarization when messages exceed window
    /// - Extract user profiles periodically
    memory: Option<Arc<MemoryManager>>,
}

impl ConversationManager {
    // ========================================================================
    // Construction
    // ========================================================================

    /// Create a new ConversationManager.
    ///
    /// - `dir`: Directory for persistence file (`conversations.json`)
    /// - `store`: Shared LlmStore reference (provider/model registry)
    /// - `router`: Shared ProviderRouter reference
    pub fn new(
        dir: impl AsRef<Path>,
        store: Arc<LlmStore>,
        router: Arc<ProviderRouter>,
    ) -> LlmResult<Self> {
        let dir = dir.as_ref();
        std::fs::create_dir_all(dir).map_err(|e| {
            LlmError::StorageError(format!(
                "Failed to create conversation dir {}: {e}",
                dir.display()
            ))
        })?;

        let file_path = dir.join("conversations.json");
        let mgr = Self {
            conversations: DashMap::new(),
            messages: DashMap::new(),
            active_streams: DashMap::new(),
            store,
            router,
            file_path,
            dirty: AtomicBool::new(false),
            memory: None,
        };

        mgr.load()?;
        Ok(mgr)
    }

    /// Set the memory manager for long-term conversational memory (Phase-84).
    ///
    /// When set, conversations with `enable_memory: true` in their config
    /// will automatically recall past context, generate summaries, and
    /// extract user profiles.
    pub fn set_memory_manager(&mut self, memory: Arc<MemoryManager>) {
        self.memory = Some(memory);
    }

    /// Get a reference to the memory manager (if configured).
    pub fn memory_manager(&self) -> Option<&MemoryManager> {
        self.memory.as_deref()
    }

    // ========================================================================
    // T79.3: Conversation CRUD
    // ========================================================================

    /// Create a new conversation.
    ///
    /// Sets `owner_id` to the provided user. Returns the conversation ID.
    pub fn create_conversation(
        &self,
        owner_id: &str,
        provider_id: &str,
        model_id: &str,
        model_name: &str,
    ) -> LlmResult<LlmConversation> {
        let conv = LlmConversation::new(owner_id, provider_id, model_id, model_name);
        let id = conv.id.clone();
        debug!(conv_id = %id, owner = %owner_id, model = %model_name, "Creating conversation");
        self.conversations.insert(id.clone(), conv.clone());
        self.messages.insert(id, Vec::new());
        self.mark_dirty_and_save()?;
        Ok(conv)
    }

    /// Create a conversation with a builder-style config.
    pub fn create_conversation_full(
        &self,
        mut conv: LlmConversation,
    ) -> LlmResult<LlmConversation> {
        let id = conv.id.clone();
        // Ensure timestamps are set
        if conv.created_at == 0 {
            let now = now_ms();
            conv.created_at = now;
            conv.updated_at = now;
        }
        debug!(conv_id = %id, owner = %conv.owner_id, "Creating conversation (full)");
        self.conversations.insert(id.clone(), conv.clone());
        self.messages.insert(id, Vec::new());
        self.mark_dirty_and_save()?;
        Ok(conv)
    }

    /// Get a conversation by ID (with permission check).
    pub fn get_conversation(
        &self,
        conv_id: &str,
        viewer_id: &str,
    ) -> LlmResult<LlmConversation> {
        let conv = self
            .conversations
            .get(conv_id)
            .ok_or_else(|| LlmError::NotFound(format!("Conversation '{conv_id}' not found")))?;

        if !conv.can_read(viewer_id) {
            return Err(LlmError::PermissionDenied(format!(
                "User '{viewer_id}' cannot read conversation '{conv_id}'"
            )));
        }

        Ok(conv.value().clone())
    }

    /// List conversations with filters (respects visibility).
    pub fn list_conversations(&self, query: &ConversationQuery) -> Vec<LlmConversation> {
        let viewer_id = query.viewer_id.as_deref().unwrap_or("");

        let mut results: Vec<LlmConversation> = self
            .conversations
            .iter()
            .filter(|r| {
                let c = r.value();

                // Visibility check
                if !viewer_id.is_empty() && !c.can_read(viewer_id) {
                    return false;
                }

                // Owner filter
                if let Some(ref oid) = query.owner_id {
                    if c.owner_id != *oid {
                        return false;
                    }
                }

                // Favorite filter
                if let Some(fav) = query.is_favorite {
                    if c.is_favorite != fav {
                        return false;
                    }
                }

                // Archived filter
                if let Some(arch) = query.is_archived {
                    if c.is_archived != arch {
                        return false;
                    }
                }

                // Title search
                if let Some(ref search) = query.title_search {
                    if !c.title.to_lowercase().contains(&search.to_lowercase()) {
                        return false;
                    }
                }

                // Model filter
                if let Some(ref mid) = query.model_id {
                    if c.model_id != *mid {
                        return false;
                    }
                }

                // Provider filter
                if let Some(ref pid) = query.provider_id {
                    if c.provider_id != *pid {
                        return false;
                    }
                }

                true
            })
            .map(|r| r.value().clone())
            .collect();

        // Sort by updated_at descending (most recent first)
        results.sort_by(|a, b| b.updated_at.cmp(&a.updated_at));

        // Pagination
        let total = results.len();
        let offset = query.offset.min(total);
        let end = (offset + query.limit).min(total);
        results[offset..end].to_vec()
    }

    /// Update a conversation (permission check).
    pub fn update_conversation(
        &self,
        conv_id: &str,
        user_id: &str,
        update_fn: impl FnOnce(&mut LlmConversation),
    ) -> LlmResult<LlmConversation> {
        let mut entry = self
            .conversations
            .get_mut(conv_id)
            .ok_or_else(|| LlmError::NotFound(format!("Conversation '{conv_id}' not found")))?;

        if !entry.can_edit(user_id) {
            return Err(LlmError::PermissionDenied(format!(
                "User '{user_id}' cannot edit conversation '{conv_id}'"
            )));
        }

        update_fn(entry.value_mut());
        entry.touch();
        let conv = entry.value().clone();
        drop(entry);
        self.mark_dirty_and_save()?;
        Ok(conv)
    }

    /// Update conversation config (convenience wrapper for memory/RAG settings).
    pub fn update_conversation_config(
        &self,
        conv_id: &str,
        user_id: &str,
        config_fn: impl FnOnce(&mut ConversationConfig),
    ) -> LlmResult<LlmConversation> {
        self.update_conversation(conv_id, user_id, |conv| {
            config_fn(&mut conv.config);
        })
    }

    /// Delete a conversation and all its messages (permission check).
    pub fn delete_conversation(&self, conv_id: &str, user_id: &str) -> LlmResult<bool> {
        // Check permissions before deleting
        {
            let conv = self
                .conversations
                .get(conv_id)
                .ok_or_else(|| {
                    LlmError::NotFound(format!("Conversation '{conv_id}' not found"))
                })?;
            if !conv.can_delete(user_id) {
                return Err(LlmError::PermissionDenied(format!(
                    "User '{user_id}' cannot delete conversation '{conv_id}'"
                )));
            }
        }

        // Cancel any active stream
        self.active_streams.remove(conv_id);

        let existed = self.conversations.remove(conv_id).is_some();
        self.messages.remove(conv_id);

        if existed {
            info!(conv_id = %conv_id, user = %user_id, "Deleted conversation");
            self.mark_dirty_and_save()?;
        }
        Ok(existed)
    }

    // ========================================================================
    // T79.4: Message Persistence + History Query + Context Window
    // ========================================================================

    /// Add a message to a conversation.
    pub fn add_message(&self, msg: StoredMessage) -> LlmResult<StoredMessage> {
        let conv_id = msg.conversation_id.clone();
        let mut msgs = self.messages.entry(conv_id.clone()).or_default();
        msgs.push(msg.clone());
        drop(msgs);

        // Update conversation counters if there's token usage
        if let Some(ref usage) = msg.token_usage {
            if let Some(mut conv) = self.conversations.get_mut(&conv_id) {
                conv.add_message_stats(usage);
            }
        } else if let Some(mut conv) = self.conversations.get_mut(&conv_id) {
            conv.message_count += 1;
            conv.touch();
        }

        self.mark_dirty_and_save()?;
        Ok(msg)
    }

    /// Get messages for a conversation (with pagination + time range).
    pub fn get_messages(&self, query: &MessageQuery) -> LlmResult<Vec<StoredMessage>> {
        let msgs = self
            .messages
            .get(&query.conversation_id)
            .map(|r| r.value().clone())
            .unwrap_or_default();

        let filtered: Vec<StoredMessage> = msgs
            .into_iter()
            .filter(|m| {
                if let Some(after) = query.after_timestamp {
                    if m.created_at <= after {
                        return false;
                    }
                }
                if let Some(before) = query.before_timestamp {
                    if m.created_at >= before {
                        return false;
                    }
                }
                true
            })
            .collect();

        // Pagination
        let total = filtered.len();
        let offset = query.offset.min(total);
        let end = (offset + query.limit).min(total);
        Ok(filtered[offset..end].to_vec())
    }

    /// Get the last N messages as context for LLM input.
    ///
    /// Respects `ConversationConfig.context_window` if set,
    /// otherwise returns all messages. Prepends system prompt if set.
    pub fn build_context(
        &self,
        conv_id: &str,
    ) -> LlmResult<Vec<LlmMessage>> {
        let conv = self
            .conversations
            .get(conv_id)
            .ok_or_else(|| LlmError::NotFound(format!("Conversation '{conv_id}' not found")))?;

        if conv.config.disable_context {
            // Stateless mode: only return the last user message
            let msgs = self.messages.get(conv_id);
            if let Some(ref msgs) = msgs {
                if let Some(last) = msgs.last() {
                    let mut result = Vec::new();
                    if let Some(ref sp) = conv.system_prompt {
                        result.push(LlmMessage::system(sp.clone()));
                    }
                    result.push(last.to_llm_message());
                    return Ok(result);
                }
            }
            return Ok(Vec::new());
        }

        let all_msgs = self
            .messages
            .get(conv_id)
            .map(|r| r.value().clone())
            .unwrap_or_default();

        // Apply context window limit
        let window = conv.config.context_window.unwrap_or(usize::MAX);
        let start = all_msgs.len().saturating_sub(window);
        let windowed = &all_msgs[start..];

        let mut result = Vec::with_capacity(windowed.len() + 1);

        // Prepend system prompt
        if let Some(ref sp) = conv.system_prompt {
            result.push(LlmMessage::system(sp.clone()));
        }

        // Convert stored messages to LlmMessage
        for msg in windowed {
            result.push(msg.to_llm_message());
        }

        Ok(result)
    }

    // ========================================================================
    // T79.5: send_message / send_with_rag
    // ========================================================================

    /// Send a user message and get an LLM response (non-streaming).
    ///
    /// Flow:
    /// 1. Store user message
    /// 2. Build context (system prompt + history)
    /// 3. Build ChatConfig from ConversationConfig
    /// 4. Route to provider via ProviderRouter
    /// 5. Call LlmClient.chat()
    /// 6. Calculate cost from ModelPricing
    /// 7. Store assistant response
    /// 8. Update conversation counters
    pub async fn send_message(
        &self,
        conv_id: &str,
        user_id: &str,
        content: &str,
    ) -> LlmResult<StoredMessage> {
        // Permission check
        {
            let conv = self
                .conversations
                .get(conv_id)
                .ok_or_else(|| {
                    LlmError::NotFound(format!("Conversation '{conv_id}' not found"))
                })?;
            if !conv.can_edit(user_id) {
                return Err(LlmError::PermissionDenied(format!(
                    "User '{user_id}' cannot send messages to conversation '{conv_id}'"
                )));
            }
        }

        // 1. Store user message
        let user_msg = StoredMessage::user(conv_id, content);
        self.add_message(user_msg)?;

        // 2. Build context
        let mut context = self.build_context(conv_id)?;

        // 2.5 (Phase-84) Inject memory context if enabled
        let memory_enabled = {
            let conv = self.conversations.get(conv_id).unwrap();
            conv.config.enable_memory
        };
        if memory_enabled {
            if let Some(ref mem) = self.memory {
                match mem.recall(conv_id, user_id, None).await {
                    Ok(mem_ctx) => {
                        if let Some(mem_system_msg) = mem_ctx.to_system_message() {
                            // Inject memory context as the second system message
                            // (after the original system prompt, before history)
                            let insert_pos = if !context.is_empty()
                                && context[0].role == MessageRole::System
                            {
                                1
                            } else {
                                0
                            };
                            context.insert(
                                insert_pos,
                                LlmMessage::system(format!(
                                    "[Long-term Memory]\n{}",
                                    mem_system_msg
                                )),
                            );
                            debug!(
                                conv_id = %conv_id,
                                tokens = mem_ctx.estimated_tokens,
                                "Injected memory context"
                            );
                        }
                    }
                    Err(e) => {
                        warn!(conv_id = %conv_id, error = %e, "Failed to recall memories, continuing without");
                    }
                }
            }
        }

        // 3. Build ChatConfig
        let config = self.build_chat_config(conv_id)?;

        // 4-5. Route and call LLM
        let start = Instant::now();
        let model_name = {
            let conv = self.conversations.get(conv_id).unwrap();
            conv.model_name.clone()
        };

        let (client, _provider_id) = self.router.route(&model_name)?;
        let mut response = client.chat(&context, &config).await?;
        let duration = start.elapsed().as_millis() as u64;
        response.duration_ms = Some(duration);

        // 6. Calculate cost
        if let Some(ref mut usage) = response.usage {
            self.calculate_cost(conv_id, usage);
        }

        // 7. Store assistant response
        let assistant_msg = StoredMessage::from_response(conv_id, &response);
        let stored = self.add_message(assistant_msg)?;

        // 7.5 (Phase-84) Memorize the exchange + trigger summarization/profile
        if memory_enabled {
            if let Some(ref mem) = self.memory {
                // Store QA pair as semantic memory (no embedding yet — caller can provide)
                if let Err(e) = mem
                    .memorize(conv_id, user_id, content, &response.content, None)
                    .await
                {
                    warn!(conv_id = %conv_id, error = %e, "Failed to memorize exchange");
                }

                // Trigger summarization if needed
                if let Some(all_msgs) = self.messages.get(conv_id) {
                    let msgs = all_msgs.value();
                    let (route_client, _) = self.router.route(&model_name)?;
                    if let Err(e) = mem
                        .maybe_summarize(conv_id, user_id, msgs, route_client.as_ref())
                        .await
                    {
                        warn!(conv_id = %conv_id, error = %e, "Failed to summarize");
                    }

                    // Trigger profile extraction if needed
                    if let Err(e) = mem
                        .maybe_extract_profile(conv_id, user_id, msgs, route_client.as_ref())
                        .await
                    {
                        warn!(conv_id = %conv_id, error = %e, "Failed to extract profile");
                    }
                }
            }
        }

        debug!(
            conv_id = %conv_id,
            model = %response.model,
            duration_ms = duration,
            "Message sent successfully"
        );

        Ok(stored)
    }

    /// Send a user message with streaming response.
    ///
    /// Returns a stream of `StreamEvent`s. Creates a CancellationToken
    /// for this conversation that can be used to stop generation.
    pub async fn send_message_stream(
        &self,
        conv_id: &str,
        user_id: &str,
        content: &str,
    ) -> LlmResult<Pin<Box<dyn futures_core::Stream<Item = LlmResult<StreamEvent>> + Send>>> {
        // Permission check
        {
            let conv = self.conversations.get(conv_id).ok_or_else(|| {
                LlmError::NotFound(format!("Conversation '{conv_id}' not found"))
            })?;
            if !conv.can_edit(user_id) {
                return Err(LlmError::PermissionDenied(format!(
                    "User '{user_id}' cannot send messages to conversation '{conv_id}'"
                )));
            }
        }

        // Store user message
        let user_msg = StoredMessage::user(conv_id, content);
        self.add_message(user_msg)?;

        // Build context and config
        let mut context = self.build_context(conv_id)?;
        let config = self.build_chat_config(conv_id)?;

        // (Phase-84) Inject memory context if enabled
        let memory_enabled = {
            let conv = self.conversations.get(conv_id).unwrap();
            conv.config.enable_memory
        };
        if memory_enabled {
            if let Some(ref mem) = self.memory {
                match mem.recall(conv_id, user_id, None).await {
                    Ok(mem_ctx) => {
                        if let Some(mem_system_msg) = mem_ctx.to_system_message() {
                            let insert_pos = if !context.is_empty()
                                && context[0].role == MessageRole::System
                            {
                                1
                            } else {
                                0
                            };
                            context.insert(
                                insert_pos,
                                LlmMessage::system(format!(
                                    "[Long-term Memory]\n{}",
                                    mem_system_msg
                                )),
                            );
                            debug!(
                                conv_id = %conv_id,
                                tokens = mem_ctx.estimated_tokens,
                                "Injected memory context (stream)"
                            );
                        }
                    }
                    Err(e) => {
                        warn!(conv_id = %conv_id, error = %e, "Failed to recall memories for stream");
                    }
                }
            }
        }

        // Create cancellation token
        let cancel_token = CancellationToken::new();
        self.active_streams
            .insert(conv_id.to_string(), cancel_token.clone());

        // Route and get stream
        let model_name = {
            let conv = self.conversations.get(conv_id).unwrap();
            conv.model_name.clone()
        };
        let (client, _provider_id) = self.router.route(&model_name)?;
        let inner_stream = client.chat_stream(&context, &config).await?;

        // Wrap stream with cancellation support + auto-save
        let conv_id_owned = conv_id.to_string();
        let conversations = self.conversations.clone();
        let messages = self.messages.clone();
        let active_streams = self.active_streams.clone();
        let store_ref = self.store.clone();
        let file_path = self.file_path.clone();

        let stream = CancellableStream::new(
            inner_stream,
            cancel_token,
            conv_id_owned,
            conversations,
            messages,
            active_streams,
            store_ref,
            file_path,
        );

        Ok(Box::pin(stream))
    }

    // ========================================================================
    // T79.6: Token Usage + Cost Calculation
    // ========================================================================

    /// Calculate and set cost on a TokenUsage based on model pricing.
    fn calculate_cost(&self, conv_id: &str, usage: &mut TokenUsage) {
        let model_id = self
            .conversations
            .get(conv_id)
            .map(|c| c.model_id.clone())
            .unwrap_or_default();

        if let Some(model) = self.store.get_model(&model_id) {
            if let Some(ref pricing) = model.pricing {
                let cost = pricing.calculate_cost(usage.prompt_tokens, usage.completion_tokens);
                usage.cost = Some(cost);
            }
        }
    }

    /// Get aggregated usage statistics for a user.
    pub fn get_usage_stats(&self, user_id: &str) -> UsageStats {
        let mut stats = UsageStats::default();

        for entry in self.conversations.iter() {
            let conv = entry.value();
            if conv.owner_id != user_id {
                continue;
            }
            stats.conversation_count += 1;

            if let Some(msgs) = self.messages.get(&conv.id) {
                for msg in msgs.iter() {
                    if let Some(ref usage) = msg.token_usage {
                        stats.total_prompt_tokens += usage.prompt_tokens as u64;
                        stats.total_completion_tokens += usage.completion_tokens as u64;
                        stats.total_tokens += usage.total_tokens as u64;
                        if let Some(cost) = usage.cost {
                            stats.total_cost += cost;
                        }
                        stats.message_count += 1;
                    }
                }
            }
        }

        stats
    }

    /// Get usage stats filtered by model.
    pub fn get_usage_stats_by_model(&self, user_id: &str, model_id: &str) -> UsageStats {
        let mut stats = UsageStats::default();

        for entry in self.conversations.iter() {
            let conv = entry.value();
            if conv.owner_id != user_id || conv.model_id != model_id {
                continue;
            }
            stats.conversation_count += 1;

            if let Some(msgs) = self.messages.get(&conv.id) {
                for msg in msgs.iter() {
                    if let Some(ref usage) = msg.token_usage {
                        stats.total_prompt_tokens += usage.prompt_tokens as u64;
                        stats.total_completion_tokens += usage.completion_tokens as u64;
                        stats.total_tokens += usage.total_tokens as u64;
                        if let Some(cost) = usage.cost {
                            stats.total_cost += cost;
                        }
                        stats.message_count += 1;
                    }
                }
            }
        }

        stats
    }

    // ========================================================================
    // T79.7: CancellationToken + Stop Generation
    // ========================================================================

    /// Stop an active streaming generation for a conversation.
    ///
    /// The stream will end with a partial message saved automatically.
    pub fn stop_generation(&self, conv_id: &str) -> bool {
        if let Some((_key, token)) = self.active_streams.remove(conv_id) {
            info!(conv_id = %conv_id, "Stopping generation");
            token.cancel();
            true
        } else {
            false
        }
    }

    /// Check if a conversation has an active generation stream.
    pub fn is_generating(&self, conv_id: &str) -> bool {
        self.active_streams
            .get(conv_id)
            .map(|t| !t.is_cancelled())
            .unwrap_or(false)
    }

    // ========================================================================
    // T79.8: Model Switch + Share/Favorite/Archive
    // ========================================================================

    /// Switch the model used by a conversation.
    pub fn switch_model(
        &self,
        conv_id: &str,
        user_id: &str,
        new_provider_id: &str,
        new_model_id: &str,
        new_model_name: &str,
    ) -> LlmResult<LlmConversation> {
        self.update_conversation(conv_id, user_id, |conv| {
            conv.provider_id = new_provider_id.to_string();
            conv.model_id = new_model_id.to_string();
            conv.model_name = new_model_name.to_string();
        })
    }

    /// Toggle favorite status.
    pub fn toggle_favorite(
        &self,
        conv_id: &str,
        user_id: &str,
    ) -> LlmResult<bool> {
        let mut new_state = false;
        self.update_conversation(conv_id, user_id, |conv| {
            conv.is_favorite = !conv.is_favorite;
            new_state = conv.is_favorite;
        })?;
        Ok(new_state)
    }

    /// Toggle archive status.
    pub fn toggle_archive(
        &self,
        conv_id: &str,
        user_id: &str,
    ) -> LlmResult<bool> {
        let mut new_state = false;
        self.update_conversation(conv_id, user_id, |conv| {
            conv.is_archived = !conv.is_archived;
            new_state = conv.is_archived;
        })?;
        Ok(new_state)
    }

    /// Share a conversation with specific users.
    pub fn share_conversation(
        &self,
        conv_id: &str,
        owner_id: &str,
        user_ids: Vec<String>,
    ) -> LlmResult<LlmConversation> {
        use crate::types::ownership::Visibility;
        self.update_conversation(conv_id, owner_id, |conv| {
            conv.visibility = Visibility::Shared { user_ids };
        })
    }

    /// Make a conversation public.
    pub fn make_public(
        &self,
        conv_id: &str,
        owner_id: &str,
    ) -> LlmResult<LlmConversation> {
        use crate::types::ownership::Visibility;
        self.update_conversation(conv_id, owner_id, |conv| {
            conv.visibility = Visibility::Public;
        })
    }

    /// Make a conversation private again.
    pub fn make_private(
        &self,
        conv_id: &str,
        owner_id: &str,
    ) -> LlmResult<LlmConversation> {
        use crate::types::ownership::Visibility;
        self.update_conversation(conv_id, owner_id, |conv| {
            conv.visibility = Visibility::Private;
        })
    }

    /// Update conversation title.
    pub fn set_title(
        &self,
        conv_id: &str,
        user_id: &str,
        title: &str,
    ) -> LlmResult<LlmConversation> {
        self.update_conversation(conv_id, user_id, |conv| {
            conv.title = title.to_string();
        })
    }

    /// Update system prompt.
    pub fn set_system_prompt(
        &self,
        conv_id: &str,
        user_id: &str,
        prompt: Option<String>,
    ) -> LlmResult<LlmConversation> {
        self.update_conversation(conv_id, user_id, |conv| {
            conv.system_prompt = prompt;
        })
    }

    /// Update conversation config.
    pub fn set_config(
        &self,
        conv_id: &str,
        user_id: &str,
        config: ConversationConfig,
    ) -> LlmResult<LlmConversation> {
        self.update_conversation(conv_id, user_id, |conv| {
            conv.config = config;
        })
    }

    // ========================================================================
    // Internal helpers
    // ========================================================================

    /// Build a `ChatConfig` from the conversation's `ConversationConfig`.
    fn build_chat_config(&self, conv_id: &str) -> LlmResult<ChatConfig> {
        let conv = self
            .conversations
            .get(conv_id)
            .ok_or_else(|| LlmError::NotFound(format!("Conversation '{conv_id}' not found")))?;

        let cc = &conv.config;
        Ok(ChatConfig {
            temperature: cc.temperature,
            max_tokens: cc.max_tokens,
            top_p: cc.top_p,
            top_k: cc.top_k,
            stream: false,
            enable_thinking: cc.enable_thinking,
            response_format: cc.response_format.clone(),
            ..Default::default()
        })
    }

    fn mark_dirty_and_save(&self) -> LlmResult<()> {
        self.dirty.store(true, Ordering::Release);
        self.save()
    }

    fn save(&self) -> LlmResult<()> {
        if !self.dirty.load(Ordering::Acquire) {
            return Ok(());
        }

        let data = ConversationData {
            conversations: self
                .conversations
                .iter()
                .map(|r| r.value().clone())
                .collect(),
            messages: self
                .messages
                .iter()
                .flat_map(|r| r.value().clone())
                .collect(),
        };

        let json = serde_json::to_string_pretty(&data).map_err(|e| {
            LlmError::StorageError(format!("Failed to serialize conversations: {e}"))
        })?;

        std::fs::write(&self.file_path, &json).map_err(|e| {
            LlmError::StorageError(format!(
                "Failed to write {}: {e}",
                self.file_path.display()
            ))
        })?;

        self.dirty.store(false, Ordering::Release);
        Ok(())
    }

    fn load(&self) -> LlmResult<()> {
        if !self.file_path.exists() {
            return Ok(());
        }

        let content = std::fs::read_to_string(&self.file_path).map_err(|e| {
            LlmError::StorageError(format!(
                "Failed to read {}: {e}",
                self.file_path.display()
            ))
        })?;

        if content.trim().is_empty() {
            return Ok(());
        }

        let data: ConversationData = serde_json::from_str(&content).map_err(|e| {
            LlmError::StorageError(format!(
                "Failed to parse {}: {e}",
                self.file_path.display()
            ))
        })?;

        for conv in data.conversations {
            self.conversations.insert(conv.id.clone(), conv);
        }

        // Group messages by conversation_id
        for msg in data.messages {
            self.messages
                .entry(msg.conversation_id.clone())
                .or_default()
                .push(msg);
        }

        // Sort messages within each conversation by created_at
        for mut entry in self.messages.iter_mut() {
            entry.value_mut().sort_by_key(|m| m.created_at);
        }

        debug!(
            convs = self.conversations.len(),
            msgs = self.messages.iter().map(|r| r.value().len()).sum::<usize>(),
            "Loaded conversations from {}",
            self.file_path.display()
        );

        Ok(())
    }

    /// Get the total number of conversations.
    pub fn conversation_count(&self) -> usize {
        self.conversations.len()
    }

    /// Get the total number of messages across all conversations.
    pub fn message_count(&self) -> usize {
        self.messages.iter().map(|r| r.value().len()).sum()
    }
}

// ============================================================================
// CancellableStream — wraps a stream with cancellation + auto-save
// ============================================================================

/// A stream wrapper that:
/// - Monitors a `CancellationToken` to stop generation
/// - Accumulates text/thinking content from stream events
/// - Auto-saves the assistant message when the stream ends
struct CancellableStream {
    inner: Pin<Box<dyn futures_core::Stream<Item = LlmResult<StreamEvent>> + Send>>,
    cancel_token: CancellationToken,
    conv_id: String,
    accumulated_text: String,
    accumulated_thinking: String,
    usage: Option<TokenUsage>,
    model: String,
    finished: bool,
    conversations: DashMap<String, LlmConversation>,
    messages: DashMap<String, Vec<StoredMessage>>,
    active_streams: DashMap<String, CancellationToken>,
    _store: Arc<LlmStore>,
    _file_path: PathBuf,
}

impl CancellableStream {
    fn new(
        inner: Pin<Box<dyn futures_core::Stream<Item = LlmResult<StreamEvent>> + Send>>,
        cancel_token: CancellationToken,
        conv_id: String,
        conversations: DashMap<String, LlmConversation>,
        messages: DashMap<String, Vec<StoredMessage>>,
        active_streams: DashMap<String, CancellationToken>,
        store: Arc<LlmStore>,
        file_path: PathBuf,
    ) -> Self {
        Self {
            inner,
            cancel_token,
            conv_id,
            accumulated_text: String::new(),
            accumulated_thinking: String::new(),
            usage: None,
            model: String::new(),
            finished: false,
            conversations,
            messages,
            active_streams,
            _store: store,
            _file_path: file_path,
        }
    }

    /// Save the accumulated content as a StoredMessage.
    fn save_accumulated(&mut self, status: MessageStatus) {
        if self.finished {
            return;
        }
        self.finished = true;

        let msg = StoredMessage {
            id: uuid::Uuid::new_v4().to_string(),
            conversation_id: self.conv_id.clone(),
            role: MessageRole::Assistant,
            content: std::mem::take(&mut self.accumulated_text),
            content_parts: Vec::new(),
            thinking: if self.accumulated_thinking.is_empty() {
                None
            } else {
                Some(std::mem::take(&mut self.accumulated_thinking))
            },
            tool_calls: Vec::new(),
            tool_call_id: None,
            name: None,
            attachments: Vec::new(),
            token_usage: self.usage.take(),
            duration_ms: None,
            model_id: if self.model.is_empty() {
                None
            } else {
                Some(self.model.clone())
            },
            status,
            created_at: now_ms(),
        };

        // Update conversation counters
        if let Some(ref usage) = msg.token_usage {
            if let Some(mut conv) = self.conversations.get_mut(&self.conv_id) {
                conv.add_message_stats(usage);
            }
        } else if let Some(mut conv) = self.conversations.get_mut(&self.conv_id) {
            conv.message_count += 1;
            conv.touch();
        }

        // Add message
        self.messages
            .entry(self.conv_id.clone())
            .or_default()
            .push(msg);

        // Remove from active streams
        self.active_streams.remove(&self.conv_id);
    }
}

impl futures_core::Stream for CancellableStream {
    type Item = LlmResult<StreamEvent>;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        // Check cancellation
        if self.cancel_token.is_cancelled() {
            self.save_accumulated(MessageStatus::Partial);
            return std::task::Poll::Ready(Some(Ok(StreamEvent::Done {
                message_id: String::new(),
                finish_reason: super::types::FinishReason::Stop,
            })));
        }

        match self.inner.as_mut().poll_next(cx) {
            std::task::Poll::Ready(Some(Ok(event))) => {
                // Accumulate content from events
                match &event {
                    StreamEvent::TextDelta { content } => {
                        self.accumulated_text.push_str(content);
                    }
                    StreamEvent::ThinkingDelta { content } => {
                        self.accumulated_thinking.push_str(content);
                    }
                    StreamEvent::Usage(usage) => {
                        self.usage = Some(usage.clone());
                    }
                    StreamEvent::Done { .. } => {
                        self.save_accumulated(MessageStatus::Sent);
                    }
                    StreamEvent::Error { message, .. } => {
                        warn!(conv_id = %self.conv_id, err = %message, "Stream error");
                        self.save_accumulated(MessageStatus::Error);
                    }
                    _ => {}
                }
                std::task::Poll::Ready(Some(Ok(event)))
            }
            std::task::Poll::Ready(Some(Err(e))) => {
                self.save_accumulated(MessageStatus::Error);
                std::task::Poll::Ready(Some(Err(e)))
            }
            std::task::Poll::Ready(None) => {
                self.save_accumulated(MessageStatus::Sent);
                std::task::Poll::Ready(None)
            }
            std::task::Poll::Pending => std::task::Poll::Pending,
        }
    }
}

// ============================================================================
// Helpers
// ============================================================================

fn now_ms() -> i64 {
    use std::time::{SystemTime, UNIX_EPOCH};
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as i64
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::llm::store::LlmStore;
    use crate::llm::router::{ProviderRouter, RoutingStrategy};

    fn make_test_manager(dir: &Path) -> ConversationManager {
        let store = Arc::new(LlmStore::open(dir.join("store")).unwrap());
        let router = Arc::new(ProviderRouter::new(RoutingStrategy::Priority));
        ConversationManager::new(dir.join("convs"), store, router).unwrap()
    }

    #[test]
    fn test_create_conversation() {
        let tmp = tempfile::tempdir().unwrap();
        let mgr = make_test_manager(tmp.path());

        let conv = mgr
            .create_conversation("user-1", "provider-1", "model-1", "GPT-4o")
            .unwrap();

        assert_eq!(conv.owner_id, "user-1");
        assert_eq!(conv.provider_id, "provider-1");
        assert_eq!(conv.model_id, "model-1");
        assert_eq!(conv.model_name, "GPT-4o");
        assert_eq!(conv.message_count, 0);
        assert_eq!(mgr.conversation_count(), 1);
    }

    #[test]
    fn test_get_conversation_permission() {
        let tmp = tempfile::tempdir().unwrap();
        let mgr = make_test_manager(tmp.path());

        let conv = mgr
            .create_conversation("user-1", "p", "m", "Model")
            .unwrap();

        // Owner can read
        assert!(mgr.get_conversation(&conv.id, "user-1").is_ok());

        // Other user cannot read (private)
        assert!(mgr.get_conversation(&conv.id, "user-2").is_err());
    }

    #[test]
    fn test_list_conversations_visibility() {
        let tmp = tempfile::tempdir().unwrap();
        let mgr = make_test_manager(tmp.path());

        // User 1 creates 2 conversations
        mgr.create_conversation("user-1", "p", "m", "Model").unwrap();
        mgr.create_conversation("user-1", "p", "m", "Model").unwrap();
        // User 2 creates 1 conversation
        mgr.create_conversation("user-2", "p", "m", "Model").unwrap();

        // User 1 sees only their own (private)
        let q = ConversationQuery::for_user("user-1");
        let list = mgr.list_conversations(&q);
        assert_eq!(list.len(), 2);

        // User 2 sees only theirs
        let q = ConversationQuery::for_user("user-2");
        let list = mgr.list_conversations(&q);
        assert_eq!(list.len(), 1);
    }

    #[test]
    fn test_update_conversation() {
        let tmp = tempfile::tempdir().unwrap();
        let mgr = make_test_manager(tmp.path());

        let conv = mgr
            .create_conversation("user-1", "p", "m", "Model")
            .unwrap();

        // Owner can update
        let updated = mgr
            .update_conversation(&conv.id, "user-1", |c| {
                c.title = "Renamed".to_string();
            })
            .unwrap();
        assert_eq!(updated.title, "Renamed");

        // Other user cannot
        assert!(mgr
            .update_conversation(&conv.id, "user-2", |_| {})
            .is_err());
    }

    #[test]
    fn test_delete_conversation() {
        let tmp = tempfile::tempdir().unwrap();
        let mgr = make_test_manager(tmp.path());

        let conv = mgr
            .create_conversation("user-1", "p", "m", "Model")
            .unwrap();

        // Other user cannot delete
        assert!(mgr.delete_conversation(&conv.id, "user-2").is_err());

        // Owner can delete
        assert!(mgr.delete_conversation(&conv.id, "user-1").unwrap());
        assert_eq!(mgr.conversation_count(), 0);
    }

    #[test]
    fn test_add_and_get_messages() {
        let tmp = tempfile::tempdir().unwrap();
        let mgr = make_test_manager(tmp.path());

        let conv = mgr
            .create_conversation("user-1", "p", "m", "Model")
            .unwrap();

        // Add messages
        mgr.add_message(StoredMessage::user(&conv.id, "Hello!")).unwrap();
        mgr.add_message(StoredMessage::user(&conv.id, "How are you?")).unwrap();

        let query = MessageQuery::new(&conv.id);
        let msgs = mgr.get_messages(&query).unwrap();
        assert_eq!(msgs.len(), 2);
        assert_eq!(msgs[0].content, "Hello!");
        assert_eq!(msgs[1].content, "How are you?");
    }

    #[test]
    fn test_message_pagination() {
        let tmp = tempfile::tempdir().unwrap();
        let mgr = make_test_manager(tmp.path());

        let conv = mgr
            .create_conversation("user-1", "p", "m", "Model")
            .unwrap();

        for i in 0..10 {
            mgr.add_message(StoredMessage::user(&conv.id, format!("msg {i}")))
                .unwrap();
        }

        let mut query = MessageQuery::new(&conv.id);
        query.offset = 3;
        query.limit = 4;
        let msgs = mgr.get_messages(&query).unwrap();
        assert_eq!(msgs.len(), 4);
        assert_eq!(msgs[0].content, "msg 3");
    }

    #[test]
    fn test_build_context() {
        let tmp = tempfile::tempdir().unwrap();
        let mgr = make_test_manager(tmp.path());

        let mut conv = LlmConversation::new("user-1", "p", "m", "Model")
            .with_system_prompt("You are helpful.");
        conv.config.context_window = Some(2);
        let conv = mgr.create_conversation_full(conv).unwrap();

        // Add 5 messages
        for i in 0..5 {
            mgr.add_message(StoredMessage::user(&conv.id, format!("msg {i}")))
                .unwrap();
        }

        let context = mgr.build_context(&conv.id).unwrap();
        // system prompt + last 2 messages
        assert_eq!(context.len(), 3);
        assert_eq!(context[0].role, MessageRole::System);
        assert_eq!(context[1].content, "msg 3");
        assert_eq!(context[2].content, "msg 4");
    }

    #[test]
    fn test_build_context_stateless() {
        let tmp = tempfile::tempdir().unwrap();
        let mgr = make_test_manager(tmp.path());

        let mut conv = LlmConversation::new("user-1", "p", "m", "Model")
            .with_system_prompt("You are helpful.");
        conv.config.disable_context = true;
        let conv = mgr.create_conversation_full(conv).unwrap();

        for i in 0..5 {
            mgr.add_message(StoredMessage::user(&conv.id, format!("msg {i}")))
                .unwrap();
        }

        let context = mgr.build_context(&conv.id).unwrap();
        // system prompt + last message only
        assert_eq!(context.len(), 2);
        assert_eq!(context[0].role, MessageRole::System);
        assert_eq!(context[1].content, "msg 4");
    }

    #[test]
    fn test_switch_model() {
        let tmp = tempfile::tempdir().unwrap();
        let mgr = make_test_manager(tmp.path());

        let conv = mgr
            .create_conversation("user-1", "p1", "m1", "Model1")
            .unwrap();

        let updated = mgr
            .switch_model(&conv.id, "user-1", "p2", "m2", "Model2")
            .unwrap();
        assert_eq!(updated.provider_id, "p2");
        assert_eq!(updated.model_id, "m2");
        assert_eq!(updated.model_name, "Model2");
    }

    #[test]
    fn test_toggle_favorite() {
        let tmp = tempfile::tempdir().unwrap();
        let mgr = make_test_manager(tmp.path());

        let conv = mgr
            .create_conversation("user-1", "p", "m", "M")
            .unwrap();

        let is_fav = mgr.toggle_favorite(&conv.id, "user-1").unwrap();
        assert!(is_fav);

        let is_fav = mgr.toggle_favorite(&conv.id, "user-1").unwrap();
        assert!(!is_fav);
    }

    #[test]
    fn test_toggle_archive() {
        let tmp = tempfile::tempdir().unwrap();
        let mgr = make_test_manager(tmp.path());

        let conv = mgr
            .create_conversation("user-1", "p", "m", "M")
            .unwrap();

        let archived = mgr.toggle_archive(&conv.id, "user-1").unwrap();
        assert!(archived);
    }

    #[test]
    fn test_share_conversation() {
        let tmp = tempfile::tempdir().unwrap();
        let mgr = make_test_manager(tmp.path());

        let conv = mgr
            .create_conversation("user-1", "p", "m", "M")
            .unwrap();

        // Share with user-2
        mgr.share_conversation(&conv.id, "user-1", vec!["user-2".into()])
            .unwrap();

        // user-2 can now read
        assert!(mgr.get_conversation(&conv.id, "user-2").is_ok());
        // user-3 cannot
        assert!(mgr.get_conversation(&conv.id, "user-3").is_err());
    }

    #[test]
    fn test_usage_stats() {
        let tmp = tempfile::tempdir().unwrap();
        let mgr = make_test_manager(tmp.path());

        let conv = mgr
            .create_conversation("user-1", "p", "m", "M")
            .unwrap();

        let mut msg1 = StoredMessage::user(&conv.id, "Hi");
        msg1.token_usage = Some(TokenUsage {
            prompt_tokens: 10,
            completion_tokens: 20,
            total_tokens: 30,
            cost: Some(0.001),
            ..Default::default()
        });
        mgr.add_message(msg1).unwrap();

        let mut msg2 = StoredMessage::user(&conv.id, "Again");
        msg2.token_usage = Some(TokenUsage {
            prompt_tokens: 15,
            completion_tokens: 25,
            total_tokens: 40,
            cost: Some(0.002),
            ..Default::default()
        });
        mgr.add_message(msg2).unwrap();

        let stats = mgr.get_usage_stats("user-1");
        assert_eq!(stats.total_tokens, 70);
        assert!((stats.total_cost - 0.003).abs() < f64::EPSILON);
        assert_eq!(stats.message_count, 2);
        assert_eq!(stats.conversation_count, 1);
    }

    #[test]
    fn test_stop_generation() {
        let tmp = tempfile::tempdir().unwrap();
        let mgr = make_test_manager(tmp.path());

        let conv = mgr
            .create_conversation("user-1", "p", "m", "M")
            .unwrap();

        // No active stream → false
        assert!(!mgr.stop_generation(&conv.id));
        assert!(!mgr.is_generating(&conv.id));

        // Simulate an active stream
        let token = CancellationToken::new();
        mgr.active_streams.insert(conv.id.clone(), token);

        assert!(mgr.is_generating(&conv.id));
        assert!(mgr.stop_generation(&conv.id));
        assert!(!mgr.is_generating(&conv.id));
    }

    #[test]
    fn test_persistence() {
        let tmp = tempfile::tempdir().unwrap();

        // Create data
        {
            let mgr = make_test_manager(tmp.path());
            let conv = mgr
                .create_conversation("user-1", "p", "m", "Model")
                .unwrap();
            mgr.add_message(StoredMessage::user(&conv.id, "Hello")).unwrap();
            mgr.add_message(StoredMessage::user(&conv.id, "World")).unwrap();
        }

        // Reload and verify
        {
            let mgr = make_test_manager(tmp.path());
            assert_eq!(mgr.conversation_count(), 1);
            assert_eq!(mgr.message_count(), 2);
        }
    }

    #[test]
    fn test_set_title_and_prompt() {
        let tmp = tempfile::tempdir().unwrap();
        let mgr = make_test_manager(tmp.path());

        let conv = mgr
            .create_conversation("user-1", "p", "m", "M")
            .unwrap();

        let c = mgr.set_title(&conv.id, "user-1", "New Title").unwrap();
        assert_eq!(c.title, "New Title");

        let c = mgr
            .set_system_prompt(&conv.id, "user-1", Some("Be concise.".into()))
            .unwrap();
        assert_eq!(c.system_prompt.as_deref(), Some("Be concise."));
    }
}
