//! RAG Pipeline core types
//!
//! Defines the data structures used throughout the RAG (Retrieval-Augmented Generation)
//! pipeline: configuration, messages, retrieval results, and response types.

use std::collections::HashMap;

use serde::{Deserialize, Serialize};

// ============================================================================
// Configuration
// ============================================================================

/// Configuration for the RAG pipeline.
///
/// Controls retrieval, context assembly, and generation behavior.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RagConfig {
    // ── Retrieval ──────────────────────────────────────────────────
    /// Collections to search in (None = all)
    #[serde(default)]
    pub collections: Option<Vec<String>>,

    /// Maximum number of retrieval results (top-K)
    #[serde(default = "default_top_k")]
    pub top_k: usize,

    /// Weight for vector search (0.0–1.0)
    #[serde(default = "default_vector_weight")]
    pub vector_weight: f64,

    /// Weight for BM25 full-text search (0.0–1.0)
    #[serde(default = "default_bm25_weight")]
    pub bm25_weight: f64,

    /// Minimum relevance score threshold (results below this are discarded)
    #[serde(default)]
    pub min_score: Option<f64>,

    /// Whether to use reranker after initial retrieval
    #[serde(default = "default_true")]
    pub enable_reranker: bool,

    /// Reranker strategy: "rrf" | "linear" | "cross-encoder"
    #[serde(default = "default_reranker_strategy")]
    pub reranker_strategy: String,

    // ── Context Assembly ──────────────────────────────────────────
    /// Maximum tokens allowed for assembled context
    #[serde(default = "default_max_context_tokens")]
    pub max_context_tokens: usize,

    /// Maximum number of chat history turns to include
    #[serde(default = "default_max_history_turns")]
    pub max_history_turns: usize,

    /// System prompt template (use `{context}` as placeholder for retrieved docs)
    #[serde(default = "default_system_template")]
    pub system_template: String,

    /// User prompt template (use `{question}` as placeholder)
    #[serde(default = "default_user_template")]
    pub user_template: String,

    // ── Generation ────────────────────────────────────────────────
    /// Generator backend: "microgpt" | "external"
    #[serde(default = "default_generator_backend")]
    pub generator_backend: String,

    /// Maximum generation tokens
    #[serde(default = "default_max_generation_tokens")]
    pub max_generation_tokens: usize,

    /// Temperature for sampling (0.0 = greedy, 1.0 = creative)
    #[serde(default = "default_temperature")]
    pub temperature: f32,

    /// Whether to include source references in the response
    #[serde(default = "default_true")]
    pub include_sources: bool,
}

fn default_top_k() -> usize {
    10
}
fn default_vector_weight() -> f64 {
    0.7
}
fn default_bm25_weight() -> f64 {
    0.3
}
fn default_true() -> bool {
    true
}
fn default_reranker_strategy() -> String {
    "rrf".to_string()
}
fn default_max_context_tokens() -> usize {
    4096
}
fn default_max_history_turns() -> usize {
    5
}
fn default_system_template() -> String {
    "You are a helpful AI assistant that answers questions based on the provided context.\n\
     Use only the information from the context below to answer. If the context doesn't contain \
     enough information, say so clearly.\n\n\
     ## Context\n\n{context}"
        .to_string()
}
fn default_user_template() -> String {
    "{question}".to_string()
}
fn default_generator_backend() -> String {
    "external".to_string()
}
fn default_max_generation_tokens() -> usize {
    1024
}
fn default_temperature() -> f32 {
    0.7
}

impl Default for RagConfig {
    fn default() -> Self {
        Self {
            collections: None,
            top_k: default_top_k(),
            vector_weight: default_vector_weight(),
            bm25_weight: default_bm25_weight(),
            min_score: None,
            enable_reranker: true,
            reranker_strategy: default_reranker_strategy(),
            max_context_tokens: default_max_context_tokens(),
            max_history_turns: default_max_history_turns(),
            system_template: default_system_template(),
            user_template: default_user_template(),
            generator_backend: default_generator_backend(),
            max_generation_tokens: default_max_generation_tokens(),
            temperature: default_temperature(),
            include_sources: true,
        }
    }
}

// ============================================================================
// Chat Message
// ============================================================================

/// Role in a chat conversation.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ChatRole {
    /// System instructions / prompt
    System,
    /// User query
    User,
    /// Assistant (AI) reply
    Assistant,
}

/// A single message in a chat conversation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatMessage {
    /// Role of the speaker
    pub role: ChatRole,
    /// Message content
    pub content: String,
    /// Timestamp (Unix milliseconds, auto-set if None)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub timestamp: Option<i64>,
}

impl ChatMessage {
    /// Create a system message.
    pub fn system(content: impl Into<String>) -> Self {
        Self {
            role: ChatRole::System,
            content: content.into(),
            timestamp: Some(now_ms()),
        }
    }

    /// Create a user message.
    pub fn user(content: impl Into<String>) -> Self {
        Self {
            role: ChatRole::User,
            content: content.into(),
            timestamp: Some(now_ms()),
        }
    }

    /// Create an assistant message.
    pub fn assistant(content: impl Into<String>) -> Self {
        Self {
            role: ChatRole::Assistant,
            content: content.into(),
            timestamp: Some(now_ms()),
        }
    }

    /// Approximate token count (rough: 1 token ≈ 4 chars for English).
    pub fn approx_tokens(&self) -> usize {
        // Simple heuristic; for CJK text, 1 token ≈ 1.5 chars
        let char_count = self.content.len();
        (char_count + 3) / 4
    }
}

/// Convert from existing `perception::Message` for interop.
impl From<&crate::perception::Message> for ChatMessage {
    fn from(msg: &crate::perception::Message) -> Self {
        let role = match msg.role {
            crate::perception::MessageRole::System => ChatRole::System,
            crate::perception::MessageRole::User => ChatRole::User,
            crate::perception::MessageRole::Assistant => ChatRole::Assistant,
        };
        Self {
            role,
            content: msg.content.clone(),
            timestamp: None,
        }
    }
}

/// Convert to `perception::Message` for ReasoningEngine calls.
impl From<&ChatMessage> for crate::perception::Message {
    fn from(msg: &ChatMessage) -> Self {
        match msg.role {
            ChatRole::System => Self::system(&msg.content),
            ChatRole::User => Self::user(&msg.content),
            ChatRole::Assistant => Self::assistant(&msg.content),
        }
    }
}

// ============================================================================
// Retrieval Result
// ============================================================================

/// A single document retrieved by the Retriever.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RetrievalResult {
    /// Document / node ID
    pub id: String,

    /// Main text content
    pub content: String,

    /// Relevance score (higher = more relevant, typically 0.0–1.0)
    pub score: f64,

    /// Collection this document belongs to
    #[serde(skip_serializing_if = "Option::is_none")]
    pub collection: Option<String>,

    /// Score breakdown by channel (e.g., "vector": 0.85, "bm25": 0.6)
    #[serde(default)]
    pub score_details: HashMap<String, f64>,

    /// Arbitrary metadata from the document
    #[serde(default)]
    pub metadata: HashMap<String, serde_json::Value>,
}

impl RetrievalResult {
    /// Create a simple retrieval result.
    pub fn new(id: impl Into<String>, content: impl Into<String>, score: f64) -> Self {
        Self {
            id: id.into(),
            content: content.into(),
            score,
            collection: None,
            score_details: HashMap::new(),
            metadata: HashMap::new(),
        }
    }

    /// Set collection name.
    pub fn with_collection(mut self, collection: impl Into<String>) -> Self {
        self.collection = Some(collection.into());
        self
    }
}

// ============================================================================
// Source Reference
// ============================================================================

/// A reference to a source document included in the RAG response.
///
/// Provides provenance information so the user can verify the answer.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SourceReference {
    /// Document / node ID
    pub id: String,

    /// Short excerpt from the source (first ~200 chars)
    pub excerpt: String,

    /// Relevance score (from retrieval)
    pub score: f64,

    /// Collection name
    #[serde(skip_serializing_if = "Option::is_none")]
    pub collection: Option<String>,

    /// Additional metadata
    #[serde(default)]
    pub metadata: HashMap<String, serde_json::Value>,
}

impl From<&RetrievalResult> for SourceReference {
    fn from(r: &RetrievalResult) -> Self {
        let excerpt = if r.content.len() > 200 {
            format!("{}...", &r.content[..200])
        } else {
            r.content.clone()
        };
        Self {
            id: r.id.clone(),
            excerpt,
            score: r.score,
            collection: r.collection.clone(),
            metadata: r.metadata.clone(),
        }
    }
}

// ============================================================================
// RAG Response
// ============================================================================

/// The final response from the RAG pipeline.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RagResponse {
    /// Generated answer text
    pub answer: String,

    /// Source references (provenance for the answer)
    #[serde(default)]
    pub sources: Vec<SourceReference>,

    /// Session ID (for multi-turn conversations)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub session_id: Option<String>,

    /// Token usage statistics
    #[serde(skip_serializing_if = "Option::is_none")]
    pub usage: Option<RagTokenUsage>,

    /// Number of documents retrieved
    pub retrieval_count: usize,

    /// Total processing time in milliseconds
    #[serde(skip_serializing_if = "Option::is_none")]
    pub latency_ms: Option<u64>,

    /// Additional metadata (e.g., model name, reranker used)
    #[serde(default)]
    pub metadata: HashMap<String, serde_json::Value>,
}

impl RagResponse {
    /// Create a minimal response.
    pub fn new(answer: impl Into<String>) -> Self {
        Self {
            answer: answer.into(),
            sources: Vec::new(),
            session_id: None,
            usage: None,
            retrieval_count: 0,
            latency_ms: None,
            metadata: HashMap::new(),
        }
    }
}

/// Token usage information for the RAG pipeline.
///
/// Extended in Phase-82 (T82.3) with `cost` and `reasoning_tokens` fields
/// for unified tracking with LLM Gateway's `TokenUsage`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RagTokenUsage {
    /// Tokens in the prompt (system + context + history + question)
    pub prompt_tokens: usize,
    /// Tokens generated in the answer
    pub completion_tokens: usize,
    /// Total tokens
    pub total_tokens: usize,
    /// Approximate tokens used for context (retrieved documents)
    pub context_tokens: usize,
    /// Cost of this request (from LLM Gateway's `ModelPricing`).
    ///
    /// Populated when using `LlmGatewayGenerator` and the model has pricing
    /// configured. Enables unified cost tracking via `ConversationManager`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cost: Option<f64>,
    /// Reasoning/thinking tokens (for deep reasoning models like o1, QwQ).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub reasoning_tokens: Option<usize>,
}

// ============================================================================
// Phase-82 T82.3: RagTokenUsage ↔ LLM Gateway TokenUsage conversions
// ============================================================================

#[cfg(feature = "llm-gateway")]
impl From<&crate::llm::types::TokenUsage> for RagTokenUsage {
    /// Convert from LLM Gateway `TokenUsage` → `RagTokenUsage`.
    ///
    /// `context_tokens` is set to 0 and should be filled by the pipeline.
    fn from(t: &crate::llm::types::TokenUsage) -> Self {
        Self {
            prompt_tokens: t.prompt_tokens as usize,
            completion_tokens: t.completion_tokens as usize,
            total_tokens: t.total_tokens as usize,
            context_tokens: 0, // Filled by RagPipeline
            cost: t.cost,
            reasoning_tokens: t.reasoning_tokens.map(|v| v as usize),
        }
    }
}

#[cfg(feature = "llm-gateway")]
impl From<&RagTokenUsage> for crate::llm::types::TokenUsage {
    /// Convert from `RagTokenUsage` → LLM Gateway `TokenUsage`.
    ///
    /// This allows RAG usage data to be fed into `ConversationManager.add_message()`
    /// for centralized token/cost tracking.
    fn from(r: &RagTokenUsage) -> Self {
        Self {
            prompt_tokens: r.prompt_tokens as u32,
            completion_tokens: r.completion_tokens as u32,
            total_tokens: r.total_tokens as u32,
            cost: r.cost,
            reasoning_tokens: r.reasoning_tokens.map(|v| v as u32),
            cached_tokens: None,
        }
    }
}

// ============================================================================
// Chat Session
// ============================================================================

/// A chat session holding conversation history.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatSession {
    /// Unique session identifier
    pub session_id: String,

    /// Conversation history (ordered by time)
    pub messages: Vec<ChatMessage>,

    /// Creation timestamp (Unix milliseconds)
    pub created_at: i64,

    /// Last activity timestamp (Unix milliseconds)
    pub updated_at: i64,

    /// Custom session metadata
    #[serde(default)]
    pub metadata: HashMap<String, serde_json::Value>,
}

impl ChatSession {
    /// Create a new empty session.
    pub fn new(session_id: impl Into<String>) -> Self {
        let now = now_ms();
        Self {
            session_id: session_id.into(),
            messages: Vec::new(),
            created_at: now,
            updated_at: now,
            metadata: HashMap::new(),
        }
    }

    /// Add a message to the session.
    pub fn push(&mut self, message: ChatMessage) {
        self.updated_at = now_ms();
        self.messages.push(message);
    }

    /// Get the last N user+assistant turns (excludes system messages).
    pub fn recent_turns(&self, max_turns: usize) -> Vec<&ChatMessage> {
        let non_system: Vec<&ChatMessage> = self
            .messages
            .iter()
            .filter(|m| m.role != ChatRole::System)
            .collect();

        // Each "turn" is a (user, assistant) pair = 2 messages
        let keep = max_turns * 2;
        if non_system.len() <= keep {
            non_system
        } else {
            non_system[non_system.len() - keep..].to_vec()
        }
    }

    /// Approximate total tokens in the session.
    pub fn total_tokens(&self) -> usize {
        self.messages.iter().map(|m| m.approx_tokens()).sum()
    }
}

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

/// Current time in Unix milliseconds.
fn now_ms() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as i64
}

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

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

    #[test]
    fn test_chat_message_creation() {
        let msg = ChatMessage::system("You are helpful.");
        assert_eq!(msg.role, ChatRole::System);
        assert!(msg.timestamp.is_some());

        let msg = ChatMessage::user("Hello?");
        assert_eq!(msg.role, ChatRole::User);

        let msg = ChatMessage::assistant("Hi there!");
        assert_eq!(msg.role, ChatRole::Assistant);
    }

    #[test]
    fn test_chat_message_approx_tokens() {
        let msg = ChatMessage::user("Hello world"); // 11 chars → ~3 tokens
        assert!(msg.approx_tokens() > 0);
    }

    #[test]
    fn test_chat_session_recent_turns() {
        let mut session = ChatSession::new("test-session");
        session.push(ChatMessage::system("System prompt"));
        session.push(ChatMessage::user("Q1"));
        session.push(ChatMessage::assistant("A1"));
        session.push(ChatMessage::user("Q2"));
        session.push(ChatMessage::assistant("A2"));
        session.push(ChatMessage::user("Q3"));
        session.push(ChatMessage::assistant("A3"));

        // Last 2 turns = 4 messages (Q2,A2,Q3,A3)
        let recent = session.recent_turns(2);
        assert_eq!(recent.len(), 4);
        assert_eq!(recent[0].content, "Q2");
        assert_eq!(recent[3].content, "A3");

        // Last 10 turns — should return all non-system (6 messages)
        let all = session.recent_turns(10);
        assert_eq!(all.len(), 6);
    }

    #[test]
    fn test_retrieval_result() {
        let r = RetrievalResult::new("doc-1", "Some content about Rust", 0.95)
            .with_collection("default");
        assert_eq!(r.id, "doc-1");
        assert_eq!(r.score, 0.95);
        assert_eq!(r.collection.as_deref(), Some("default"));
    }

    #[test]
    fn test_source_reference_from_retrieval() {
        let r = RetrievalResult::new("doc-1", "A".repeat(300), 0.9);
        let src: SourceReference = SourceReference::from(&r);
        assert!(src.excerpt.len() < 210); // truncated to ~200 + "..."
        assert!(src.excerpt.ends_with("..."));
    }

    #[test]
    fn test_rag_response() {
        let resp = RagResponse::new("Based on the context, the answer is 42.");
        assert!(!resp.answer.is_empty());
        assert!(resp.sources.is_empty());
    }

    #[test]
    fn test_rag_config_default() {
        let config = RagConfig::default();
        assert_eq!(config.top_k, 10);
        assert_eq!(config.vector_weight, 0.7);
        assert_eq!(config.bm25_weight, 0.3);
        assert!(config.enable_reranker);
        assert_eq!(config.reranker_strategy, "rrf");
        assert!(config.system_template.contains("{context}"));
    }

    #[test]
    fn test_chat_message_interop() {
        use crate::perception::{Message, MessageRole};

        // perception::Message → ChatMessage
        let pm = Message::user("hello");
        let cm = ChatMessage::from(&pm);
        assert_eq!(cm.role, ChatRole::User);
        assert_eq!(cm.content, "hello");

        // ChatMessage → perception::Message
        let cm = ChatMessage::assistant("world");
        let pm: Message = Message::from(&cm);
        assert!(matches!(pm.role, MessageRole::Assistant));
        assert_eq!(pm.content, "world");
    }
}
