use crate::bus::{
    ChatMessage, SYSTEM_CONTEXT_AGENT_PROMPT, SYSTEM_CONTEXT_HISTORY_COMPACTION,
    SYSTEM_CONTEXT_SCHEDULED_PROMPT,
};
use crate::config::LLMProviderConfig;
use crate::providers::{ChatCompletionRequest, LLMProvider, Message, create_provider};
use crate::text::{char_count, take_prefix_chars};
use crate::agent::{AgentError, AgentRuntimeConfig};

const TOKEN_ESTIMATE_SAFETY_MULTIPLIER: f64 = 1.2;
const CJK_CHARS_PER_TOKEN: f64 = 2.0;
const OTHER_CHARS_PER_TOKEN: f64 = 4.0;
const JSON_OVERHEAD_PER_MESSAGE: usize = 50;

/// System context marker for the heavy compression (older segment) summary.
pub const SYSTEM_CONTEXT_HISTORY_COMPACTION_OLDER: &str = "history_compaction_older";
/// System context marker for the light compression (newer segment) summary.
pub const SYSTEM_CONTEXT_HISTORY_COMPACTION_NEWER: &str = "history_compaction_newer";

/// Default threshold ratio: compress when estimated tokens exceed 70% of context window.
const DEFAULT_THRESHOLD_RATIO: f64 = 0.7;

/// Default budget split: older segment gets 30% of summary budget, newer gets 70%.
const OLDER_BUDGET_RATIO: f64 = 0.3;
const NEWER_BUDGET_RATIO: f64 = 0.7;

// ============================================================================
// HistoryUnit — atomic message units for compression
// ============================================================================

/// An indivisible unit of conversation history. Compression operates at the
/// unit level, so tool call sequences can never be split across boundaries.
#[derive(Debug, Clone)]
enum HistoryUnit {
    /// Always-preserved system messages (agent_prompt, scheduled_prompt).
    SystemGuard(ChatMessage),
    /// A user message.
    UserMessage(ChatMessage),
    /// An atomic tool call sequence: assistant(tool_calls) + all following
    /// tool-result messages. Must be preserved or compressed as a whole.
    ToolRound {
        assistant: ChatMessage,
        results: Vec<ChatMessage>,
    },
    /// A plain assistant text response (no tool_calls).
    AssistantText(ChatMessage),
}

impl HistoryUnit {
    /// Estimate tokens for this unit alone.
    fn estimate_tokens(&self) -> usize {
        match self {
            HistoryUnit::SystemGuard(msg) | HistoryUnit::UserMessage(msg) | HistoryUnit::AssistantText(msg) => {
                estimate_tokens(std::slice::from_ref(msg))
            }
            HistoryUnit::ToolRound { assistant, results } => {
                let mut all = vec![assistant.clone()];
                all.extend(results.clone());
                estimate_tokens(&all)
            }
        }
    }

    /// Expand this unit back into a flat list of ChatMessages.
    #[allow(dead_code)]
    fn into_messages(self) -> Vec<ChatMessage> {
        match self {
            HistoryUnit::SystemGuard(msg)
            | HistoryUnit::UserMessage(msg)
            | HistoryUnit::AssistantText(msg) => vec![msg],
            HistoryUnit::ToolRound { assistant, results } => {
                let mut msgs = vec![assistant];
                msgs.extend(results);
                msgs
            }
        }
    }

    /// Returns true if this unit is a ToolRound.
    #[allow(dead_code)]
    fn is_tool_round(&self) -> bool {
        matches!(self, HistoryUnit::ToolRound { .. })
    }
}

// ============================================================================
// Unit parser — one forward pass, O(n)
// ============================================================================

fn is_system_guard(msg: &ChatMessage) -> bool {
    msg.role == "system"
        && (msg.has_system_context(SYSTEM_CONTEXT_AGENT_PROMPT)
            || msg.has_system_context(SYSTEM_CONTEXT_SCHEDULED_PROMPT))
}

fn is_assistant_with_tool_calls(msg: &ChatMessage) -> bool {
    msg.role == "assistant"
        && msg
            .tool_calls
            .as_ref()
            .map_or(false, |calls| !calls.is_empty())
}

/// Parse a flat message list into atomic units. Orphaned tool results
/// (without a preceding assistant with tool_calls) are silently dropped.
fn parse_to_units(messages: &[ChatMessage]) -> Vec<HistoryUnit> {
    let mut units = Vec::new();
    let mut i = 0;

    while i < messages.len() {
        let msg = &messages[i];

        if is_system_guard(msg) {
            units.push(HistoryUnit::SystemGuard(msg.clone()));
            i += 1;
        } else if msg.role == "user" {
            units.push(HistoryUnit::UserMessage(msg.clone()));
            i += 1;
        } else if is_assistant_with_tool_calls(msg) {
            let assistant = msg.clone();
            let mut results = Vec::new();
            i += 1;
            while i < messages.len() && messages[i].role == "tool" {
                results.push(messages[i].clone());
                i += 1;
            }
            units.push(HistoryUnit::ToolRound { assistant, results });
        } else if msg.role == "assistant" {
            units.push(HistoryUnit::AssistantText(msg.clone()));
            i += 1;
        } else if msg.role == "tool" {
            // Orphaned tool result — drop with a warning
            tracing::warn!(
                tool_call_id = ?msg.tool_call_id,
                message_index = i,
                "Dropping orphaned tool result during unit parsing"
            );
            i += 1;
        } else {
            // system messages that aren't guards — keep as plain messages
            units.push(HistoryUnit::AssistantText(msg.clone()));
            i += 1;
        }
    }

    units
}

/// Estimate total tokens from a slice of units.
fn estimate_tokens_from_units(units: &[HistoryUnit]) -> usize {
    units.iter().map(|u| u.estimate_tokens()).sum()
}

/// Check if a character is CJK (Chinese, Japanese, Korean)
fn is_cjk_char(c: char) -> bool {
    matches!(c,
        '\u{4E00}'..='\u{9FFF}' |  // CJK Unified Ideographs
        '\u{3040}'..='\u{309F}' |  // Hiragana
        '\u{30A0}'..='\u{30FF}' |  // Katakana
        '\u{AC00}'..='\u{D7AF}' |  // Korean Hangul
        '\u{3400}'..='\u{4DBF}' |  // CJK Extension A
        '\u{20000}'..='\u{2A6DF}'  // CJK Extension B
    )
}

/// Token estimation using weighted character counting based on language
pub fn estimate_tokens(messages: &[ChatMessage]) -> usize {
    let mut cjk_count = 0usize;
    let mut other_count = 0usize;
    let mut media_refs_count = 0usize;

    for msg in messages {
        // Count content characters
        for c in msg.content.chars() {
            if is_cjk_char(c) {
                cjk_count += 1;
            } else {
                other_count += 1;
            }
        }
        // Count media references
        media_refs_count += msg.media_refs.len();
    }

    // Weighted token calculation: CJK chars need more tokens per character
    let content_tokens = (cjk_count as f64 / CJK_CHARS_PER_TOKEN)
        + (other_count as f64 / OTHER_CHARS_PER_TOKEN);

    // JSON serialization overhead for message structure (fields, brackets, etc.)
    let json_overhead = messages.len() * JSON_OVERHEAD_PER_MESSAGE;

    // Media references add to JSON size (each path is a string in the array)
    let media_overhead = media_refs_count * 20; // Each media ref adds ~20 chars to JSON

    // Apply safety multiplier
    ((content_tokens + json_overhead as f64 + media_overhead as f64)
        * TOKEN_ESTIMATE_SAFETY_MULTIPLIER) as usize
}

/// Configuration for context compression.
#[derive(Debug, Clone)]
pub struct ContextCompressionConfig {
    /// Preserve the latest N complete user turns in full.
    pub retain_last_user_turns: usize,
    /// Maximum characters in summary
    pub summary_max_chars: usize,
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct UserTurnRange {
    start: usize,
    end_exclusive: usize,
}

#[derive(Debug, Clone)]
pub struct HistoryCompactionPlan {
    pub preserved_system_messages: Vec<ChatMessage>,
    pub summary_message: ChatMessage,
    pub preserved_messages: Vec<ChatMessage>,
    pub compressed_turns: usize,
    pub preserved_turns: usize,
}

impl Default for ContextCompressionConfig {
    fn default() -> Self {
        Self {
            retain_last_user_turns: 3,
            summary_max_chars: 20_000,
        }
    }
}

/// Context compressor that reduces message history when it exceeds token limits.
#[derive(Clone)]
pub struct ContextCompressor {
    config: ContextCompressionConfig,
    context_window: usize,
    /// Threshold ratio to trigger compression (70% of context window).
    threshold_ratio: f64,
}

impl ContextCompressor {
    #[cfg(test)]
    fn summary_char_budget_for_context_window(context_window: usize) -> usize {
        const SUMMARY_RATIO: f64 = 0.1;
        const CHARS_PER_TOKEN: f64 = 2.5;
        const MIN_SUMMARY_CHARS: usize = 1_500;
        const MAX_SUMMARY_CHARS: usize = 50_000;

        ((context_window as f64 * SUMMARY_RATIO * CHARS_PER_TOKEN) as usize)
            .clamp(MIN_SUMMARY_CHARS, MAX_SUMMARY_CHARS)
    }

    // =========================================================================
    // Transcript building (shared helpers, unchanged from original)
    // =========================================================================

    fn format_transcript_entry(message: &ChatMessage) -> String {
        let role = match message.role.as_str() {
            "assistant" => "Assistant",
            "tool" => "Tool",
            _ => message.role.as_str(),
        };
        let name = message
            .tool_name
            .as_ref()
            .map(|n| format!(" ({})", n))
            .unwrap_or_default();
        format!("{}: {}{}", role, message.content, name)
    }

    fn build_transcript(messages: &[ChatMessage]) -> String {
        messages
            .iter()
            .map(Self::format_transcript_entry)
            .collect::<Vec<_>>()
            .join("\n\n")
    }

    fn split_text_chunks(text: &str, max_chars: usize) -> Vec<String> {
        if text.is_empty() {
            return Vec::new();
        }

        let chunk_size = max_chars.max(1);
        let chars: Vec<char> = text.chars().collect();
        chars
            .chunks(chunk_size)
            .map(|chunk| chunk.iter().collect())
            .collect()
    }

    fn chunk_messages_for_summary(messages: &[ChatMessage], max_chars: usize) -> Vec<String> {
        if messages.is_empty() {
            return Vec::new();
        }

        let chunk_limit = max_chars.max(1);
        let mut chunks = Vec::new();
        let mut current = String::new();

        for entry in messages.iter().map(Self::format_transcript_entry) {
            let separator = if current.is_empty() { "" } else { "\n\n" };
            let candidate = format!("{}{}{}", current, separator, entry);

            if !current.is_empty() && char_count(&candidate) > chunk_limit {
                chunks.push(current);
                current = String::new();
            }

            if char_count(&entry) > chunk_limit {
                if !current.is_empty() {
                    chunks.push(current);
                    current = String::new();
                }
                chunks.extend(Self::split_text_chunks(&entry, chunk_limit));
                continue;
            }

            if current.is_empty() {
                current = entry;
            } else {
                current.push_str("\n\n");
                current.push_str(&entry);
            }
        }

        if !current.is_empty() {
            chunks.push(current);
        }

        chunks
    }

    fn build_summary_prompt(transcript: &str, target_chars: usize) -> String {
        format!(
            r#"You are a conversation compaction engine. Summarize the following conversation segment.

PRESERVE:
- Each user question or request in full or as a near-verbatim restatement
- All identifiers (UUIDs, hashes, file paths, URLs)
- Actions taken (tool calls, file operations, commands)
- Key information obtained (results, data, errors)
- Decisions and user preferences
- Current task status

OMIT:
- Reproducing full tool output verbatim unless it is essential
- Repeated greetings or filler

Do not assume tool content was pre-trimmed. You may receive long tool outputs; keep the important results, errors, and artifacts.

Be concise, aim for {} characters or less.

---

{}

"#,
            target_chars, transcript
        )
    }

    // =========================================================================
    // Two-segment compression: heavy (older) + light (newer) prompts
    // =========================================================================

    /// Prompt for the older segment — aggressive compression, keep only essentials.
    fn build_heavy_summary_prompt(transcript: &str, target_chars: usize) -> String {
        format!(
            r#"You are a conversation compaction engine. Aggressively summarize the following OLDER conversation segment. These events happened earlier in the session and will be marked as "较早的操作记录（已压缩）".

A newer portion (not shown here) follows this segment. The model can use tools to re-read files for exact data, so you do NOT need to preserve full tool outputs.

=== MUST PRESERVE (keep verbatim or near-verbatim) ===
- All file paths (absolute and relative)
- All URLs
- Key numeric results and conclusions (one sentence per result)
- Error messages (exact text if important)
- User preferences and decisions
- The original user task or goal

=== SHOULD CONDENSE ===
- Each tool call → one line summary: "tool_name: 关键参数 → 结果概要"
- Multiple similar operations → group: "对 ./data/ 目录下的 10 个 CSV 文件执行了统计分析"
- Verbose tool output → drop entirely unless it contains a critical result or error

=== SHOULD DROP ===
- Full tool output (hundreds of lines of raw data)
- Repeated attempts of the same operation
- Greetings, filler text, markdown formatting

Be concise, aim for {} characters or less. Output the summary in Chinese if the original conversation was in Chinese.

---

OLDER SEGMENT (events from earlier in the session):

{}

"#,
            target_chars, transcript
        )
    }

    /// Prompt for the newer segment — lighter compression, keep more detail.
    fn build_light_summary_prompt(transcript: &str, target_chars: usize) -> String {
        format!(
            r#"You are a conversation compaction engine. Lightly summarize the following RECENT conversation segment. These events happened just before the current moment and will be marked as "近期操作记录（轻度压缩，保留了更多细节）".

An older segment (already summarized separately) precedes this. Keep enough detail so the model can continue the task seamlessly without re-reading files for data that appears in these recent results.

=== MUST PRESERVE (keep with more detail than a normal summary) ===
- All file paths, URLs, and identifiers
- The exact sequence of recent operations (step by step)
- Tool parameters (especially file paths, search queries, command strings)
- Key outputs and results (shortened but keep the substance)
- Error messages (complete, not truncated)
- Current task status and what should happen next

=== SHOULD CONDENSE ===
- Tool outputs → truncate to the most meaningful parts (key data, conclusions, not raw output)
- Long text → keep the essence but not every word
- Repeated similar outputs → note the pattern but keep a representative example

=== SHOULD DROP ===
- Completely irrelevant debug output
- Boilerplate text with no informational value
- Trivial operations that have no bearing on the task

Be concise, aim for {} characters or less. Output the summary in Chinese if the original conversation was in Chinese.

---

RECENT SEGMENT (events from just before the current moment):

{}

"#,
            target_chars, transcript
        )
    }

    async fn summarize_transcript(
        &self,
        provider: &dyn LLMProvider,
        transcript: &str,
        target_chars: usize,
    ) -> Result<String, AgentError> {
        let request = ChatCompletionRequest {
            messages: vec![
                Message::system("You are a helpful assistant."),
                Message::user(Self::build_summary_prompt(transcript, target_chars)),
            ],
            temperature: Some(0.3),
            max_tokens: Some(1000),
            tools: None,
        };

        let response = provider
            .chat(request)
            .await
            .map_err(|e| AgentError::LlmError(e.to_string()))?;
        Ok(response.content)
    }

    async fn summarize_chunked_transcript(
        &self,
        provider: &dyn LLMProvider,
        messages: &[ChatMessage],
        transcript: &str,
    ) -> Result<String, AgentError> {
        let target_chars = self.config.summary_max_chars.max(1);
        let mut layer = Self::chunk_messages_for_summary(messages, target_chars);

        if layer.is_empty() {
            layer.push(transcript.to_string());
        }

        for _ in 0..6 {
            if layer.len() == 1 && char_count(&layer[0]) <= target_chars {
                return self
                    .summarize_transcript(provider, &layer[0], target_chars)
                    .await;
            }

            let per_chunk_target = (target_chars / layer.len().max(1))
                .max(500)
                .min(target_chars);
            let mut summaries = Vec::with_capacity(layer.len());
            for chunk in &layer {
                summaries.push(
                    self.summarize_transcript(provider, chunk, per_chunk_target)
                        .await?,
                );
            }

            if summaries.len() == 1 {
                let summary = summaries.pop().unwrap_or_default();
                if char_count(&summary) <= target_chars {
                    return Ok(summary);
                }
                layer = Self::split_text_chunks(&summary, target_chars);
                continue;
            }

            let merged = summaries.join("\n\n");
            if char_count(&merged) <= target_chars {
                return self
                    .summarize_transcript(provider, &merged, target_chars)
                    .await;
            }

            layer = Self::split_text_chunks(&merged, target_chars);
        }

        Ok(take_prefix_chars(transcript, target_chars))
    }

    // =========================================================================
    // Two-segment compression
    // =========================================================================

    /// Find a safe split point in the unit list. The split ensures:
    /// - Accumulated tokens up to `ratio` of total are placed in the older segment.
    /// - The split never lands on a SystemGuard (they're always preserved).
    /// - The split lands on a unit boundary (ToolRound, AssistantText, or UserMessage).
    fn find_safe_split_point(&self, units: &[HistoryUnit], ratio: f64) -> usize {
        let total_tokens = estimate_tokens_from_units(units);
        let target = (total_tokens as f64 * ratio) as usize;
        let mut accumulated = 0;

        for (i, unit) in units.iter().enumerate() {
            // Never split on a SystemGuard — they must be preserved
            if matches!(unit, HistoryUnit::SystemGuard(_)) {
                continue;
            }
            accumulated += unit.estimate_tokens();
            if accumulated >= target {
                return i + 1;
            }
        }

        // Fallback: put everything in the older segment
        units.len()
    }

    /// Summarize a transcript with a custom prompt builder function.
    async fn summarize_with_prompt(
        &self,
        provider: &dyn LLMProvider,
        transcript: &str,
        target_chars: usize,
        build_prompt: fn(&str, usize) -> String,
    ) -> Result<String, AgentError> {
        let request = ChatCompletionRequest {
            messages: vec![
                Message::system("You are a helpful assistant."),
                Message::user(build_prompt(transcript, target_chars)),
            ],
            temperature: Some(0.3),
            max_tokens: Some(1000),
            tools: None,
        };

        let response = provider
            .chat(request)
            .await
            .map_err(|e| AgentError::LlmError(e.to_string()))?;
        Ok(response.content)
    }

    /// Summarize a segment of units with chunked fallback.
    async fn summarize_units_segment(
        &self,
        provider: &dyn LLMProvider,
        messages: &[ChatMessage],
        transcript: &str,
        target_chars: usize,
        build_prompt: fn(&str, usize) -> String,
    ) -> Result<String, AgentError> {
        if messages.is_empty() {
            return Ok(String::new());
        }

        if char_count(transcript) <= target_chars {
            return self
                .summarize_with_prompt(provider, transcript, target_chars, build_prompt)
                .await;
        }

        // Use chunked summarization for oversized transcripts
        let target = target_chars.max(1);
        let mut layer = Self::chunk_messages_for_summary(messages, target);

        if layer.is_empty() {
            layer.push(transcript.to_string());
        }

        for _ in 0..6 {
            if layer.len() == 1 && char_count(&layer[0]) <= target {
                return self
                    .summarize_with_prompt(provider, &layer[0], target, build_prompt)
                    .await;
            }

            let per_chunk_target = (target / layer.len().max(1)).max(500).min(target);
            let mut summaries = Vec::with_capacity(layer.len());
            for chunk in &layer {
                summaries.push(
                    self.summarize_with_prompt(provider, chunk, per_chunk_target, build_prompt)
                        .await?,
                );
            }

            if summaries.len() == 1 {
                let summary = summaries.pop().unwrap_or_default();
                if char_count(&summary) <= target {
                    return Ok(summary);
                }
                layer = Self::split_text_chunks(&summary, target);
                continue;
            }

            let merged = summaries.join("\n\n");
            if char_count(&merged) <= target {
                return self
                    .summarize_with_prompt(provider, &merged, target, build_prompt)
                    .await;
            }

            layer = Self::split_text_chunks(&merged, target);
        }

        Ok(take_prefix_chars(transcript, target))
    }

    /// Main entry point for two-segment compression.
    ///
    /// Splits history into older and newer segments, then summarizes each
    /// with LLM using different prompts and budget allocations. The result
    /// contains no tool_calls or tool-result messages — only system summaries
    /// and the original user message — eliminating any risk of API 400 errors
    /// from orphaned tool call sequences.
    pub async fn compress_two_segment(
        &self,
        history: &[ChatMessage],
        provider_config: &LLMProviderConfig,
    ) -> Result<Vec<ChatMessage>, AgentError> {
        let tokens = estimate_tokens(history);
        if tokens <= self.threshold() {
            tracing::info!(
                tokens = tokens,
                threshold = self.threshold(),
                msg_count = history.len(),
                "Two-segment compression not needed (under threshold)"
            );
            return Ok(history.to_vec());
        }

        tracing::info!(
            tokens = tokens,
            threshold = self.threshold(),
            msg_count = history.len(),
            "Starting two-segment compression"
        );

        // Step 1: Parse into atomic units
        let units = parse_to_units(history);

        // Step 2: Separate system guards + user messages from compressible units
        let mut system_guards: Vec<ChatMessage> = Vec::new();
        let mut user_messages: Vec<ChatMessage> = Vec::new();
        let mut compressible: Vec<HistoryUnit> = Vec::new();

        for unit in units {
            match unit {
                HistoryUnit::SystemGuard(msg) => system_guards.push(msg),
                HistoryUnit::UserMessage(msg) => user_messages.push(msg),
                other => compressible.push(other),
            }
        }

        // Keep only the latest user message in full; older ones go into compression
        let latest_user_msg = user_messages.pop();

        // Step 3: Find split point in compressible units
        let total_compressible = estimate_tokens_from_units(&compressible);
        let split_point = if total_compressible > 0 {
            self.find_safe_split_point(&compressible, 0.5)
        } else {
            0
        };
        let older_units: Vec<&HistoryUnit> = compressible[..split_point].iter().collect();
        let newer_units: Vec<&HistoryUnit> = compressible[split_point..].iter().collect();

        // Step 4: Build transcripts
        let older_messages: Vec<ChatMessage> = older_units
            .iter()
            .flat_map(|u| match u {
                HistoryUnit::ToolRound { assistant, results } => {
                    let mut msgs = vec![assistant.clone()];
                    msgs.extend(results.clone());
                    msgs
                }
                HistoryUnit::AssistantText(msg) => vec![msg.clone()],
                HistoryUnit::UserMessage(msg) => {
                    // Older user messages go into the compressible transcript
                    vec![msg.clone()]
                }
                _ => vec![],
            })
            .collect();

        let newer_messages: Vec<ChatMessage> = newer_units
            .iter()
            .flat_map(|u| match u {
                HistoryUnit::ToolRound { assistant, results } => {
                    let mut msgs = vec![assistant.clone()];
                    msgs.extend(results.clone());
                    msgs
                }
                HistoryUnit::AssistantText(msg) => vec![msg.clone()],
                _ => vec![],
            })
            .collect();

        // Include older user messages in the older transcript
        let older_user_msgs: Vec<ChatMessage> = user_messages
            .iter()
            .map(|m| {
                let mut msg = m.clone();
                msg.role = "user".to_string();
                msg
            })
            .collect();
        let all_older_messages: Vec<ChatMessage> = older_user_msgs
            .iter()
            .chain(older_messages.iter())
            .cloned()
            .collect();

        let older_transcript = Self::build_transcript(&all_older_messages);
        let newer_transcript = Self::build_transcript(&newer_messages);

        // Step 5: Budget allocation
        let total_budget = self.config.summary_max_chars;
        let older_budget = (total_budget as f64 * OLDER_BUDGET_RATIO) as usize;
        let newer_budget = (total_budget as f64 * NEWER_BUDGET_RATIO) as usize;

        // Step 6: Create provider and run both summaries (can be parallel)
        let runtime_config = AgentRuntimeConfig::from(provider_config.clone());
        let provider = create_provider(runtime_config.provider)
            .map_err(|e| AgentError::ProviderCreation(e.to_string()))?;

        let (older_result, newer_result) = if older_units.is_empty() && newer_units.is_empty() {
            (Ok(String::new()), Ok(String::new()))
        } else if older_units.is_empty() {
            let result = self
                .summarize_units_segment(
                    provider.as_ref(),
                    &newer_messages,
                    &newer_transcript,
                    newer_budget,
                    Self::build_light_summary_prompt,
                )
                .await;
            (Ok(String::new()), result)
        } else if newer_units.is_empty() {
            let result = self
                .summarize_units_segment(
                    provider.as_ref(),
                    &all_older_messages,
                    &older_transcript,
                    older_budget,
                    Self::build_heavy_summary_prompt,
                )
                .await;
            (result, Ok(String::new()))
        } else {
            let older_fut = self.summarize_units_segment(
                provider.as_ref(),
                &all_older_messages,
                &older_transcript,
                older_budget,
                Self::build_heavy_summary_prompt,
            );
            let newer_fut = self.summarize_units_segment(
                provider.as_ref(),
                &newer_messages,
                &newer_transcript,
                newer_budget,
                Self::build_light_summary_prompt,
            );
            tokio::join!(older_fut, newer_fut)
        };

        let older_summary = older_result?;
        let newer_summary = newer_result?;

        // Step 7: Assemble compressed history
        let mut compressed: Vec<ChatMessage> = Vec::with_capacity(4);

        // System guards first
        compressed.extend(system_guards);

        // Latest user message
        if let Some(user_msg) = latest_user_msg {
            compressed.push(user_msg);
        }

        // Heavy compression summary (older)
        if !older_summary.is_empty() {
            compressed.push(ChatMessage::system_with_context(
                format!("## 较早的操作记录（已压缩）\n\n{}", older_summary),
                Some(SYSTEM_CONTEXT_HISTORY_COMPACTION_OLDER.to_string()),
            ));
        }

        // Light compression summary (newer)
        if !newer_summary.is_empty() {
            compressed.push(ChatMessage::system_with_context(
                format!(
                    "## 近期操作记录（轻度压缩，保留了更多细节）\n\n{}\n\n---\n以上为近期操作记录。如需精确数据，可使用工具重新读取相关文件。",
                    newer_summary
                ),
                Some(SYSTEM_CONTEXT_HISTORY_COMPACTION_NEWER.to_string()),
            ));
        }

        tracing::info!(
            original_tokens = tokens,
            final_tokens = estimate_tokens(&compressed),
            final_msg_count = compressed.len(),
            older_units = older_units.len(),
            newer_units = newer_units.len(),
            "Two-segment compression completed"
        );

        Ok(compressed)
    }

    // =========================================================================
    // Legacy methods (kept for backward compatibility during migration)
    // =========================================================================

    /// Create a new compressor with the given context window size.
    pub fn new(context_window: usize) -> Self {
        Self {
            config: ContextCompressionConfig::default(),
            context_window,
            threshold_ratio: DEFAULT_THRESHOLD_RATIO,
        }
    }

    pub fn from_provider_config(provider_config: &LLMProviderConfig) -> Self {
        Self::from_runtime_config(&AgentRuntimeConfig::from(provider_config.clone()))
    }

    pub fn from_runtime_config(config: &AgentRuntimeConfig) -> Self {
        Self::with_config(
            config.context_window_tokens,
            ContextCompressionConfig {
                summary_max_chars: config.context_summary_char_budget,
                ..ContextCompressionConfig::default()
            },
        )
    }

    /// Create with custom configuration.
    pub fn with_config(context_window: usize, config: ContextCompressionConfig) -> Self {
        Self {
            config,
            context_window,
            threshold_ratio: DEFAULT_THRESHOLD_RATIO,
        }
    }

    /// Get the compression threshold in tokens (70% of context window).
    fn threshold(&self) -> usize {
        (self.context_window as f64 * self.threshold_ratio) as usize
    }

    pub fn should_compress(&self, history: &[ChatMessage]) -> bool {
        estimate_tokens(history) > self.threshold()
    }

    fn user_turn_ranges(&self, history: &[ChatMessage]) -> Vec<UserTurnRange> {
        let user_indices: Vec<usize> = history
            .iter()
            .enumerate()
            .filter(|(_, message)| message.role == "user")
            .map(|(index, _)| index)
            .collect();

        user_indices
            .iter()
            .enumerate()
            .map(|(turn_index, start)| UserTurnRange {
                start: *start,
                end_exclusive: user_indices
                    .get(turn_index + 1)
                    .copied()
                    .unwrap_or(history.len()),
            })
            .collect()
    }

    fn should_preserve_system_message(&self, message: &ChatMessage) -> bool {
        message.role == "system"
            && (message.has_system_context(SYSTEM_CONTEXT_AGENT_PROMPT)
                || message.has_system_context(SYSTEM_CONTEXT_SCHEDULED_PROMPT))
    }

    fn split_prefix_messages(
        &self,
        history: &[ChatMessage],
    ) -> (Vec<ChatMessage>, Vec<ChatMessage>) {
        let preserved_system_messages = history
            .iter()
            .filter(|message| self.should_preserve_system_message(message))
            .cloned()
            .collect();
        let summary_source = history
            .iter()
            .filter(|message| !self.should_preserve_system_message(message))
            .cloned()
            .collect();

        (preserved_system_messages, summary_source)
    }

    pub async fn build_compaction_plan(
        &self,
        history: &[ChatMessage],
        provider_config: &LLMProviderConfig,
    ) -> Result<Option<HistoryCompactionPlan>, AgentError> {
        if !self.should_compress(history) {
            return Ok(None);
        }

        let turn_ranges = self.user_turn_ranges(history);
        if turn_ranges.len() <= self.config.retain_last_user_turns {
            return Ok(None);
        }

        let preserved_turn_start =
            turn_ranges[turn_ranges.len() - self.config.retain_last_user_turns].start;
        if preserved_turn_start == 0 {
            return Ok(None);
        }

        let (preserved_system_messages, summary_source) =
            self.split_prefix_messages(&history[..preserved_turn_start]);

        let summary = self
            .summarize_segment(&summary_source, provider_config)
            .await?;

        // Sanitize preserved messages: the boundary between the summarized
        // and preserved sections can split an assistant tool_calls message
        // from its tool result messages, creating orphaned sequences that
        // would cause API 400 errors.
        let mut preserved_messages = history[preserved_turn_start..].to_vec();
        let removed =
            crate::bus::message::sanitize_incomplete_tool_call_sequences(&mut preserved_messages);
        if removed > 0 {
            tracing::warn!(
                removed_count = removed,
                preserved_turn_start,
                "Compaction plan: removed incomplete tool call sequences from preserved messages"
            );
        }

        Ok(Some(HistoryCompactionPlan {
            preserved_system_messages,
            summary_message: ChatMessage::system_with_context(
                format!("[Compressed History]\n\n{}", summary),
                Some(SYSTEM_CONTEXT_HISTORY_COMPACTION.to_string()),
            ),
            preserved_messages,
            compressed_turns: turn_ranges.len() - self.config.retain_last_user_turns,
            preserved_turns: self.config.retain_last_user_turns,
        }))
    }

    /// Main entry point - compresses history if over threshold.
    pub async fn compress_if_needed(
        &self,
        history: Vec<ChatMessage>,
        provider_config: &LLMProviderConfig,
    ) -> Result<Vec<ChatMessage>, AgentError> {
        let tokens = estimate_tokens(&history);
        if tokens <= self.threshold() {
            #[cfg(debug_assertions)]
            tracing::info!(
                tokens = tokens,
                threshold = self.threshold(),
                msg_count = history.len(),
                "Context compression not needed"
            );
            return Ok(history);
        }

        tracing::info!(
            tokens = tokens,
            threshold = self.threshold(),
            msg_count = history.len(),
            "Starting context compression"
        );

        let mut current_history = match self
            .build_compaction_plan(&history, provider_config)
            .await?
        {
            Some(plan) => {
                let mut compressed = Vec::with_capacity(
                    plan.preserved_system_messages.len() + plan.preserved_messages.len() + 1,
                );
                compressed.extend(plan.preserved_system_messages);
                compressed.push(plan.summary_message);
                compressed.extend(plan.preserved_messages);
                compressed
            }
            None => history,
        };

        // Post-compression sanitization: compression can split an assistant
        // tool_calls message from its tool result messages at the boundary
        // between the summarized and preserved sections. This pass removes
        // any orphaned tool_calls or tool results to ensure the message
        // sequence is always valid for the LLM API.
        let removed =
            crate::bus::message::sanitize_incomplete_tool_call_sequences(&mut current_history);
        if removed > 0 {
            tracing::warn!(
                removed_count = removed,
                remaining_messages = current_history.len(),
                "Post-compression sanitization removed incomplete tool call sequences"
            );
        }

        tracing::info!(
            final_tokens = estimate_tokens(&current_history),
            final_msg_count = current_history.len(),
            "Context compression completed"
        );

        Ok(current_history)
    }

    /// Summarize a segment of messages using LLM.
    async fn summarize_segment(
        &self,
        messages: &[ChatMessage],
        provider_config: &LLMProviderConfig,
    ) -> Result<String, AgentError> {
        if messages.is_empty() {
            return Ok(String::new());
        }

        let runtime_config = AgentRuntimeConfig::from(provider_config.clone());
        let provider = create_provider(runtime_config.provider)
            .map_err(|e| AgentError::ProviderCreation(e.to_string()))?;
        let transcript = Self::build_transcript(messages);

        let result = if char_count(&transcript) <= self.config.summary_max_chars {
            self.summarize_transcript(
                provider.as_ref(),
                &transcript,
                self.config.summary_max_chars,
            )
            .await
        } else {
            self.summarize_chunked_transcript(provider.as_ref(), messages, &transcript)
                .await
        };

        match result {
            Ok(summary) => Ok(summary),
            Err(e) => {
                tracing::warn!(error = %e, "LLM summarization failed, using truncated transcript");
                Ok(take_prefix_chars(
                    &transcript,
                    self.config.summary_max_chars,
                ))
            }
        }
    }
}

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

    #[test]
    fn test_estimate_tokens() {
        let messages = vec![
            ChatMessage::user("Hello"),
            ChatMessage::assistant("Hi there!"),
            ChatMessage::user("How are you?"),
        ];

        let tokens = estimate_tokens(&messages);
        // Content: "Hello" (5) + "Hi there!" (9) + "How are you?" (12) = 26 chars
        // English: 26 / 4 = 6.5 tokens for content
        // JSON overhead: 3 * 50 = 150
        // Total before multiplier: 156.5
        // After 1.2x: ~188 tokens
        assert!(
            tokens > 100 && tokens < 300,
            "Expected ~150-250 tokens for English content, got {}",
            tokens
        );
    }

    #[test]
    fn test_estimate_tokens_chinese_content() {
        let messages = vec![
            ChatMessage::user("你好，这是一个中文测试"),
            ChatMessage::assistant("这是一个中文回复消息"),
        ];

        let tokens = estimate_tokens(&messages);
        // Content: ~20 CJK chars, CJK uses 2 chars/token = ~10 tokens for content
        // JSON overhead: 2 * 50 = 100
        // Total before multiplier: ~110
        // After 1.2x: ~132 tokens
        assert!(
            tokens > 80 && tokens < 200,
            "Expected ~100-180 tokens for Chinese content, got {}",
            tokens
        );
    }

    #[test]
    fn test_estimate_tokens_mixed_content() {
        let messages = vec![
            ChatMessage::user("Hello 世界 this is 测试"),
        ];

        let tokens = estimate_tokens(&messages);
        // Content: 18 English chars + 4 CJK chars
        // English: 18 / 4 = 4.5, CJK: 4 / 2 = 2, content tokens = 6.5
        // JSON overhead: 1 * 50 = 50
        // Total before multiplier: 56.5
        // After 1.2x: ~68 tokens
        assert!(
            tokens > 40 && tokens < 120,
            "Expected ~50-100 tokens for mixed content, got {}",
            tokens
        );
    }

    #[test]
    fn test_chinese_tokens_higher_than_english() {
        // Use more characters to make the content difference significant
        // compared to JSON overhead (50 tokens per message)
        let english = vec![ChatMessage::user(&"abcdefghij".repeat(20))]; // 200 English chars
        let chinese = vec![ChatMessage::user(&"这是一个测试消息字".repeat(20))]; // 200 CJK chars (10 chars * 20)

        let english_tokens = estimate_tokens(&english);
        let chinese_tokens = estimate_tokens(&chinese);

        // 200 English chars: 200/4 = 50 content tokens
        // 200 CJK chars: 200/2 = 100 content tokens
        // With same JSON overhead, Chinese should use ~1.5x tokens
        assert!(
            chinese_tokens > english_tokens * 130 / 100, // At least 1.3x
            "Chinese (200 chars) should use more tokens than English (200 chars): EN={} CN={}",
            english_tokens,
            chinese_tokens
        );
    }

    #[test]
    fn test_estimate_tokens_includes_image_media_refs() {
        let temp_dir = tempfile::tempdir().unwrap();
        let image_path = temp_dir.path().join("demo.jpg");
        std::fs::write(&image_path, vec![0_u8; 12_000]).unwrap();

        let plain = vec![ChatMessage::user("hello")];
        let with_image = vec![ChatMessage::user_with_media(
            "hello",
            vec![image_path.to_string_lossy().to_string()],
        )];

        assert!(estimate_tokens(&with_image) > estimate_tokens(&plain));
    }

    #[test]
    fn test_should_compress() {
        let compressor = ContextCompressor::new(20);
        // Need more content to trigger compression with new weighted calculation
        // 200 English chars / 4 = 50 tokens, plus overhead
        let messages = vec![ChatMessage::user(&"x".repeat(400))];
        assert!(compressor.should_compress(&messages));
    }

    #[test]
    fn test_user_turn_ranges_follow_user_boundaries() {
        let compressor = ContextCompressor::new(100_000);
        let history = vec![
            ChatMessage::system("system"),
            ChatMessage::user("u1"),
            ChatMessage::assistant("a1"),
            ChatMessage::tool("call-1", "bash", "t1"),
            ChatMessage::user("u2"),
            ChatMessage::assistant("a2"),
            ChatMessage::user("u3"),
        ];

        let turns = compressor.user_turn_ranges(&history);
        assert_eq!(
            turns,
            vec![
                UserTurnRange {
                    start: 1,
                    end_exclusive: 4
                },
                UserTurnRange {
                    start: 4,
                    end_exclusive: 6
                },
                UserTurnRange {
                    start: 6,
                    end_exclusive: 7
                },
            ]
        );
    }

    #[test]
    fn test_split_prefix_messages_preserves_key_system_messages() {
        let compressor = ContextCompressor::new(50);
        let prefix = vec![
            ChatMessage::system_with_context(
                "agent prompt",
                Some(SYSTEM_CONTEXT_AGENT_PROMPT.to_string()),
            ),
            ChatMessage::user("u1"),
            ChatMessage::assistant("a1"),
            ChatMessage::system_with_context(
                "scheduled prompt",
                Some(SYSTEM_CONTEXT_SCHEDULED_PROMPT.to_string()),
            ),
        ];

        let (preserved_system_messages, summary_source) = compressor.split_prefix_messages(&prefix);
        assert_eq!(preserved_system_messages.len(), 2);
        assert_eq!(summary_source.len(), 2);
        assert!(preserved_system_messages[0].has_system_context(SYSTEM_CONTEXT_AGENT_PROMPT));
        assert!(preserved_system_messages[1].has_system_context(SYSTEM_CONTEXT_SCHEDULED_PROMPT));
    }

    #[test]
    fn test_threshold() {
        let compressor = ContextCompressor::new(128_000);
        assert_eq!(compressor.threshold(), 89_600); // 70% of 128_000
    }

    #[test]
    fn test_summary_char_budget_for_context_window_scales_and_clamps() {
        assert_eq!(
            ContextCompressor::summary_char_budget_for_context_window(4_096),
            1_500
        );
        assert_eq!(
            ContextCompressor::summary_char_budget_for_context_window(65_536),
            16_384
        );
        assert_eq!(
            ContextCompressor::summary_char_budget_for_context_window(128_000),
            32_000
        );
        assert_eq!(
            ContextCompressor::summary_char_budget_for_context_window(400_000),
            50_000
        );
    }

    #[test]
    fn test_chunk_messages_for_summary_keeps_message_boundaries_when_possible() {
        let messages = vec![
            ChatMessage::user("alpha"),
            ChatMessage::assistant("beta"),
            ChatMessage::user("gamma"),
        ];

        let chunks = ContextCompressor::chunk_messages_for_summary(&messages, 30);

        assert_eq!(chunks.len(), 2);
        assert!(chunks.iter().all(|chunk| char_count(chunk) <= 30));
        assert_eq!(chunks[0], "user: alpha\n\nAssistant: beta");
        assert_eq!(chunks[1], "user: gamma");
    }

    #[test]
    fn test_chunk_messages_for_summary_splits_oversized_message() {
        let messages = vec![ChatMessage::user(&"x".repeat(25))];

        let chunks = ContextCompressor::chunk_messages_for_summary(&messages, 10);

        assert!(chunks.len() > 1);
        assert!(chunks.iter().all(|chunk| char_count(chunk) <= 10));
        assert_eq!(chunks.concat(), "user: xxxxxxxxxxxxxxxxxxxxxxxxx");
    }

    // =========================================================================
    // HistoryUnit / parse_to_units tests
    // =========================================================================

    #[test]
    fn test_parse_to_units_system_guards_preserved() {
        let messages = vec![
            ChatMessage::system_with_context(
                "agent prompt",
                Some(SYSTEM_CONTEXT_AGENT_PROMPT.to_string()),
            ),
            ChatMessage::user("hello"),
            ChatMessage::assistant("hi"),
        ];

        let units = parse_to_units(&messages);
        // SystemGuard + UserMessage + AssistantText = 3
        assert_eq!(units.len(), 3);
        assert!(matches!(units[0], HistoryUnit::SystemGuard(_)));
        assert!(matches!(units[1], HistoryUnit::UserMessage(_)));
        assert!(matches!(units[2], HistoryUnit::AssistantText(_)));
    }

    #[test]
    fn test_parse_to_units_tool_round_atomic() {
        let messages = vec![
            ChatMessage::user("read a file"),
            ChatMessage::assistant_with_tool_calls(
                "let me read",
                vec![crate::domain::messages::ToolCall {
                    id: "call_1".to_string(),
                    name: "file_read".to_string(),
                    arguments: serde_json::json!({"path": "/data/test.txt"}),
                }],
            ),
            ChatMessage::tool("call_1", "file_read", "file contents here"),
            ChatMessage::assistant("the file says hello"),
        ];

        let units = parse_to_units(&messages);
        assert_eq!(units.len(), 3); // UserMessage + ToolRound + AssistantText
        assert!(matches!(units[0], HistoryUnit::UserMessage(_)));
        assert!(matches!(units[1], HistoryUnit::ToolRound { .. }));
        assert!(matches!(units[2], HistoryUnit::AssistantText(_)));

        // Verify ToolRound contains both assistant and results
        if let HistoryUnit::ToolRound { assistant, results } = &units[1] {
            assert_eq!(assistant.role, "assistant");
            assert!(assistant.tool_calls.is_some());
            assert_eq!(results.len(), 1);
            assert_eq!(results[0].role, "tool");
            assert_eq!(results[0].tool_call_id, Some("call_1".to_string()));
        }
    }

    #[test]
    fn test_parse_to_units_orphaned_tool_result_dropped() {
        let messages = vec![
            ChatMessage::user("hello"),
            // Orphaned tool result — no preceding assistant with tool_calls
            ChatMessage::tool("orphan_1", "bash", "some output"),
            ChatMessage::assistant("done"),
        ];

        let units = parse_to_units(&messages);
        // Orphaned tool result should be dropped, leaving UserMessage + AssistantText
        assert_eq!(units.len(), 2);
        assert!(matches!(units[0], HistoryUnit::UserMessage(_)));
        assert!(matches!(units[1], HistoryUnit::AssistantText(_)));
    }

    #[test]
    fn test_parse_to_units_multiple_tool_results_same_round() {
        let messages = vec![
            ChatMessage::user("do multiple things"),
            ChatMessage::assistant_with_tool_calls(
                "doing things",
                vec![
                    crate::domain::messages::ToolCall {
                        id: "call_a".to_string(),
                        name: "bash".to_string(),
                        arguments: serde_json::json!({"command": "ls"}),
                    },
                    crate::domain::messages::ToolCall {
                        id: "call_b".to_string(),
                        name: "file_read".to_string(),
                        arguments: serde_json::json!({"path": "/data/test.txt"}),
                    },
                ],
            ),
            ChatMessage::tool("call_a", "bash", "file1.txt\nfile2.txt"),
            ChatMessage::tool("call_b", "file_read", "hello world"),
        ];

        let units = parse_to_units(&messages);
        assert_eq!(units.len(), 2); // UserMessage + ToolRound
        if let HistoryUnit::ToolRound { results, .. } = &units[1] {
            assert_eq!(results.len(), 2);
            assert_eq!(results[0].tool_call_id, Some("call_a".to_string()));
            assert_eq!(results[1].tool_call_id, Some("call_b".to_string()));
        } else {
            panic!("Expected ToolRound");
        }
    }

    #[test]
    fn test_parse_to_units_empty_input() {
        let units = parse_to_units(&[]);
        assert!(units.is_empty());
    }

    #[test]
    fn test_find_safe_split_point() {
        let compressor = ContextCompressor::new(100_000);
        // Build a list of units with known token sizes
        let messages: Vec<ChatMessage> = (0..10)
            .map(|i| ChatMessage::assistant(&format!("message content number {}", i)))
            .collect();
        let units = parse_to_units(&messages);

        // All units are AssistantText, split at 50% token ratio
        let split = compressor.find_safe_split_point(&units, 0.5);
        // Should split somewhere in the middle (not 0, not len())
        assert!(split > 0 && split < units.len(),
            "split {} should be between 0 and {}", split, units.len());
    }

    #[test]
    fn test_compress_two_segment_no_tool_calls_in_output() {
        // This test verifies the critical invariant:
        // compress_two_segment output MUST NOT contain any tool_calls or tool messages.
        let compressor = ContextCompressor::new(128_000);

        // Build a simple history that's well under the threshold
        // so compress_two_segment returns it unchanged (no LLM call needed).
        let history = vec![
            ChatMessage::system_with_context(
                "You are a helpful assistant.",
                Some(SYSTEM_CONTEXT_AGENT_PROMPT.to_string()),
            ),
            ChatMessage::user("hello"),
            ChatMessage::assistant("hi there"),
        ];

        // Dummy config — won't be used because history is under threshold.
        let config = LLMProviderConfig {
            provider_type: "openai".to_string(),
            name: "test".to_string(),
            base_url: "http://localhost".to_string(),
            api_key: "sk-test".to_string(),
            extra_headers: std::collections::HashMap::new(),
            llm_timeout_secs: 120,
            memory_maintenance_timeout_secs: 300,
            model_id: "test-model".to_string(),
            temperature: None,
            max_tokens: None,
            context_window_tokens: Some(128_000),
            model_extra: std::collections::HashMap::new(),
            max_tool_iterations: 100,
            tool_result_max_chars: 100_000,
            context_tool_result_trim_chars: 2_000,
            max_images_in_context: 10,
            max_image_age_rounds: 50,
        };

        let runtime = tokio::runtime::Runtime::new().unwrap();
        let result = runtime.block_on(compressor.compress_two_segment(&history, &config));
        assert!(result.is_ok());
        let compressed = result.unwrap();
        // Under threshold: should return unchanged (3 messages)
        assert_eq!(compressed.len(), 3);
        // Critical invariant: NO tool_calls or tool_call_id anywhere
        for msg in &compressed {
            assert!(msg.tool_calls.is_none(),
                "compress_two_segment output should never contain tool_calls");
            assert!(msg.tool_call_id.is_none(),
                "compress_two_segment output should never contain tool_call_id");
        }
    }
}