//! Dialogue system for conversational AI
//!
//! Provides QA and multi-turn conversation support with context management.

use serde::{Deserialize, Serialize};
use std::collections::VecDeque;

/// Dialogue context for multi-turn conversations
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct DialogueContext {
    /// Conversation history
    pub history: VecDeque<DialogueTurn>,
    /// Maximum number of turns to keep in history
    pub max_history: usize,
    /// System prompt/context
    pub system_prompt: Option<String>,
    /// Session metadata
    pub metadata: DialogueMetadata,
}

/// A single turn in a dialogue
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DialogueTurn {
    /// Role of the speaker
    pub role: Role,
    /// Content of the message
    pub content: String,
    /// Timestamp of the turn
    pub timestamp: chrono::DateTime<chrono::Utc>,
    /// Optional metadata for this turn
    pub metadata: Option<TurnMetadata>,
}

/// Speaker role in a conversation
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum Role {
    /// User/human input
    User,
    /// AI assistant response
    Assistant,
    /// System message/prompt
    System,
    /// Tool/function call result
    Tool,
}

/// Metadata for dialogue sessions
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct DialogueMetadata {
    /// Session identifier
    pub session_id: Option<String>,
    /// User identifier
    pub user_id: Option<String>,
    /// Additional custom attributes
    pub attributes: std::collections::HashMap<String, String>,
}

/// Metadata for individual turns
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct TurnMetadata {
    /// Generation confidence score
    pub confidence: Option<f32>,
    /// Number of tokens in this turn
    pub token_count: Option<usize>,
    /// Processing time in milliseconds
    pub latency_ms: Option<f64>,
}

impl DialogueContext {
    /// Create a new dialogue context with specified history limit
    pub fn new(max_history: usize) -> Self {
        Self {
            history: VecDeque::new(),
            max_history,
            system_prompt: None,
            metadata: DialogueMetadata::default(),
        }
    }

    /// Create a new dialogue context with a system prompt
    pub fn with_system_prompt(max_history: usize, system_prompt: impl Into<String>) -> Self {
        Self {
            history: VecDeque::new(),
            max_history,
            system_prompt: Some(system_prompt.into()),
            metadata: DialogueMetadata::default(),
        }
    }

    /// Add a turn to the conversation
    pub fn add_turn(&mut self, role: Role, content: impl Into<String>) {
        let turn = DialogueTurn {
            role,
            content: content.into(),
            timestamp: chrono::Utc::now(),
            metadata: None,
        };
        self.history.push_back(turn);

        // Trim history if needed
        while self.history.len() > self.max_history {
            self.history.pop_front();
        }
    }

    /// Add a turn with metadata
    pub fn add_turn_with_metadata(
        &mut self,
        role: Role,
        content: impl Into<String>,
        metadata: TurnMetadata,
    ) {
        let turn = DialogueTurn {
            role,
            content: content.into(),
            timestamp: chrono::Utc::now(),
            metadata: Some(metadata),
        };
        self.history.push_back(turn);

        while self.history.len() > self.max_history {
            self.history.pop_front();
        }
    }

    /// Add a user message
    pub fn add_user_message(&mut self, content: impl Into<String>) {
        self.add_turn(Role::User, content);
    }

    /// Add an assistant response
    pub fn add_assistant_message(&mut self, content: impl Into<String>) {
        self.add_turn(Role::Assistant, content);
    }

    /// Clear conversation history
    pub fn clear(&mut self) {
        self.history.clear();
    }

    /// Get the last N turns
    pub fn last_n_turns(&self, n: usize) -> impl Iterator<Item = &DialogueTurn> {
        self.history.iter().rev().take(n).rev()
    }

    /// Format the conversation for model input
    pub fn format_for_model(&self) -> String {
        let mut formatted = String::new();

        if let Some(ref system) = self.system_prompt {
            formatted.push_str(&format!("System: {system}\n\n"));
        }

        for turn in &self.history {
            let role_str = match turn.role {
                Role::User => "User",
                Role::Assistant => "Assistant",
                Role::System => "System",
                Role::Tool => "Tool",
            };
            formatted.push_str(&format!("{}: {}\n", role_str, turn.content));
        }

        formatted
    }

    /// Get total token estimate (rough approximation)
    pub fn estimate_tokens(&self) -> usize {
        let mut total = 0;

        if let Some(ref system) = self.system_prompt {
            total += system.split_whitespace().count() * 4 / 3; // Rough token estimate
        }

        for turn in &self.history {
            total += turn.content.split_whitespace().count() * 4 / 3;
        }

        total
    }

    /// Check if the conversation is empty
    pub fn is_empty(&self) -> bool {
        self.history.is_empty()
    }

    /// Get the number of turns
    pub fn len(&self) -> usize {
        self.history.len()
    }
}

impl DialogueTurn {
    /// Create a new dialogue turn
    pub fn new(role: Role, content: impl Into<String>) -> Self {
        Self {
            role,
            content: content.into(),
            timestamp: chrono::Utc::now(),
            metadata: None,
        }
    }
}

/// Question-Answering context and utilities
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QAContext {
    /// The question being asked
    pub question: String,
    /// Relevant context passages
    pub context_passages: Vec<String>,
    /// Expected answer type
    pub answer_type: AnswerType,
}

/// Types of expected answers
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum AnswerType {
    /// Short extractive answer
    Extractive,
    /// Longer abstractive answer
    Abstractive,
    /// Yes/No answer
    Boolean,
    /// Numeric answer
    Numeric,
    /// List/enumeration
    List,
}

impl QAContext {
    /// Create a new QA context
    pub fn new(question: impl Into<String>) -> Self {
        Self {
            question: question.into(),
            context_passages: Vec::new(),
            answer_type: AnswerType::Abstractive,
        }
    }

    /// Add a context passage
    pub fn with_context(mut self, passage: impl Into<String>) -> Self {
        self.context_passages.push(passage.into());
        self
    }

    /// Set the expected answer type
    pub fn with_answer_type(mut self, answer_type: AnswerType) -> Self {
        self.answer_type = answer_type;
        self
    }

    /// Format for model input
    pub fn format_for_model(&self) -> String {
        let mut formatted = String::new();

        if !self.context_passages.is_empty() {
            formatted.push_str("Context:\n");
            for (i, passage) in self.context_passages.iter().enumerate() {
                formatted.push_str(&format!("[{}] {}\n", i + 1, passage));
            }
            formatted.push('\n');
        }

        formatted.push_str(&format!("Question: {}\n", self.question));
        formatted.push_str("Answer:");

        formatted
    }
}

/// Conversation manager for handling multiple dialogue sessions
pub struct ConversationTracker {
    /// Active conversations by session ID
    sessions: std::collections::HashMap<String, DialogueContext>,
    /// Default max history for new sessions
    default_max_history: usize,
}

impl ConversationTracker {
    /// Create a new conversation tracker
    pub fn new(default_max_history: usize) -> Self {
        Self {
            sessions: std::collections::HashMap::new(),
            default_max_history,
        }
    }

    /// Get or create a session
    pub fn get_or_create(&mut self, session_id: &str) -> &mut DialogueContext {
        self.sessions
            .entry(session_id.to_string())
            .or_insert_with(|| DialogueContext::new(self.default_max_history))
    }

    /// Get an existing session
    pub fn get(&self, session_id: &str) -> Option<&DialogueContext> {
        self.sessions.get(session_id)
    }

    /// Get a mutable reference to an existing session
    pub fn get_mut(&mut self, session_id: &str) -> Option<&mut DialogueContext> {
        self.sessions.get_mut(session_id)
    }

    /// Remove a session
    pub fn remove(&mut self, session_id: &str) -> Option<DialogueContext> {
        self.sessions.remove(session_id)
    }

    /// Get all active session IDs
    pub fn session_ids(&self) -> impl Iterator<Item = &String> {
        self.sessions.keys()
    }

    /// Get the number of active sessions
    pub fn session_count(&self) -> usize {
        self.sessions.len()
    }
}

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

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

    #[test]
    fn test_dialogue_context() {
        let mut ctx = DialogueContext::new(5);

        ctx.add_user_message("Hello");
        ctx.add_assistant_message("Hi there!");

        assert_eq!(ctx.len(), 2);
        assert!(!ctx.is_empty());
    }

    #[test]
    fn test_max_history() {
        let mut ctx = DialogueContext::new(3);

        for i in 0..5 {
            ctx.add_user_message(format!("Message {}", i));
        }

        assert_eq!(ctx.len(), 3);
    }

    #[test]
    fn test_qa_context() {
        let qa = QAContext::new("What is Rust?")
            .with_context("Rust is a systems programming language.")
            .with_answer_type(AnswerType::Abstractive);

        let formatted = qa.format_for_model();
        assert!(formatted.contains("What is Rust?"));
        assert!(formatted.contains("systems programming"));
    }
}