//! Context Assembler — 将检索结果组装为 LLM prompt
//!
//! 负责：
//! 1. 将检索到的文档格式化为上下文字符串
//! 2. 应用 token 预算约束（截断过长内容）
//! 3. 拼接 system prompt + 上下文 + 对话历史 + 用户问题
//!
//! # Token Budget
//!
//! ```text
//! ┌──────────────────────────────────────────────┐
//! │              max_context_tokens               │
//! ├──────────┬──────────────┬─────────────────────┤
//! │  System  │  Retrieved   │  History + Question  │
//! │  Prompt  │  Documents   │                      │
//! │ (~fixed) │ (truncated)  │    (windowed)        │
//! └──────────┴──────────────┴─────────────────────┘
//! ```

use super::error::RagResult;
use super::traits::ContextAssembler;
use super::types::{ChatMessage, ChatRole, RagConfig, RetrievalResult};

// ============================================================================
// DefaultContextAssembler
// ============================================================================

/// Default implementation of `ContextAssembler`.
///
/// Formats retrieved documents with numbered references, applies token budget,
/// and constructs the final message list for the generator.
pub struct DefaultContextAssembler;

impl DefaultContextAssembler {
    /// Create a new assembler.
    pub fn new() -> Self {
        Self
    }

    /// Format a single retrieval result as a context block.
    fn format_document(idx: usize, doc: &RetrievalResult) -> String {
        let collection_info = doc
            .collection
            .as_deref()
            .map(|c| format!(" [{}]", c))
            .unwrap_or_default();

        format!(
            "[{}]{} (score: {:.2})\n{}",
            idx + 1,
            collection_info,
            doc.score,
            doc.content
        )
    }

    /// Estimate token count for a string.
    ///
    /// Uses a simple heuristic: ~4 chars per token for English,
    /// ~1.5 chars per token for CJK.
    pub(crate) fn estimate_tokens(text: &str) -> usize {
        // Count CJK characters for better estimate
        let cjk_count = text.chars().filter(|c| is_cjk(*c)).count();
        let non_cjk_len = text.len() - cjk_count * 3; // UTF-8 CJK is ~3 bytes
        let cjk_tokens = (cjk_count as f64 / 1.5).ceil() as usize;
        let non_cjk_tokens = (non_cjk_len + 3) / 4;
        cjk_tokens + non_cjk_tokens
    }

    /// Truncate text to fit within a token budget.
    fn truncate_to_tokens(text: &str, max_tokens: usize) -> String {
        let est = Self::estimate_tokens(text);
        if est <= max_tokens {
            return text.to_string();
        }

        // Approximate char limit
        let ratio = max_tokens as f64 / est as f64;
        let char_limit = (text.len() as f64 * ratio * 0.95) as usize; // 5% safety margin
        if char_limit >= text.len() {
            return text.to_string();
        }

        // Truncate at char boundary
        let truncated: String = text.chars().take(char_limit).collect();
        format!("{}...[truncated]", truncated)
    }

    /// Build the context string from retrieved documents within token budget.
    pub(crate) fn build_context(docs: &[RetrievalResult], max_tokens: usize) -> (String, usize) {
        let mut parts: Vec<String> = Vec::new();
        let mut total_tokens = 0;

        for (idx, doc) in docs.iter().enumerate() {
            let block = Self::format_document(idx, doc);
            let block_tokens = Self::estimate_tokens(&block);

            if total_tokens + block_tokens > max_tokens && !parts.is_empty() {
                // Exceeded budget — try to fit partial
                let remaining = max_tokens.saturating_sub(total_tokens);
                if remaining > 50 {
                    // Enough room for a meaningful truncation
                    let truncated = Self::truncate_to_tokens(&block, remaining);
                    let truncated_tokens = Self::estimate_tokens(&truncated);
                    parts.push(truncated);
                    total_tokens += truncated_tokens;
                }
                break;
            }

            total_tokens += block_tokens;
            parts.push(block);
        }

        let context = parts.join("\n\n---\n\n");
        (context, total_tokens)
    }
}

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

impl ContextAssembler for DefaultContextAssembler {
    fn assemble(
        &self,
        question: &str,
        retrieved: &[RetrievalResult],
        history: &[ChatMessage],
        config: &RagConfig,
    ) -> RagResult<(Vec<ChatMessage>, usize)> {
        let max_tokens = config.max_context_tokens;

        // 1. Estimate fixed overhead: system template shell + user question
        let template_shell_tokens =
            Self::estimate_tokens(&config.system_template.replace("{context}", ""));
        let question_tokens = Self::estimate_tokens(question);

        // 2. Estimate history tokens
        let history_window: Vec<&ChatMessage> = history
            .iter()
            .filter(|m| m.role != ChatRole::System)
            .collect::<Vec<_>>();
        let history_msgs: Vec<&ChatMessage> = if history_window.len() > config.max_history_turns * 2
        {
            history_window[history_window.len() - config.max_history_turns * 2..].to_vec()
        } else {
            history_window
        };
        let history_tokens: usize = history_msgs.iter().map(|m| m.approx_tokens()).sum();

        // 3. Calculate budget for context documents
        let overhead = template_shell_tokens + question_tokens + history_tokens + 20; // 20 token margin
        let context_budget = max_tokens.saturating_sub(overhead);

        // 4. Build context from retrieved documents
        let (context_str, context_tokens) = Self::build_context(retrieved, context_budget);

        // 5. Assemble messages
        let mut messages: Vec<ChatMessage> = Vec::new();

        // System message with context
        let system_content = config.system_template.replace("{context}", &context_str);
        messages.push(ChatMessage::system(system_content));

        // History messages
        for msg in &history_msgs {
            messages.push((*msg).clone());
        }

        // User question
        let user_content = config.user_template.replace("{question}", question);
        messages.push(ChatMessage::user(user_content));

        let total_tokens = context_tokens + overhead;

        Ok((messages, total_tokens))
    }
}

// ============================================================================
// TemplateAssembler (Phase-82: T82.5)
// ============================================================================

/// Context assembler that resolves prompts from `PromptTemplateStore`.
///
/// Instead of using the hardcoded `system_template` / `user_template` from
/// `RagConfig`, this assembler looks up templates by ID from the centralized
/// `PromptTemplateStore`. Falls back to `RagConfig` templates if the store
/// doesn't have the specified template.
///
/// # Template Variables
///
/// The following variables are automatically injected:
/// - `{{context}}` — assembled retrieved documents
/// - `{{question}}` — user's question
/// - `{{history_count}}` — number of history messages
///
/// # Requires
///
/// - `llm-gateway` feature enabled
///
/// # Example
///
/// ```ignore
/// use aifastdb_core::rag::assembler::TemplateAssembler;
///
/// let assembler = TemplateAssembler::new(template_store)
///     .with_system_template("rag-strict")
///     .with_user_template("rag-user-custom");
///
/// let pipeline = RagPipeline::builder()
///     .assembler(assembler)
///     // ...
///     .build()?;
/// ```
#[cfg(feature = "llm-gateway")]
pub struct TemplateAssembler {
    /// Shared reference to the prompt template store.
    store: std::sync::Arc<crate::llm::prompt_template::PromptTemplateStore>,
    /// Template ID for the system prompt (default: "rag-default").
    system_template_id: String,
    /// Optional user template ID (if None, uses `RagConfig.user_template`).
    user_template_id: Option<String>,
}

#[cfg(feature = "llm-gateway")]
impl TemplateAssembler {
    /// Create a new template-based assembler.
    ///
    /// By default, uses the "rag-default" template from the store.
    pub fn new(
        store: std::sync::Arc<crate::llm::prompt_template::PromptTemplateStore>,
    ) -> Self {
        Self {
            store,
            system_template_id: "rag-default".to_string(),
            user_template_id: None,
        }
    }

    /// Set the system template ID (looked up in `PromptTemplateStore`).
    pub fn with_system_template(mut self, template_id: impl Into<String>) -> Self {
        self.system_template_id = template_id.into();
        self
    }

    /// Set the user template ID (looked up in `PromptTemplateStore`).
    pub fn with_user_template(mut self, template_id: impl Into<String>) -> Self {
        self.user_template_id = Some(template_id.into());
        self
    }

    /// Resolve and render a system template with the given context.
    fn render_system_template(
        &self,
        context: &str,
        question: &str,
        config: &RagConfig,
    ) -> String {
        let mut vars = std::collections::HashMap::new();
        vars.insert("context".to_string(), context.to_string());
        vars.insert("question".to_string(), question.to_string());

        // Try to render from store
        match self.store.render_partial(&self.system_template_id, &vars) {
            Ok(rendered) => rendered,
            Err(e) => {
                tracing::warn!(
                    template_id = %self.system_template_id,
                    error = %e,
                    "Template not found, falling back to RagConfig.system_template"
                );
                config.system_template.replace("{context}", context)
            }
        }
    }

    /// Resolve and render a user template.
    fn render_user_template(&self, question: &str, config: &RagConfig) -> String {
        if let Some(ref template_id) = self.user_template_id {
            let mut vars = std::collections::HashMap::new();
            vars.insert("question".to_string(), question.to_string());

            match self.store.render_partial(template_id, &vars) {
                Ok(rendered) => return rendered,
                Err(e) => {
                    tracing::warn!(
                        template_id = %template_id,
                        error = %e,
                        "User template not found, falling back to RagConfig.user_template"
                    );
                }
            }
        }

        config.user_template.replace("{question}", question)
    }
}

#[cfg(feature = "llm-gateway")]
impl ContextAssembler for TemplateAssembler {
    fn assemble(
        &self,
        question: &str,
        retrieved: &[RetrievalResult],
        history: &[ChatMessage],
        config: &RagConfig,
    ) -> RagResult<(Vec<ChatMessage>, usize)> {
        let max_tokens = config.max_context_tokens;

        // 1. Estimate template overhead (without context)
        let template_shell_tokens =
            DefaultContextAssembler::estimate_tokens(
                &self.render_system_template("", question, config),
            );
        let question_tokens =
            DefaultContextAssembler::estimate_tokens(question);

        // 2. History windowing
        let history_window: Vec<&ChatMessage> = history
            .iter()
            .filter(|m| m.role != ChatRole::System)
            .collect();
        let history_msgs: Vec<&ChatMessage> =
            if history_window.len() > config.max_history_turns * 2 {
                history_window[history_window.len() - config.max_history_turns * 2..].to_vec()
            } else {
                history_window
            };
        let history_tokens: usize = history_msgs.iter().map(|m| m.approx_tokens()).sum();

        // 3. Context budget
        let overhead = template_shell_tokens + question_tokens + history_tokens + 20;
        let context_budget = max_tokens.saturating_sub(overhead);

        // 4. Build context string
        let (context_str, context_tokens) =
            DefaultContextAssembler::build_context(retrieved, context_budget);

        // 5. Assemble messages using templates from the store
        let mut messages: Vec<ChatMessage> = Vec::new();

        let system_content = self.render_system_template(&context_str, question, config);
        messages.push(ChatMessage::system(system_content));

        for msg in &history_msgs {
            messages.push((*msg).clone());
        }

        let user_content = self.render_user_template(question, config);
        messages.push(ChatMessage::user(user_content));

        let total_tokens = context_tokens + overhead;

        Ok((messages, total_tokens))
    }
}

/// Check if a character is CJK (Chinese/Japanese/Korean).
fn is_cjk(c: char) -> bool {
    matches!(c,
        '\u{4E00}'..='\u{9FFF}'   // CJK Unified Ideographs
        | '\u{3400}'..='\u{4DBF}' // CJK Extension A
        | '\u{F900}'..='\u{FAFF}' // CJK Compatibility Ideographs
        | '\u{3000}'..='\u{303F}' // CJK Symbols and Punctuation
        | '\u{3040}'..='\u{309F}' // Hiragana
        | '\u{30A0}'..='\u{30FF}' // Katakana
        | '\u{AC00}'..='\u{D7AF}' // Hangul Syllables
    )
}

// ============================================================================
// GraphRagAssembler (Phase-86: T86.5)
// ============================================================================

/// Graph-aware context assembler for Graph RAG.
///
/// Extends `DefaultContextAssembler` with dedicated sections for:
/// - **Entity descriptions**: Structured entity information
/// - **Relationship list**: Formatted graph relationships
/// - **Community summaries**: Pre-generated sub-graph summaries
///
/// Token budget allocation:
/// - 60% for retrieved documents (vector/BM25)
/// - 25% for graph context (entities + relations)
/// - 15% for community summaries
pub struct GraphRagAssembler {
    /// Fraction of context budget for document content (default 0.60)
    doc_budget_ratio: f64,
    /// Fraction of context budget for graph context (default 0.25)
    graph_budget_ratio: f64,
    /// Fraction of context budget for community summaries (default 0.15)
    community_budget_ratio: f64,
}

impl GraphRagAssembler {
    /// Create a new GraphRagAssembler with default budget ratios.
    pub fn new() -> Self {
        Self {
            doc_budget_ratio: 0.60,
            graph_budget_ratio: 0.25,
            community_budget_ratio: 0.15,
        }
    }

    /// Set custom budget ratios (will be normalized).
    pub fn with_ratios(mut self, doc: f64, graph: f64, community: f64) -> Self {
        let total = doc + graph + community;
        if total > 0.0 {
            self.doc_budget_ratio = doc / total;
            self.graph_budget_ratio = graph / total;
            self.community_budget_ratio = community / total;
        }
        self
    }

    /// Partition retrieval results into categories by source.
    fn partition_results(
        results: &[RetrievalResult],
    ) -> (Vec<&RetrievalResult>, Vec<&RetrievalResult>, Vec<&RetrievalResult>) {
        let mut doc_results = Vec::new();
        let mut graph_results = Vec::new();
        let mut community_results = Vec::new();

        for r in results {
            let source = r
                .metadata
                .get("source")
                .and_then(|v| v.as_str())
                .unwrap_or("");
            let collection = r.collection.as_deref().unwrap_or("");

            if source == "community_summary" || collection == "graph_community" {
                community_results.push(r);
            } else if source == "graph_retriever"
                || source == "graph_context_expander"
                || collection == "graph"
            {
                graph_results.push(r);
            } else {
                doc_results.push(r);
            }
        }

        (doc_results, graph_results, community_results)
    }

    /// Build context string for a slice of results within a token budget.
    fn build_section(
        header: &str,
        results: &[&RetrievalResult],
        max_tokens: usize,
    ) -> (String, usize) {
        if results.is_empty() || max_tokens < 20 {
            return (String::new(), 0);
        }

        let header_tokens = DefaultContextAssembler::estimate_tokens(header);
        let remaining = max_tokens.saturating_sub(header_tokens + 2);

        // Build content from results
        let owned: Vec<RetrievalResult> = results.iter().map(|r| (*r).clone()).collect();
        let (content, content_tokens) = DefaultContextAssembler::build_context(&owned, remaining);

        if content.is_empty() {
            return (String::new(), 0);
        }

        let section = format!("{}\n{}", header, content);
        let total = header_tokens + content_tokens;
        (section, total)
    }
}

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

impl ContextAssembler for GraphRagAssembler {
    fn assemble(
        &self,
        question: &str,
        retrieved: &[RetrievalResult],
        history: &[ChatMessage],
        config: &RagConfig,
    ) -> RagResult<(Vec<ChatMessage>, usize)> {
        let max_tokens = config.max_context_tokens;

        // 1. Estimate fixed overhead
        let template_shell_tokens =
            DefaultContextAssembler::estimate_tokens(&config.system_template.replace("{context}", ""));
        let question_tokens = DefaultContextAssembler::estimate_tokens(question);

        // History windowing
        let history_window: Vec<&ChatMessage> = history
            .iter()
            .filter(|m| m.role != ChatRole::System)
            .collect();
        let history_msgs: Vec<&ChatMessage> =
            if history_window.len() > config.max_history_turns * 2 {
                history_window[history_window.len() - config.max_history_turns * 2..].to_vec()
            } else {
                history_window
            };
        let history_tokens: usize = history_msgs.iter().map(|m| m.approx_tokens()).sum();

        let overhead = template_shell_tokens + question_tokens + history_tokens + 30;
        let context_budget = max_tokens.saturating_sub(overhead);

        // 2. Partition results
        let (doc_results, graph_results, community_results) = Self::partition_results(retrieved);

        // 3. Allocate budgets
        let doc_budget = (context_budget as f64 * self.doc_budget_ratio) as usize;
        let graph_budget = (context_budget as f64 * self.graph_budget_ratio) as usize;
        let community_budget = (context_budget as f64 * self.community_budget_ratio) as usize;

        // 4. Build each section
        let (doc_section, doc_tokens) =
            Self::build_section("## Retrieved Documents", &doc_results, doc_budget);
        let (graph_section, graph_tokens) =
            Self::build_section("## Graph Context (Entities & Relations)", &graph_results, graph_budget);
        let (community_section, community_tokens) =
            Self::build_section("## Community Summaries", &community_results, community_budget);

        // 5. Combine sections
        let mut context_parts = Vec::new();
        if !doc_section.is_empty() {
            context_parts.push(doc_section);
        }
        if !graph_section.is_empty() {
            context_parts.push(graph_section);
        }
        if !community_section.is_empty() {
            context_parts.push(community_section);
        }

        let context_str = context_parts.join("\n\n");
        let total_context_tokens = doc_tokens + graph_tokens + community_tokens;

        // 6. Assemble messages
        let mut messages: Vec<ChatMessage> = Vec::new();

        let system_content = config.system_template.replace("{context}", &context_str);
        messages.push(ChatMessage::system(system_content));

        for msg in &history_msgs {
            messages.push((*msg).clone());
        }

        let user_content = config.user_template.replace("{question}", question);
        messages.push(ChatMessage::user(user_content));

        let total_tokens = total_context_tokens + overhead;

        Ok((messages, total_tokens))
    }
}

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

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

    #[test]
    fn test_format_document() {
        let doc = RetrievalResult::new("doc-1", "Hello world content", 0.95)
            .with_collection("notes");
        let formatted = DefaultContextAssembler::format_document(0, &doc);
        assert!(formatted.contains("[1]"));
        assert!(formatted.contains("[notes]"));
        assert!(formatted.contains("0.95"));
        assert!(formatted.contains("Hello world content"));
    }

    #[test]
    fn test_estimate_tokens_english() {
        let text = "Hello world this is a test"; // 26 chars → ~7 tokens
        let tokens = DefaultContextAssembler::estimate_tokens(text);
        assert!(tokens > 0);
        assert!(tokens < 15);
    }

    #[test]
    fn test_estimate_tokens_cjk() {
        let text = "这是一个中文测试句子"; // 9 CJK chars → ~6 tokens
        let tokens = DefaultContextAssembler::estimate_tokens(text);
        assert!(tokens > 0);
        assert!(tokens < 20);
    }

    #[test]
    fn test_truncate_to_tokens_short() {
        let text = "Short text";
        let result = DefaultContextAssembler::truncate_to_tokens(text, 100);
        assert_eq!(result, "Short text"); // No truncation needed
    }

    #[test]
    fn test_truncate_to_tokens_long() {
        let text = "A".repeat(1000); // ~250 tokens
        let result = DefaultContextAssembler::truncate_to_tokens(&text, 50);
        assert!(result.len() < 1000);
        assert!(result.ends_with("...[truncated]"));
    }

    #[test]
    fn test_build_context() {
        let docs = vec![
            RetrievalResult::new("1", "First document about Rust", 0.9),
            RetrievalResult::new("2", "Second document about Python", 0.8),
            RetrievalResult::new("3", "Third document about Go", 0.7),
        ];

        let (context, tokens) = DefaultContextAssembler::build_context(&docs, 1000);
        assert!(context.contains("[1]"));
        assert!(context.contains("[2]"));
        assert!(context.contains("[3]"));
        assert!(context.contains("---")); // separator
        assert!(tokens > 0);
    }

    #[test]
    fn test_build_context_budget_limit() {
        let docs: Vec<RetrievalResult> = (0..100)
            .map(|i| RetrievalResult::new(format!("doc-{}", i), "A".repeat(200), 0.9 - i as f64 * 0.01))
            .collect();

        // Very small budget — should only include a few docs
        let (context, tokens) = DefaultContextAssembler::build_context(&docs, 100);
        assert!(tokens <= 120); // Allow some margin
        // Should not include all 100 docs
        assert!(!context.contains("[100]"));
    }

    #[test]
    fn test_assemble_basic() {
        let assembler = DefaultContextAssembler::new();
        let config = RagConfig::default();

        let docs = vec![
            RetrievalResult::new("1", "Rust is a systems language", 0.9),
            RetrievalResult::new("2", "Rust has zero-cost abstractions", 0.8),
        ];

        let (messages, tokens) = assembler
            .assemble("What is Rust?", &docs, &[], &config)
            .unwrap();

        // Should have: system + user = 2 messages (no history)
        assert_eq!(messages.len(), 2);
        assert_eq!(messages[0].role, ChatRole::System);
        assert_eq!(messages[1].role, ChatRole::User);

        // System should contain the context
        assert!(messages[0].content.contains("Rust is a systems language"));
        assert!(messages[0].content.contains("zero-cost abstractions"));

        // User should contain the question
        assert!(messages[1].content.contains("What is Rust?"));

        assert!(tokens > 0);
    }

    #[test]
    fn test_assemble_with_history() {
        let assembler = DefaultContextAssembler::new();
        let config = RagConfig {
            max_history_turns: 2,
            ..Default::default()
        };

        let docs = vec![RetrievalResult::new("1", "Some content", 0.9)];

        let history = vec![
            ChatMessage::user("First question"),
            ChatMessage::assistant("First answer"),
            ChatMessage::user("Second question"),
            ChatMessage::assistant("Second answer"),
            ChatMessage::user("Third question"),
            ChatMessage::assistant("Third answer"),
        ];

        let (messages, _tokens) = assembler
            .assemble("Current question", &docs, &history, &config)
            .unwrap();

        // system + 4 history (last 2 turns) + user = 6
        assert_eq!(messages.len(), 6);
        assert_eq!(messages[0].role, ChatRole::System);
        // History should be the last 2 turns (4 messages)
        assert_eq!(messages[1].content, "Second question");
        assert_eq!(messages[2].content, "Second answer");
        assert_eq!(messages[3].content, "Third question");
        assert_eq!(messages[4].content, "Third answer");
        assert_eq!(messages[5].role, ChatRole::User);
    }

    #[test]
    fn test_assemble_empty_retrieval() {
        let assembler = DefaultContextAssembler::new();
        let config = RagConfig::default();

        let (messages, _) = assembler
            .assemble("Hello?", &[], &[], &config)
            .unwrap();

        // Should still produce system + user
        assert_eq!(messages.len(), 2);
        // Context area should be empty in system message
        assert!(messages[0].content.contains("## Context"));
    }

    #[test]
    fn test_is_cjk() {
        assert!(is_cjk('中'));
        assert!(is_cjk('日'));
        assert!(is_cjk('한'));
        assert!(!is_cjk('A'));
        assert!(!is_cjk('1'));
    }

    // ── GraphRagAssembler tests ──────────────────────────────────────

    #[test]
    fn test_graph_rag_assembler_defaults() {
        let assembler = GraphRagAssembler::new();
        assert!((assembler.doc_budget_ratio - 0.60).abs() < 0.01);
        assert!((assembler.graph_budget_ratio - 0.25).abs() < 0.01);
        assert!((assembler.community_budget_ratio - 0.15).abs() < 0.01);
    }

    #[test]
    fn test_graph_rag_assembler_custom_ratios() {
        let assembler = GraphRagAssembler::new().with_ratios(1.0, 1.0, 1.0);
        // All equal → each ~0.33
        assert!((assembler.doc_budget_ratio - 1.0 / 3.0).abs() < 0.01);
        assert!((assembler.graph_budget_ratio - 1.0 / 3.0).abs() < 0.01);
    }

    #[test]
    fn test_graph_rag_partition_results() {
        let mut doc = RetrievalResult::new("d1", "A document", 0.9);
        // No special metadata → classified as doc

        let mut graph = RetrievalResult::new("g1", "Entity info", 0.8);
        graph.metadata.insert("source".into(), serde_json::json!("graph_retriever"));
        graph.collection = Some("graph".into());

        let mut community = RetrievalResult::new("c1", "Community summary", 0.7);
        community.metadata.insert("source".into(), serde_json::json!("community_summary"));
        community.collection = Some("graph_community".into());

        let results = vec![doc.clone(), graph.clone(), community.clone()];
        let (docs, graphs, communities) = GraphRagAssembler::partition_results(&results);
        assert_eq!(docs.len(), 1);
        assert_eq!(graphs.len(), 1);
        assert_eq!(communities.len(), 1);
        assert_eq!(docs[0].id, "d1");
        assert_eq!(graphs[0].id, "g1");
        assert_eq!(communities[0].id, "c1");
    }

    #[test]
    fn test_graph_rag_assemble_mixed() {
        let assembler = GraphRagAssembler::new();
        let config = RagConfig::default();

        let mut graph_result = RetrievalResult::new(
            "entity-1",
            "[Entity] Alice (type: person)\nRelations:\n  Alice --[knows]--> Bob (weight: 0.9)",
            0.85,
        );
        graph_result.metadata.insert("source".into(), serde_json::json!("graph_retriever"));
        graph_result.collection = Some("graph".into());

        let doc_result = RetrievalResult::new("doc-1", "Alice is a researcher at MIT.", 0.9);

        let results = vec![doc_result, graph_result];

        let (messages, tokens) = assembler
            .assemble("Who is Alice?", &results, &[], &config)
            .unwrap();

        assert_eq!(messages.len(), 2); // system + user
        assert!(messages[0].content.contains("Retrieved Documents"));
        assert!(messages[0].content.contains("Alice is a researcher"));
        assert!(messages[0].content.contains("Graph Context"));
        assert!(messages[0].content.contains("Alice"));
        assert!(messages[0].content.contains("knows"));
        assert!(tokens > 0);
    }

    #[test]
    fn test_graph_rag_assemble_only_docs() {
        let assembler = GraphRagAssembler::new();
        let config = RagConfig::default();

        let results = vec![
            RetrievalResult::new("d1", "Document content A", 0.9),
            RetrievalResult::new("d2", "Document content B", 0.8),
        ];

        let (messages, _) = assembler
            .assemble("What?", &results, &[], &config)
            .unwrap();

        assert_eq!(messages.len(), 2);
        assert!(messages[0].content.contains("Retrieved Documents"));
        // Should not contain graph section headers (no graph results)
        assert!(!messages[0].content.contains("Graph Context"));
        assert!(!messages[0].content.contains("Community Summaries"));
    }

    #[test]
    fn test_graph_rag_assemble_empty() {
        let assembler = GraphRagAssembler::new();
        let config = RagConfig::default();

        let (messages, _) = assembler
            .assemble("Hello?", &[], &[], &config)
            .unwrap();

        assert_eq!(messages.len(), 2);
    }
}
