lingshu-core 0.10.0

//! # Context Compression — prevents context-window overflow
//!
//! WHY: Long conversations accumulate tokens until they exceed the
//! model's context window. Rather than hard-truncating (which loses
//! important early context), we summarize old messages while preserving
//! the most recent ones verbatim.
//!
//! ```text
//!   [system] [msg1] [msg2] ... [msgN-20] [msgN-19] ... [msgN]
//!    ↑ keep    └──── prune tools ───────┘
//!               └─── llm_summarize ─────┘  └── keep last 20 ──┘
//! ```
//!
//! Pipeline (v0.4.0 — matching hermes-agent 0.4.x):
//!
//! 1. **Tool output pruning** — replace gigantic tool results in old
//!    messages with `PRUNED_TOOL_PLACEHOLDER` (cheap, no LLM needed).
//!    This alone often halves the prompt size.
//!
//! 2. **Boundary determination** — tail is token-budget based (walks
//!    backward accumulating token estimates until `threshold × target_ratio`
//!    budget is exhausted), with `protect_last_n` as a floor. Boundaries
//!    are aligned backward to avoid splitting tool_call/tool_result groups.
//!
//! 3. **LLM-powered summary** — calls the provider with a structured
//!    8-section template: Goal / Constraints & Preferences / Progress
//!    (Done / In Progress / Blocked) / Key Decisions / Relevant Files /
//!    Next Steps / Critical Context. Output is prefixed with `SUMMARY_PREFIX`
//!    so the next compression pass can locate and update it (iterative
//!    updates). Summary token budget = content_tokens × 0.20, min 2 000,
//!    max min(context_length × 0.05, 12 000).
//!
//! 4. **Structural fallback** — if the LLM call fails, a structured
//!    stat-based summary is built instead (message counts, excerpts).
//!
//! 5. **Orphan sanitization** — after assembling head + summary + tail,
//!    orphaned tool_result messages (no matching tool_call in history)
//!    are removed and orphaned tool_calls get a stub result injected.
//!
//! ## Context pressure warnings
//!
//! When estimated tokens exceed 85 % of the compression threshold the
//! function returns `CompressionStatus::PressureWarning`. After a
//! successful compression that brings usage below 85 % of threshold the
//! status reverts to `CompressionStatus::Ok`.
//!
//! ```text
//!   compress_with_llm(messages, params, provider)
//!       │
//!       ├── prune_tool_outputs(old_messages)      ← step 1 (cheap)
//!       ├── find prior SUMMARY_PREFIX block?       ← iterative update
//!       │       yes → prepend to transcript
//!       │       no  → fresh summary
//!       ├── llm_summarize(pruned_old) → Ok(text) OR Err
//!       │       ↓ on Err
//!       │   build_summary() [structural fallback]
//!       │
//!       └── [Message::system_summary(SUMMARY_PREFIX + text), ...recent]
//! ```

use std::path::Path;
use std::sync::Arc;

use lingshu_types::{Content, ContentPart, Message, Role};
use edgequake_llm::LLMProvider;

use crate::config::CompressionConfig;
use crate::model_catalog::ModelCatalog;
use crate::tool_result_spill::{SpillConfig, SpillOutcome, SpillSequence};

// ─── Constants ────────────────────────────────────────────────────────

/// Prefix for LLM-generated compaction summaries.
///
/// WHY a recognisable prefix: The next compression pass can locate this
/// message and feed it back to the LLM as "prior summary" context so the
/// model produces an *update* rather than starting from scratch. This
/// means summaries improve with each subsequent compaction.
pub const SUMMARY_PREFIX: &str =
    "[CONTEXT COMPACTION] Earlier turns were summarised to reclaim context window space.\n\n";

/// One-shot note for the model after the first context compression (FP33).
///
/// Appended to the cached system prompt during in-loop compression, or to the
/// handoff user message when `/handoff` auto-compresses before a model swap.
pub const FIRST_COMPRESSION_NOTE: &str = concat!(
    "\n\n[Note: Earlier conversation turns have been compacted into a ",
    "handoff summary to stay within the context window. The current ",
    "session state already reflects that earlier work — build on it ",
    "rather than re-doing completed steps.]"
);

/// Inline variant for handoff user messages (no leading newlines).
pub const HANDOFF_COMPRESSION_NOTE: &str =
    "[Note: Earlier turns were auto-compressed for the target model's context window.]";

/// Replacement text for pruned tool output blocks.
///
/// WHY prune first: Tool results (file contents, shell output) can be
/// thousands of tokens each. Replacing them before the LLM call keeps
/// the summarisation prompt itself small — no recursion risk.
pub const PRUNED_TOOL_PLACEHOLDER: &str = "[tool output pruned — reclaimed context window space]";

/// Context for spilling tool results to artifact files during compression.
///
/// When provided to `compress_with_llm` / `prune_tool_outputs`, large tool
/// results are written to disk artifacts instead of being replaced with a
/// generic placeholder. This preserves agent access to the full data via
/// `read_file` while still reclaiming context window space.
pub struct PruneSpillContext<'a> {
    /// Active session ID — used for artifact directory scoping.
    pub session_id: &'a str,
    /// Current working directory — artifact root.
    pub cwd: &'a Path,
    /// Spill configuration (enabled, threshold, preview_lines).
    pub config: &'a SpillConfig,
    /// Shared per-session sequence counter for unique artifact filenames.
    pub seq: &'a SpillSequence,
}

impl<'a> PruneSpillContext<'a> {
    pub fn new(
        session_id: &'a str,
        cwd: &'a Path,
        config: &'a SpillConfig,
        seq: &'a SpillSequence,
    ) -> Self {
        Self {
            session_id,
            cwd,
            config,
            seq,
        }
    }
}

/// Number of head messages (system prompt + first exchange) always preserved.
/// Matches hermes-agent's `protect_first_n = 3` constant.
const PROTECT_FIRST_N: usize = 3;

/// Minimum tokens for the LLM summary budget.
const MIN_SUMMARY_TOKENS: usize = 2_000;

/// Summary token budget as a fraction of compressed content tokens.
const SUMMARY_RATIO: f32 = 0.20;

/// Hard ceiling on summary tokens (absolute maximum).
const SUMMARY_TOKENS_CEILING: usize = 12_000;

/// Approximate characters per token for rough estimation without a tokenizer.
const CHARS_PER_TOKEN: usize = 4;

/// Stub text injected for orphaned tool_calls after compression.
const STUB_TOOL_RESULT: &str = "[Result from earlier conversation — see context summary above]";

/// 8-section structured summary template (hermes-agent 0.4.x format).
const SUMMARY_TEMPLATE: &str = "\
## Goal
[What the user is trying to accomplish]

## Constraints & Preferences
[User preferences, coding style, constraints, important decisions]

## Progress
### Done
[Completed work — include specific file paths, commands run, results obtained]
### In Progress
[Work currently underway]
### Blocked
[Any blockers or issues encountered]

## Key Decisions
[Important technical decisions and why they were made]

## Relevant Files
[Files read, modified, or created — with brief note on each]

## Next Steps
[What needs to happen next to continue the work]

## Critical Context
[Any specific values, error messages, configuration details, or data that would be lost without explicit preservation]";

/// Configuration for context compression.
#[derive(Debug, Clone)]
pub struct CompressionParams {
    /// Estimated context window size for the target model.
    pub context_window: usize,
    /// Compress when estimated tokens exceed this fraction of the window.
    /// Default 0.50 (50 %). Threshold tokens = context_window × threshold.
    pub threshold: f32,
    /// Tail budget ratio: tail_token_budget = threshold_tokens × target_ratio.
    /// Controls how many tokens the "protected recent messages" tail may use.
    /// Default 0.20. Falls back to protect_last_n when the budget would keep
    /// fewer than protect_last_n messages.
    pub target_ratio: f32,
    /// Minimum number of recent messages always kept uncompressed.
    /// Default 20. Acts as a floor when token-budget tail selection would
    /// protect fewer messages.
    pub protect_last_n: usize,
}

const DEFAULT_CONTEXT_WINDOW: usize = 128_000;

impl Default for CompressionParams {
    fn default() -> Self {
        Self {
            context_window: DEFAULT_CONTEXT_WINDOW,
            threshold: 0.50,
            target_ratio: 0.20,
            protect_last_n: 20,
        }
    }
}

impl CompressionParams {
    /// Resolve compression parameters for the active model/configuration.
    pub fn from_model_config(model: &str, cfg: &CompressionConfig) -> Self {
        let context_window = ModelCatalog::context_window_for_spec(model)
            .map(|tokens| tokens as usize)
            .unwrap_or(DEFAULT_CONTEXT_WINDOW);

        Self {
            context_window,
            threshold: cfg.threshold.clamp(0.01, 1.0),
            target_ratio: cfg.target_ratio.clamp(0.01, 1.0),
            protect_last_n: cfg.protect_last_n.max(1),
        }
    }
}

/// Result of a compression trigger check.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CompressionStatus {
    /// Token usage is below the warning threshold.
    Ok,
    /// Approaching compaction: tokens > 85 % of threshold.
    /// Emitted as a UI warning before compression fires.
    PressureWarning,
    /// Compression should fire: tokens ≥ threshold.
    NeedsCompression,
}

/// Estimate token count for a message list.
///
/// WHY ~4 chars/token: This is a rough heuristic that works well
/// for English text across GPT/Claude tokenizers. It's fast (no
/// tokenizer dependency) and good enough for the compression
/// threshold check. Exact counts come from the API response.
pub fn estimate_tokens(messages: &[Message]) -> usize {
    messages.iter().map(estimate_message_tokens).sum()
}

/// Per-message token estimate including image blocks (flat ~1500/image).
pub fn estimate_message_tokens(m: &Message) -> usize {
    use lingshu_types::{Content, ContentPart, MULTIMODAL_IMAGE_TOKEN_ESTIMATE, Role};

    let text_len = m.text_content().len();
    let mut tokens = (text_len / 4) + 4;

    if let Some(Content::Parts(parts)) = &m.content {
        tokens += parts
            .iter()
            .filter(|p| matches!(p, ContentPart::ImageUrl { .. }))
            .count()
            * MULTIMODAL_IMAGE_TOKEN_ESTIMATE;
    }

    if m.role == Role::Tool
        && m.name.as_deref() == Some("computer_use")
        && let Some(Content::Text(text)) = &m.content
        && (lingshu_types::multimodal_has_image(text)
            || lingshu_types::multimodal_disk_image_from_content(text).is_some())
    {
        tokens += MULTIMODAL_IMAGE_TOKEN_ESTIMATE;
    }

    tokens
}

/// Prune stale computer_use screenshots in place (call after each capture tool result).
pub fn maybe_prune_computer_use_screenshots(messages: &mut Vec<Message>, keep_last_n: u32) {
    if keep_last_n > 0 {
        *messages = prune_computer_use_screenshots(messages, keep_last_n);
    }
}

/// Check if compression is needed.
pub fn needs_compression(messages: &[Message], params: &CompressionParams) -> bool {
    matches!(
        check_compression_status(messages, params),
        CompressionStatus::NeedsCompression
    )
}

/// Full compression status check with pressure warning.
///
/// Returns:
/// - `Ok` — below 85 % of threshold
/// - `PressureWarning` — between 85 % and 100 % of threshold (UI warning)
/// - `NeedsCompression` — at or above threshold (compression should fire)
pub fn check_compression_status(
    messages: &[Message],
    params: &CompressionParams,
) -> CompressionStatus {
    let estimated = estimate_tokens(messages);
    check_compression_status_for_estimate(estimated, params)
}

/// Classify compression pressure from a precomputed token estimate.
pub fn check_compression_status_for_estimate(
    estimated: usize,
    params: &CompressionParams,
) -> CompressionStatus {
    let threshold_tokens = (params.context_window as f32 * params.threshold) as usize;
    let warning_tokens = (threshold_tokens as f32 * 0.85) as usize;

    if estimated >= threshold_tokens {
        CompressionStatus::NeedsCompression
    } else if estimated >= warning_tokens {
        CompressionStatus::PressureWarning
    } else {
        CompressionStatus::Ok
    }
}

/// Perform simple compression: summarize old messages into a single
/// system-level summary, keeping the last N messages intact.
///
/// Returns the compressed message list. The summary message is a
/// placeholder — in production, this would call a cheaper LLM to
/// generate a real summary.
///
/// WHY simple truncation for Phase 1: Full LLM-based summarization
/// requires an async call to a summary model and careful chunking.
/// This is deferred to Phase 2. For now, we produce a structured
/// summary stub that preserves the message structure.
pub fn compress_messages(messages: &[Message], params: &CompressionParams) -> Vec<Message> {
    if messages.len() <= params.protect_last_n {
        return messages.to_vec();
    }

    let split_point = messages.len().saturating_sub(params.protect_last_n);
    let old_messages = &messages[..split_point];
    let recent_messages = &messages[split_point..];

    // Build a structured summary of the old messages
    let summary = build_summary(old_messages);

    let mut compressed = Vec::with_capacity(1 + recent_messages.len());
    compressed.push(Message::system_summary(summary));
    compressed.extend_from_slice(recent_messages);

    compressed
}

/// Build a text summary of messages (simple extraction, no LLM).
///
/// Extracts key information: user questions, assistant conclusions,
/// tool calls made. This is a structural summary — the LLM-based
/// summary (Phase 2) will produce a more coherent narrative.
fn build_summary(messages: &[Message]) -> String {
    let mut parts = Vec::new();
    parts.push("[Context Summary — earlier messages compressed]".to_string());

    let mut user_count = 0u32;
    let mut assistant_count = 0u32;
    let mut tool_count = 0u32;

    for m in messages {
        match m.role {
            lingshu_types::Role::User => user_count += 1,
            lingshu_types::Role::Assistant => assistant_count += 1,
            lingshu_types::Role::Tool => tool_count += 1,
            lingshu_types::Role::System => {}
        }
    }

    parts.push(format!(
        "Compressed {user_count} user messages, {assistant_count} assistant \
         responses, and {tool_count} tool results."
    ));

    // Include the first user message for context
    if let Some(first_user) = messages
        .iter()
        .find(|m| m.role == lingshu_types::Role::User)
    {
        let preview = first_user.text_content();
        let truncated = if preview.len() > 200 {
            format!("{}...", crate::safe_truncate(&preview, 200))
        } else {
            preview
        };
        parts.push(format!("First user message: {truncated}"));
    }

    parts.join("\n")
}

// ─── LLM-powered compression ──────────────────────────────────────────

/// LLM-powered context compression (v0.4.0 — hermes-agent parity).
///
/// WHY LLM summarization > structural: A structural summary preserves
/// message counts but loses semantic meaning. An LLM summary produces a
/// coherent narrative the model can use to reason about earlier state.
///
/// Pipeline (6 phases — mirrors hermes-agent `context_compressor.py`):
/// 1. **Prune** — replace large tool outputs with placeholders or spill to
///    artifact files (cheap, no LLM). When `spill_ctx` is provided, results
///    exceeding the prune threshold are written to disk artifacts so the
///    agent can still access them via `read_file`.
/// 2. **Boundary** — determine head/tail by token-budget walk; align both
///    boundaries to avoid splitting tool_call/tool_result groups.
/// 3. **Prior** — extract any existing `SUMMARY_PREFIX` block for iterative update.
/// 4. **Summarise** — call LLM with 8-section template; fall back to structural
///    summary on LLM failure (never silently drops context).
/// 5. **Assemble** — head messages + summary message + tail messages.
///    Role-collision check: if the last head message has the same role as
///    the summary role (system), pick `user` instead to avoid adjacent
///    same-role messages that break strict-alternation providers (FP30).
/// 6. **Sanitize** — remove orphaned tool results; inject stub results for
///    orphaned tool_calls so the assembled list is always API-compliant.
///
/// Returns `(compressed_messages, llm_succeeded)`. The bool is `true` when the
/// LLM summarization call succeeded; `false` when it fell back to structural.
/// Callers use this to implement the circuit breaker (FP29).
pub async fn compress_with_llm(
    messages: &[Message],
    params: &CompressionParams,
    provider: &Arc<dyn LLMProvider>,
    spill_ctx: Option<&PruneSpillContext<'_>>,
) -> (Vec<Message>, bool) {
    let n = messages.len();
    // Need at least: protected head + 1 message to summarise + protected tail.
    if n <= PROTECT_FIRST_N + params.protect_last_n {
        return (messages.to_vec(), true);
    }

    // Phase 1: prune tool outputs (cheap, no LLM).
    // When spill context is available, large results are written to artifact
    // files on disk instead of being replaced with a generic placeholder.
    let pruned = prune_tool_outputs(messages, spill_ctx);

    // Phase 2: determine compression boundaries.
    // Head: always keep system prompt + first exchange (PROTECT_FIRST_N messages).
    let head_end = align_boundary_forward(&pruned, PROTECT_FIRST_N);
    // Tail: walk backward until token budget exhausted.
    let threshold_tokens = (params.context_window as f32 * params.threshold) as usize;
    let tail_token_budget = (threshold_tokens as f32 * params.target_ratio) as usize;
    let tail_start =
        find_tail_cut_by_tokens(&pruned, head_end, tail_token_budget, params.protect_last_n);

    if head_end >= tail_start {
        // Nothing in the middle — history is too short to compress.
        return (messages.to_vec(), true);
    }

    let turns_to_summarize = &pruned[head_end..tail_start];

    // Phase 3: extract prior summary for iterative update.
    let prior_summary = extract_prior_summary(messages);

    // Phase 4: LLM summarization with 8-section template.
    let (summary_text, llm_succeeded) = match llm_summarize(
        turns_to_summarize,
        params.context_window,
        provider,
        prior_summary.as_deref(),
    )
    .await
    {
        Ok(text) => (text, true),
        Err(e) => {
            tracing::warn!(error = %e, "LLM compression failed, using structural fallback");
            (build_summary(turns_to_summarize), false)
        }
    };

    // Phase 5: assemble head + summary + tail.
    //
    // FP30: Role-collision guard — mirrors hermes-agent `context_compressor.py`.
    // If the last head message is already `System` role, injecting another
    // `system_summary` (also System) would create adjacent system messages.
    // Strict-alternation providers (Gemini, Mistral) reject this.
    // Pick `User` role for the summary when the head ends with System.
    let prefixed = format!("{SUMMARY_PREFIX}{summary_text}");
    let last_head_role = pruned
        .get(head_end.saturating_sub(1))
        .map(|m| m.role.clone());
    let summary_msg = if last_head_role == Some(lingshu_types::Role::System) {
        // Head ends with a system message (e.g. a prior summary_prefix block
        // that landed in the protected head window). Use user role to avoid
        // adjacent system+system which breaks strict-alternation providers.
        Message::user(&prefixed)
    } else {
        Message::system_summary(prefixed)
    };
    let mut result = Vec::with_capacity(head_end + 1 + (n - tail_start));
    result.extend_from_slice(&pruned[..head_end]);
    result.push(summary_msg);
    result.extend_from_slice(&pruned[tail_start..]);

    // Phase 6: fix orphaned tool pairs.
    (sanitize_orphan_pairs(result), llm_succeeded)
}

/// Structural-only compression — no LLM call, just prune + summarize stats.
///
/// Used when the compression circuit breaker has tripped (FP12: "Fail once,
/// learn; fail thrice, stop"). Runs the same phases as `compress_with_llm`
/// but replaces Phase 4 (LLM summarization) with `build_summary()`.
///
/// See [specs/improve_plan/16-assessment-round3.md](../../../specs/improve_plan/16-assessment-round3.md).
pub fn compress_structural_only(
    messages: &[Message],
    params: &CompressionParams,
    spill_ctx: Option<&PruneSpillContext<'_>>,
) -> Vec<Message> {
    let n = messages.len();
    if n <= PROTECT_FIRST_N + params.protect_last_n {
        return messages.to_vec();
    }
    let pruned = prune_tool_outputs(messages, spill_ctx);
    let head_end = align_boundary_forward(&pruned, PROTECT_FIRST_N);
    let threshold_tokens = (params.context_window as f32 * params.threshold) as usize;
    let tail_token_budget = (threshold_tokens as f32 * params.target_ratio) as usize;
    let tail_start =
        find_tail_cut_by_tokens(&pruned, head_end, tail_token_budget, params.protect_last_n);
    if head_end >= tail_start {
        return messages.to_vec();
    }
    let turns_to_summarize = &pruned[head_end..tail_start];
    let summary_text = build_summary(turns_to_summarize);
    let prefixed = format!("{SUMMARY_PREFIX}{summary_text}");
    let mut result = Vec::with_capacity(head_end + 1 + (n - tail_start));
    result.extend_from_slice(&pruned[..head_end]);
    result.push(Message::system_summary(prefixed));
    result.extend_from_slice(&pruned[tail_start..]);
    sanitize_orphan_pairs(result)
}

/// Count tool results that [`prune_tool_outputs`] would replace (>200 chars).
pub fn count_long_tool_outputs(messages: &[Message]) -> usize {
    messages
        .iter()
        .filter(|m| m.role == lingshu_types::Role::Tool && m.text_content().len() > 200)
        .count()
}

/// Cheap structural prefill prune for local inference — prune/spill oversized tool results only.
///
/// Returns `(pruned_messages, long_tool_outputs_replaced)`.
pub fn structural_prefill_prune(
    messages: &[Message],
    spill_ctx: Option<&PruneSpillContext<'_>>,
) -> (Vec<Message>, usize) {
    let before = count_long_tool_outputs(messages);
    let pruned = prune_tool_outputs(messages, spill_ctx);
    let after = count_long_tool_outputs(&pruned);
    (pruned, before.saturating_sub(after))
}

/// Metrics from a successful structural tool-output prune (message tokens only).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct StructuralPruneOutcome {
    pub tools_pruned: usize,
    pub message_tokens_before: usize,
    pub message_tokens_after: usize,
    pub long_tool_outputs_remaining: usize,
}

/// Prune oversized tool results when any exist; returns `None` if nothing would change.
pub fn apply_structural_tool_output_prune(
    messages: &[Message],
    spill_ctx: Option<&PruneSpillContext<'_>>,
) -> Option<(Vec<Message>, StructuralPruneOutcome)> {
    if count_long_tool_outputs(messages) == 0 {
        return None;
    }
    let message_tokens_before = estimate_tokens(messages);
    let (pruned, tools_pruned) = structural_prefill_prune(messages, spill_ctx);
    if tools_pruned == 0 {
        return None;
    }
    let message_tokens_after = estimate_tokens(&pruned);
    let long_tool_outputs_remaining = count_long_tool_outputs(&pruned);
    Some((
        pruned,
        StructuralPruneOutcome {
            tools_pruned,
            message_tokens_before,
            message_tokens_after,
            long_tool_outputs_remaining,
        },
    ))
}

/// Replace large tool-result messages with a placeholder, or spill to disk.
///
/// WHY: Tool outputs (file contents, grep results, command output) are
/// often thousands of tokens. Replacing them with a 10-token placeholder
/// before summarisation halves the LLM input cost at no semantic loss —
/// the summary will describe *what* the tool found, not dump raw bytes.
///
/// When `spill_ctx` is `Some` and spilling is enabled, large results are
/// written to artifact files on disk instead of being replaced with a
/// generic placeholder. The message body becomes a compact stub with a
/// preview and a file path — the agent can still `read_file` the full data.
///
/// Threshold: tool results over 200 chars are pruned. This preserves
/// short "ok" / "error" responses that carry semantic meaning.
pub fn prune_tool_outputs(
    messages: &[Message],
    spill_ctx: Option<&PruneSpillContext<'_>>,
) -> Vec<Message> {
    messages
        .iter()
        .map(|m| {
            if m.role == lingshu_types::Role::Tool && m.text_content().len() > 200 {
                let tool_call_id = m.tool_call_id.as_deref().unwrap_or("unknown");
                let tool_name = m.name.as_deref().unwrap_or("tool");

                // When spill context is available, attempt to spill to artifact
                if let Some(ctx) = spill_ctx
                    && ctx.config.enabled
                {
                    let result = m.text_content();
                    match crate::tool_result_spill::maybe_spill(
                        tool_name,
                        tool_call_id,
                        result,
                        ctx.session_id,
                        ctx.cwd,
                        ctx.config,
                        ctx.seq,
                    ) {
                        SpillOutcome::Spilled { stub, .. } => {
                            return Message::tool_result(tool_call_id, tool_name, &stub);
                        }
                        SpillOutcome::Inline(_) => {
                            // Below spill threshold but above prune threshold (200 chars)
                            // — fall through to placeholder
                        }
                    }
                }

                // Default: replace with generic placeholder
                Message::tool_result(tool_call_id, tool_name, PRUNED_TOOL_PLACEHOLDER)
            } else {
                m.clone()
            }
        })
        .collect()
}

/// Strip `computer_use` screenshot images from older tool results, keeping only
/// the last `keep_last_n` captures with image parts.
///
/// WHY every turn (not only on compress): a 1280×800 PNG is ~1–1.5k image tokens;
/// three stale screenshots can dominate the context window before the compressor
/// fires. Hermes spec requires aggressive screenshot pruning regardless of
/// compression state.
pub fn prune_computer_use_screenshots(messages: &[Message], keep_last_n: u32) -> Vec<Message> {
    if keep_last_n == 0 {
        return messages
            .iter()
            .map(strip_computer_use_screenshot_images)
            .collect();
    }

    let mut screenshot_indices = Vec::new();
    for (i, m) in messages.iter().enumerate() {
        if computer_use_message_has_screenshot(m) {
            screenshot_indices.push(i);
        }
    }

    let keep = keep_last_n as usize;
    if screenshot_indices.len() <= keep {
        return messages.to_vec();
    }

    let strip_count = screenshot_indices.len().saturating_sub(keep);
    let strip_set: std::collections::HashSet<usize> = screenshot_indices
        .iter()
        .take(strip_count)
        .copied()
        .collect();

    messages
        .iter()
        .enumerate()
        .map(|(i, m)| {
            if strip_set.contains(&i) {
                strip_computer_use_screenshot_images(m)
            } else {
                m.clone()
            }
        })
        .collect()
}

fn computer_use_message_has_screenshot(msg: &Message) -> bool {
    if msg.role != Role::Tool || msg.name.as_deref() != Some("computer_use") {
        return false;
    }
    match &msg.content {
        Some(Content::Parts(parts)) => parts
            .iter()
            .any(|p| matches!(p, ContentPart::ImageUrl { .. })),
        Some(Content::Text(text)) => lingshu_types::capture_has_image_reference(text),
        _ => false,
    }
}

fn strip_computer_use_screenshot_images(msg: &Message) -> Message {
    if !computer_use_message_has_screenshot(msg) {
        return msg.clone();
    }

    let tool_call_id = msg.tool_call_id.as_deref().unwrap_or("unknown");
    let name = msg.name.as_deref().unwrap_or("computer_use");

    match &msg.content {
        Some(Content::Parts(parts)) => {
            let mut text_parts: Vec<String> = parts
                .iter()
                .filter_map(|p| match p {
                    ContentPart::Text { text } => Some(text.clone()),
                    _ => None,
                })
                .collect();
            if !text_parts.iter().any(|t| t.contains("[screenshot pruned]")) {
                text_parts.push("[screenshot pruned — retained text summary only]".into());
            }
            Message {
                role: Role::Tool,
                content: Some(Content::Text(text_parts.join("\n"))),
                tool_call_id: Some(tool_call_id.to_string()),
                name: Some(name.to_string()),
                ..Default::default()
            }
        }
        Some(Content::Text(text)) => {
            if lingshu_types::is_multimodal_tool_json(text) {
                let summary = lingshu_types::multimodal_text_summary(text).unwrap_or_default();
                let body = if summary.is_empty() {
                    "[screenshot pruned — retained text summary only]".into()
                } else {
                    format!("{summary}\n[screenshot pruned — retained text summary only]")
                };
                return Message::tool_result(tool_call_id, name, &body);
            }
            msg.clone()
        }
        _ => msg.clone(),
    }
}

/// Extract the text of the most recent SUMMARY_PREFIX block, if any.
///
/// WHY: Iterative update means the second compression pass feeds the
/// prior summary back to the LLM as existing context. The LLM can then
/// produce an *incremental update* rather than re-summarising everything
/// from scratch, which is both cheaper and more coherent.
fn extract_prior_summary(messages: &[Message]) -> Option<String> {
    messages
        .iter()
        .find(|m| {
            m.role == lingshu_types::Role::System && m.text_content().starts_with(SUMMARY_PREFIX)
        })
        .map(|m| {
            m.text_content()
                .strip_prefix(SUMMARY_PREFIX)
                .unwrap_or(&m.text_content())
                .to_string()
        })
}

// ─── Boundary alignment helpers ──────────────────────────────────────

/// Slide `idx` forward past any leading tool-result messages.
///
/// WHY: If the head boundary lands on a tool result, the preceding
/// assistant tool_call has been preserved but the result would fall into
/// the middle (summarized) region, splitting the pair. Moving forward
/// ensures we start at a clean message boundary.
fn align_boundary_forward(messages: &[Message], idx: usize) -> usize {
    let mut i = idx;
    while i < messages.len() && messages[i].role == lingshu_types::Role::Tool {
        i += 1;
    }
    i
}

/// Pull `idx` backward past any trailing tool results to the parent assistant.
///
/// WHY: If the tail-start boundary falls inside a tool_call/result group,
/// dropping the parent assistant message would create orphaned tool results
/// that the API rejects. Walking backward to the parent assistant ensures
/// the whole group is either kept or summarized together.
fn align_boundary_backward(messages: &[Message], idx: usize) -> usize {
    if idx == 0 || idx >= messages.len() {
        return idx;
    }
    // Walk backward past consecutive tool results.
    let mut check = idx.saturating_sub(1);
    while check > 0 && messages[check].role == lingshu_types::Role::Tool {
        check -= 1;
    }
    // If the parent is an assistant with tool_calls, pull boundary before it.
    if messages[check].role == lingshu_types::Role::Assistant && messages[check].has_tool_calls() {
        check
    } else {
        idx
    }
}

// ─── Token-budget tail selection ─────────────────────────────────────

/// Walk backward from the end of `messages`, accumulating token estimates,
/// and return the index where the protected tail starts.
///
/// WHY token-budget tail instead of fixed `protect_last_n`: A fixed count
/// fails on large models (20 short messages ≪ 20 K tokens) and on small
/// ones (20 long tool outputs may fill the context window). A budget-scaled
/// tail self-adjusts to model context size and message density.
///
/// Falls back to `protect_last_n` if the budget would protect the entire
/// history (small conversation) or fewer than `protect_last_n` messages.
fn find_tail_cut_by_tokens(
    messages: &[Message],
    head_end: usize,
    token_budget: usize,
    protect_last_n: usize,
) -> usize {
    let n = messages.len();
    let mut accumulated: usize = 0;
    let mut cut_idx = n;

    for i in (head_end..n).rev() {
        let msg_tokens = messages[i].text_content().len() / CHARS_PER_TOKEN + 10;
        let protected_count = n - i;
        if accumulated + msg_tokens > token_budget && protected_count >= protect_last_n {
            break;
        }
        accumulated += msg_tokens;
        cut_idx = i;
    }

    // Enforce minimum tail of `protect_last_n` messages.
    let fallback = n.saturating_sub(protect_last_n);
    let cut_idx = cut_idx.min(fallback);

    // If budget swallowed everything (small history), use fixed fallback.
    let cut_idx = if cut_idx <= head_end {
        fallback
    } else {
        cut_idx
    };

    // Align: never split a tool_call/tool_result group at the tail boundary.
    let cut_idx = align_boundary_backward(messages, cut_idx);

    // Always leave at least one message in the middle to compress.
    cut_idx.max(head_end + 1)
}

/// Ensure every assistant `tool_call` has a matching tool result before an API request.
///
/// Public wrapper used by the conversation loop (not only compression).
pub fn ensure_api_safe_tool_pairs(messages: Vec<Message>) -> Vec<Message> {
    sanitize_orphan_pairs(messages)
}

/// Fix orphaned tool_call / tool_result pairs after assembling the compressed list.
///
/// Two failure modes that this resolves:
/// 1. A tool *result* references a call_id whose parent assistant `tool_call`
///    was summarized away → API rejects "No tool_call found for call_id …".
/// 2. An assistant message has `tool_calls` whose results were dropped →
///    API rejects because every tool_call must have a matching result message.
///
/// Removes orphaned results (case 1) and injects one-line stub results for
/// orphaned calls (case 2) so the assembled list is always API-compliant.
fn sanitize_orphan_pairs(messages: Vec<Message>) -> Vec<Message> {
    use std::collections::HashSet;

    // Surviving call IDs present in assistant messages.
    let call_ids: HashSet<String> = messages
        .iter()
        .filter(|m| m.role == lingshu_types::Role::Assistant)
        .flat_map(|m| m.tool_calls.iter().flatten().map(|tc| tc.id.clone()))
        .collect();

    // Call IDs referenced by existing tool result messages.
    let result_ids: HashSet<String> = messages
        .iter()
        .filter(|m| m.role == lingshu_types::Role::Tool)
        .filter_map(|m| m.tool_call_id.clone())
        .collect();

    // Phase 1: drop orphaned tool results (result references a missing call).
    let orphaned_results: HashSet<String> = result_ids.difference(&call_ids).cloned().collect();
    let messages: Vec<Message> = if orphaned_results.is_empty() {
        messages
    } else {
        tracing::debug!(
            count = orphaned_results.len(),
            "sanitizer: dropped orphaned tool results"
        );
        messages
            .into_iter()
            .filter(|m| {
                m.role != lingshu_types::Role::Tool
                    || m.tool_call_id
                        .as_ref()
                        .map(|id| !orphaned_results.contains(id))
                        .unwrap_or(true)
            })
            .collect()
    };

    // Rebuild remaining result IDs after phase-1 filtering.
    let result_ids_after: HashSet<String> = messages
        .iter()
        .filter(|m| m.role == lingshu_types::Role::Tool)
        .filter_map(|m| m.tool_call_id.clone())
        .collect();

    // Phase 2: inject stub results for tool_calls that lost their result.
    let missing_results: HashSet<String> =
        call_ids.difference(&result_ids_after).cloned().collect();
    if missing_results.is_empty() {
        return messages;
    }

    tracing::debug!(
        count = missing_results.len(),
        "sanitizer: injected stub tool results"
    );
    let mut patched = Vec::with_capacity(messages.len() + missing_results.len());
    for m in messages {
        let is_assistant = m.role == lingshu_types::Role::Assistant;
        let tool_calls = m.tool_calls.clone();
        patched.push(m);
        if is_assistant && let Some(tcs) = tool_calls {
            for tc in tcs {
                if missing_results.contains(&tc.id) {
                    patched.push(Message::tool_result(
                        &tc.id,
                        &tc.function.name,
                        STUB_TOOL_RESULT,
                    ));
                }
            }
        }
    }
    patched
}

// ─── Summary budget & serialization ──────────────────────────────────

/// Scale the LLM summary token budget with content size and model context window.
///
/// Formula: `content_tokens × SUMMARY_RATIO`, clamped to
/// `[MIN_SUMMARY_TOKENS, min(context_window × 0.05, SUMMARY_TOKENS_CEILING)]`.
///
/// WHY scaled not fixed: Small conversations need small summaries; large-context
/// models (200 K+ tokens) deserve richer summaries. The ceiling prevents cost runaway.
fn compute_summary_budget(content_tokens: usize, context_window: usize) -> usize {
    let budget = (content_tokens as f32 * SUMMARY_RATIO) as usize;
    let ceiling = ((context_window as f32 * 0.05) as usize).min(SUMMARY_TOKENS_CEILING);
    budget.max(MIN_SUMMARY_TOKENS).min(ceiling)
}

/// Serialize conversation turns into labeled text for the summarizer LLM.
///
/// Includes tool call arguments and result content (truncated to 3 000 chars
/// per message) so the summarizer can capture file paths, commands, outputs.
/// System messages are excluded because they are not conversation history.
fn serialize_for_summary(messages: &[Message]) -> String {
    const MAX_MSG_CHARS: usize = 3_000;
    const HEAD_CHARS: usize = 2_000;
    const TAIL_CHARS: usize = 800;

    messages
        .iter()
        .filter(|m| m.role != lingshu_types::Role::System)
        .map(|m| {
            let text = m.text_content();
            let content = if text.len() > MAX_MSG_CHARS {
                let head = crate::safe_truncate(&text, HEAD_CHARS.min(text.len()));
                let tail_start =
                    crate::safe_char_start(&text, text.len().saturating_sub(TAIL_CHARS));
                format!("{}…[truncated]…{}", head, &text[tail_start..])
            } else {
                text
            };
            match m.role {
                lingshu_types::Role::Tool => {
                    let id = m.tool_call_id.as_deref().unwrap_or("");
                    format!("[TOOL RESULT {id}]: {content}")
                }
                lingshu_types::Role::Assistant => {
                    let mut line = format!("[ASSISTANT]: {content}");
                    if let Some(tcs) = &m.tool_calls {
                        let calls: Vec<String> = tcs
                            .iter()
                            .map(|tc| {
                                let args = if tc.function.arguments.len() > 500 {
                                    format!(
                                        "{}…",
                                        crate::safe_truncate(&tc.function.arguments, 400)
                                    )
                                } else {
                                    tc.function.arguments.clone()
                                };
                                format!("  {}({})", tc.function.name, args)
                            })
                            .collect();
                        line.push_str("\n[Tool calls:\n");
                        line.push_str(&calls.join("\n"));
                        line.push(']');
                    }
                    line
                }
                lingshu_types::Role::User => format!("[USER]: {content}"),
                lingshu_types::Role::System => unreachable!("filtered above"),
            }
        })
        .collect::<Vec<_>>()
        .join("\n\n")
}

// ─── LLM summarization ────────────────────────────────────────────────

/// Call the provider to produce a structured 8-section summary of old messages.
///
/// Sections: Goal / Constraints & Preferences / Progress (Done / In Progress /
/// Blocked) / Key Decisions / Relevant Files / Next Steps / Critical Context.
///
/// When `prior_summary` is `Some`, the prompt asks for an *iterative update*
/// rather than a fresh summary — cheaper, more coherent across repeated passes.
///
/// `max_tokens` = `compute_summary_budget(content_tokens, context_window) × 2`
/// to give the model headroom; the provider truncates the response if needed.
async fn llm_summarize(
    messages: &[Message],
    context_window: usize,
    provider: &Arc<dyn LLMProvider>,
    prior_summary: Option<&str>,
) -> Result<String, edgequake_llm::LlmError> {
    let content = serialize_for_summary(messages);
    let content_tokens = estimate_tokens(messages);
    let summary_budget = compute_summary_budget(content_tokens, context_window);

    let prompt = match prior_summary {
        Some(prior) => format!(
            "You are updating a context compaction summary. A previous compaction produced \
             the summary below. New conversation turns have occurred since then and need to \
             be incorporated.\n\n\
             PREVIOUS SUMMARY:\n{prior}\n\n\
             NEW TURNS TO INCORPORATE:\n{content}\n\n\
             Update the summary using this exact structure. PRESERVE all existing information \
             that is still relevant. ADD new progress. Move items from \"In Progress\" to \
             \"Done\" when completed. Remove information only if it is clearly obsolete.\n\n\
             {SUMMARY_TEMPLATE}\n\n\
             Target ~{summary_budget} tokens. Be specific — include file paths, command \
             outputs, error messages, and concrete values rather than vague descriptions.\n\n\
             Write only the summary body. Do not include any preamble or prefix."
        ),
        None => format!(
            "Create a structured handoff summary for a later assistant that will continue \
             this conversation after earlier turns are compacted.\n\n\
             TURNS TO SUMMARIZE:\n{content}\n\n\
             Use this exact structure:\n\n\
             {SUMMARY_TEMPLATE}\n\n\
             Target ~{summary_budget} tokens. Be specific — include file paths, command \
             outputs, error messages, and concrete values rather than vague descriptions. \
             The goal is to prevent the next assistant from repeating work or losing \
             important details.\n\n\
             Write only the summary body. Do not include any preamble or prefix."
        ),
    };

    let options = edgequake_llm::CompletionOptions {
        max_tokens: Some(summary_budget * 2),
        temperature: Some(0.3),
        ..Default::default()
    };
    let llm_messages = vec![edgequake_llm::ChatMessage::user(&prompt)];
    let response = provider.chat(&llm_messages, Some(&options)).await?;
    Ok(response.content.trim().to_string())
}

// ─── Tests ────────────────────────────────────────────────────────────

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

    fn make_messages(n: usize) -> Vec<Message> {
        (0..n)
            .map(|i| {
                if i % 2 == 0 {
                    Message::user(&format!("question {i}"))
                } else {
                    Message::assistant(&format!("answer {i}"))
                }
            })
            .collect()
    }

    #[test]
    fn estimate_tokens_basic() {
        let msgs = vec![Message::user("hello world")]; // 11 chars → ~2 tokens + 4 overhead
        let tokens = estimate_tokens(&msgs);
        assert!(tokens > 0);
        assert!(tokens < 20);
    }

    #[test]
    fn needs_compression_under_threshold() {
        let msgs = make_messages(5);
        let params = CompressionParams {
            context_window: 128_000,
            threshold: 0.50,
            target_ratio: 0.20,
            protect_last_n: 20,
        };
        assert!(!needs_compression(&msgs, &params));
    }

    #[test]
    fn needs_compression_over_threshold() {
        let msgs: Vec<Message> = (0..1000)
            .map(|i| Message::user(&format!("{}{}", "a".repeat(500), i)))
            .collect();
        let params = CompressionParams {
            context_window: 1000, // small window
            threshold: 0.10,
            target_ratio: 0.20,
            protect_last_n: 5,
        };
        assert!(needs_compression(&msgs, &params));
    }

    #[test]
    fn check_status_pressure_warning() {
        // threshold_tokens = 1000 * 0.50 = 500; warning_tokens = 500 * 0.85 = 425.
        // We need estimate > 425 and < 500.
        // estimate_tokens for one 1700-char message = 1700/4 + 4 = 429. ✓
        let msgs = vec![Message::user(&"x".repeat(1_700))];
        let params = CompressionParams {
            context_window: 1_000,
            threshold: 0.50,
            target_ratio: 0.20,
            protect_last_n: 5,
        };
        assert_eq!(
            check_compression_status(&msgs, &params),
            CompressionStatus::PressureWarning
        );
    }

    #[test]
    fn check_status_needs_compression() {
        let msgs: Vec<Message> = (0..1000)
            .map(|i| Message::user(&"a".repeat(500 + i)))
            .collect();
        let params = CompressionParams {
            context_window: 1_000,
            threshold: 0.10,
            target_ratio: 0.20,
            protect_last_n: 5,
        };
        assert_eq!(
            check_compression_status(&msgs, &params),
            CompressionStatus::NeedsCompression
        );
    }

    #[test]
    fn check_status_ok_below_warning() {
        let msgs = make_messages(2);
        let params = CompressionParams::default();
        assert_eq!(
            check_compression_status(&msgs, &params),
            CompressionStatus::Ok
        );
    }

    #[test]
    fn check_status_for_estimate_reuses_threshold_logic() {
        let params = CompressionParams {
            context_window: 1_000,
            threshold: 0.50,
            target_ratio: 0.20,
            protect_last_n: 5,
        };
        assert_eq!(
            check_compression_status_for_estimate(430, &params),
            CompressionStatus::PressureWarning
        );
        assert_eq!(
            check_compression_status_for_estimate(500, &params),
            CompressionStatus::NeedsCompression
        );
    }

    #[test]
    fn compression_params_from_model_config_uses_runtime_values() {
        let cfg = CompressionConfig {
            enabled: true,
            threshold: 0.75,
            target_ratio: 0.33,
            protect_last_n: 12,
            summary_model: None,
        };
        let params = CompressionParams::from_model_config("anthropic/claude-opus-4.6", &cfg);
        assert_eq!(params.threshold, 0.75);
        assert_eq!(params.target_ratio, 0.33);
        assert_eq!(params.protect_last_n, 12);
        assert_eq!(
            params.context_window,
            ModelCatalog::context_window("anthropic", "claude-opus-4.6").expect("catalog context")
                as usize
        );
    }

    #[test]
    fn compress_preserves_recent() {
        let msgs = make_messages(30);
        let params = CompressionParams {
            protect_last_n: 10,
            ..Default::default()
        };

        let compressed = compress_messages(&msgs, &params);
        // 1 summary + 10 recent = 11
        assert_eq!(compressed.len(), 11);

        // First message should be the summary
        assert_eq!(compressed[0].role, lingshu_types::Role::System);
        assert!(compressed[0].text_content().contains("Context Summary"));

        // Last message should be the last original message
        assert_eq!(
            compressed.last().expect("last").text_content(),
            msgs.last().expect("last").text_content()
        );
    }

    #[test]
    fn compress_small_history_is_noop() {
        let msgs = make_messages(5);
        let params = CompressionParams {
            protect_last_n: 20,
            ..Default::default()
        };
        let compressed = compress_messages(&msgs, &params);
        assert_eq!(compressed.len(), msgs.len());
    }

    #[test]
    fn summary_contains_counts() {
        let msgs = make_messages(10);
        let summary = build_summary(&msgs);
        assert!(summary.contains("5 user messages"));
        assert!(summary.contains("5 assistant responses"));
    }

    // ── Boundary helpers ──────────────────────────────────────────────

    #[test]
    fn align_forward_skips_leading_tool_messages() {
        let msgs = vec![
            Message::user("q"),
            Message::tool_result("c1", "t", "r1"),
            Message::tool_result("c2", "t", "r2"),
            Message::user("follow-up"),
        ];
        assert_eq!(align_boundary_forward(&msgs, 1), 3);
        assert_eq!(align_boundary_forward(&msgs, 0), 0);
        assert_eq!(align_boundary_forward(&msgs, 4), 4); // past end
    }

    #[test]
    fn align_backward_pulls_before_assistant_with_tool_calls() {
        let tc = lingshu_types::ToolCall {
            id: "c1".into(),
            r#type: "function".into(),
            function: lingshu_types::FunctionCall {
                name: "my_tool".into(),
                arguments: "{}".into(),
            },
            thought_signature: None,
        };
        let msgs = vec![
            Message::user("q"),
            Message::assistant_with_tool_calls("", vec![tc]),
            Message::tool_result("c1", "my_tool", "result"),
            Message::user("next"),
        ];
        // Boundary at index 3 should pull before the assistant (index 1).
        assert_eq!(align_boundary_backward(&msgs, 3), 1);
        // Edge cases: 0 and past-end stay unchanged.
        assert_eq!(align_boundary_backward(&msgs, 0), 0);
    }

    #[test]
    fn align_backward_noop_without_tool_calls() {
        // Assistant without tool_calls — boundary should not move.
        let msgs = vec![
            Message::user("q"),
            Message::assistant("a"),
            Message::user("next"),
        ];
        assert_eq!(align_boundary_backward(&msgs, 2), 2);
    }

    #[test]
    fn find_tail_cut_returns_more_than_head_end() {
        let msgs = make_messages(10);
        let cut = find_tail_cut_by_tokens(&msgs, 2, 0, 2);
        assert!(cut > 2, "cut={cut} must be > head_end=2");
        assert!(cut <= msgs.len());
    }

    #[test]
    fn find_tail_cut_respects_protect_last_n() {
        let msgs = make_messages(20);
        // With a huge budget, fallback to protect_last_n=5.
        let cut = find_tail_cut_by_tokens(&msgs, 0, usize::MAX, 5);
        // cut should be at most n - protect_last_n = 15
        assert!(cut <= 15, "cut={cut}");
    }

    // ── Orphan sanitization ───────────────────────────────────────────

    #[test]
    fn sanitize_removes_orphaned_tool_result() {
        // Tool result with no matching assistant tool_call → removed.
        let messages = vec![
            Message::user("do something"),
            Message::tool_result("call_999", "some_tool", "output"),
        ];
        let sanitized = sanitize_orphan_pairs(messages);
        assert_eq!(sanitized.len(), 1);
        assert_eq!(sanitized[0].role, lingshu_types::Role::User);
    }

    #[test]
    fn sanitize_injects_stub_for_missing_tool_result() {
        // Assistant with tool_call but no matching result → stub injected.
        let tc = lingshu_types::ToolCall {
            id: "call_1".into(),
            r#type: "function".into(),
            function: lingshu_types::FunctionCall {
                name: "my_tool".into(),
                arguments: "{}".into(),
            },
            thought_signature: None,
        };
        let messages = vec![
            Message::user("do something"),
            Message::assistant_with_tool_calls("", vec![tc]),
        ];
        let sanitized = sanitize_orphan_pairs(messages);
        // user + assistant + stub tool result
        assert_eq!(sanitized.len(), 3);
        assert_eq!(sanitized[2].role, lingshu_types::Role::Tool);
        assert_eq!(sanitized[2].tool_call_id.as_deref(), Some("call_1"));
        assert!(sanitized[2].text_content().contains("earlier conversation"));
    }

    #[test]
    fn sanitize_noop_on_well_formed_pairs() {
        // Well-formed assistant + result → unchanged.
        let tc = lingshu_types::ToolCall {
            id: "call_x".into(),
            r#type: "function".into(),
            function: lingshu_types::FunctionCall {
                name: "search".into(),
                arguments: "{}".into(),
            },
            thought_signature: None,
        };
        let messages = vec![
            Message::user("query"),
            Message::assistant_with_tool_calls("", vec![tc]),
            Message::tool_result("call_x", "search", "results"),
        ];
        let len = messages.len();
        let sanitized = sanitize_orphan_pairs(messages);
        assert_eq!(sanitized.len(), len);
    }

    #[test]
    fn sanitize_empty_input_is_noop() {
        let sanitized = sanitize_orphan_pairs(vec![]);
        assert!(sanitized.is_empty());
    }

    // ── Summary budget ─────────────────────────────────────────────────

    #[test]
    fn budget_clamps_to_minimum() {
        // Tiny content → floor at MIN_SUMMARY_TOKENS.
        assert_eq!(compute_summary_budget(10, 128_000), MIN_SUMMARY_TOKENS);
    }

    #[test]
    fn budget_clamps_to_ceiling_from_context() {
        // ceiling = min(128_000 * 0.05, 12_000) = min(6_400, 12_000) = 6_400
        let budget = compute_summary_budget(1_000_000, 128_000);
        assert_eq!(budget, 6_400);
    }

    #[test]
    fn budget_hard_cap_limits_huge_windows() {
        // With a very large context window the 12_000 hard cap must kick in.
        let budget = compute_summary_budget(1_000_000, 4_000_000);
        assert!(budget <= SUMMARY_TOKENS_CEILING, "budget={budget}");
    }

    // ── Serialize for summary ──────────────────────────────────────────

    #[test]
    fn serialize_labels_user_and_assistant() {
        let msgs = vec![Message::user("hello"), Message::assistant("world")];
        let text = serialize_for_summary(&msgs);
        assert!(text.contains("[USER]: hello"), "text={text}");
        assert!(text.contains("[ASSISTANT]: world"), "text={text}");
    }

    #[test]
    fn serialize_skips_system_messages() {
        let msgs = vec![Message::system("You are an AI"), Message::user("hi")];
        let text = serialize_for_summary(&msgs);
        assert!(!text.contains("You are an AI"));
        assert!(text.contains("[USER]: hi"));
    }

    #[test]
    fn serialize_truncates_long_content() {
        let long_content = "z".repeat(5_000);
        let msgs = vec![Message::user(&long_content)];
        let text = serialize_for_summary(&msgs);
        assert!(
            text.contains("[truncated]"),
            "should truncate long messages"
        );
    }

    #[test]
    fn serialize_truncates_long_unicode_content_without_panicking() {
        let prefix = "z".repeat(1_999);
        let long_content = format!("{prefix}étail{}", "y".repeat(5_000));
        let msgs = vec![Message::user(&long_content)];
        let text = serialize_for_summary(&msgs);
        assert!(text.contains("[truncated]"));
        assert!(!text.contains('🧠'));
    }

    #[test]
    fn summary_includes_first_user_message() {
        let msgs = vec![
            Message::user("What is the meaning of life?"),
            Message::assistant("42"),
        ];
        let summary = build_summary(&msgs);
        assert!(summary.contains("What is the meaning of life?"));
    }

    #[test]
    fn summary_truncates_long_first_message() {
        let long_msg = "x".repeat(500);
        let msgs = vec![Message::user(&long_msg)];
        let summary = build_summary(&msgs);
        assert!(summary.contains("..."));
        assert!(summary.len() < 600);
    }

    // ── New v0.4.0 tests ──────────────────────────────────────────────

    #[test]
    fn summary_prefix_constant_starts_correctly() {
        assert!(SUMMARY_PREFIX.starts_with("[CONTEXT COMPACTION]"));
    }

    #[test]
    fn pruned_tool_placeholder_is_short() {
        // Must fit in a single token budget line
        assert!(PRUNED_TOOL_PLACEHOLDER.len() < 100);
    }

    #[test]
    fn prune_tool_outputs_replaces_long_results() {
        let messages = vec![
            Message::user("run a command"),
            Message::tool_result("id1", "shell_exec", &"x".repeat(500)),
        ];
        let pruned = prune_tool_outputs(&messages, None);
        assert_eq!(pruned.len(), 2);
        // User message unchanged
        assert_eq!(pruned[0].text_content(), "run a command");
        // Tool result replaced with placeholder
        assert_eq!(pruned[1].text_content(), PRUNED_TOOL_PLACEHOLDER);
    }

    #[test]
    fn structural_prefill_prune_reclaims_tool_output_tokens() {
        let messages: Vec<Message> = (0..8)
            .map(|i| {
                Message::tool_result(
                    &format!("id{i}"),
                    "web_extract",
                    &format!("page body {}\n", "x".repeat(8_000)),
                )
            })
            .collect();
        let tokens_before = estimate_tokens(&messages);
        assert_eq!(count_long_tool_outputs(&messages), 8);

        let (pruned, replaced) = structural_prefill_prune(&messages, None);
        assert_eq!(replaced, 8);
        assert_eq!(count_long_tool_outputs(&pruned), 0);
        let tokens_after = estimate_tokens(&pruned);
        assert!(
            tokens_after < tokens_before / 4,
            "expected large token drop: before={tokens_before} after={tokens_after}"
        );
    }

    #[test]
    fn apply_structural_tool_output_prune_returns_none_when_nothing_long() {
        let messages = vec![Message::tool_result("id", "shell_exec", "ok")];
        assert!(apply_structural_tool_output_prune(&messages, None).is_none());
    }

    #[test]
    fn apply_structural_tool_output_prune_reports_outcome() {
        let messages = vec![Message::tool_result("id", "web_extract", &"x".repeat(500))];
        let (pruned, outcome) = apply_structural_tool_output_prune(&messages, None)
            .expect("long tool output should prune");
        assert_eq!(outcome.tools_pruned, 1);
        assert_eq!(outcome.long_tool_outputs_remaining, 0);
        assert!(outcome.message_tokens_after < outcome.message_tokens_before);
        assert_eq!(count_long_tool_outputs(&pruned), 0);
    }

    #[test]
    fn prune_tool_outputs_keeps_short_results() {
        let messages = vec![Message::tool_result("id1", "shell_exec", "ok")];
        let pruned = prune_tool_outputs(&messages, None);
        assert_eq!(pruned[0].text_content(), "ok");
    }

    #[test]
    fn prune_computer_use_screenshots_keeps_last_three() {
        use lingshu_types::{Content, ContentPart, ImageUrl};

        let make_capture = |id: &str, label: &str| Message {
            role: Role::Tool,
            content: Some(Content::Parts(vec![
                ContentPart::Text {
                    text: format!("capture {label}"),
                },
                ContentPart::ImageUrl {
                    image_url: ImageUrl {
                        url: format!("data:image/png;base64,{label}"),
                        detail: None,
                    },
                },
            ])),
            tool_call_id: Some(id.into()),
            name: Some("computer_use".into()),
            ..Default::default()
        };

        let messages: Vec<Message> = (0..4)
            .map(|i| make_capture(&format!("call_{i}"), &format!("img{i}")))
            .collect();

        let pruned = prune_computer_use_screenshots(&messages, 3);
        assert_eq!(pruned.len(), 4);
        // Oldest capture loses its image part.
        assert!(!pruned[0].text_content().contains("data:image"));
        assert!(pruned[0].text_content().contains("[screenshot pruned"));
        // Newest three retain image parts.
        for msg in &pruned[1..] {
            match &msg.content {
                Some(Content::Parts(parts)) => {
                    assert!(
                        parts
                            .iter()
                            .any(|p| matches!(p, ContentPart::ImageUrl { .. }))
                    );
                }
                other => panic!("expected Parts, got {other:?}"),
            }
        }
    }

    #[test]
    fn prune_computer_use_screenshots_zero_strips_all() {
        use lingshu_types::{Content, ContentPart, ImageUrl};

        let msg = Message {
            role: Role::Tool,
            content: Some(Content::Parts(vec![
                ContentPart::Text {
                    text: "capture".into(),
                },
                ContentPart::ImageUrl {
                    image_url: ImageUrl {
                        url: "data:image/png;base64,abc".into(),
                        detail: None,
                    },
                },
            ])),
            tool_call_id: Some("call_0".into()),
            name: Some("computer_use".into()),
            ..Default::default()
        };

        let pruned = prune_computer_use_screenshots(&[msg], 0);
        assert!(matches!(pruned[0].content, Some(Content::Text(_))));
        assert!(pruned[0].text_content().contains("[screenshot pruned"));
    }

    #[test]
    fn prune_tool_outputs_spills_when_context_provided() {
        use tempfile::TempDir;

        let tmp = TempDir::new().expect("tempdir");
        let seq = crate::tool_result_spill::SpillSequence::new();
        let spill_config = crate::tool_result_spill::SpillConfig {
            enabled: true,
            threshold: 100, // low threshold to trigger spill
            preview_lines: 3,
        };
        let spill_ctx = PruneSpillContext {
            session_id: "test-session",
            cwd: tmp.path(),
            config: &spill_config,
            seq: &seq,
        };

        // 500 chars > spill threshold (100) — should spill, not use placeholder
        let big_result: String = (1..=50)
            .map(|i| format!("line {i}"))
            .collect::<Vec<_>>()
            .join("\n");
        let messages = vec![
            Message::user("search files"),
            Message::tool_result("id1", "file_search", &big_result),
        ];
        let pruned = prune_tool_outputs(&messages, Some(&spill_ctx));
        assert_eq!(pruned.len(), 2);
        // User message unchanged
        assert_eq!(pruned[0].text_content(), "search files");
        // Tool result should be a spill stub, NOT the generic placeholder
        let result_content = pruned[1].text_content();
        assert!(
            result_content.contains("[tool_result_spill]"),
            "expected spill stub, got: {result_content}"
        );
        assert!(result_content.contains("tool: file_search"));
        assert!(result_content.contains("--- BEGIN PREVIEW"));
        assert!(!result_content.contains(PRUNED_TOOL_PLACEHOLDER));
    }

    #[test]
    fn prune_tool_outputs_falls_back_to_placeholder_when_below_spill_threshold() {
        use tempfile::TempDir;

        let tmp = TempDir::new().expect("tempdir");
        let seq = crate::tool_result_spill::SpillSequence::new();
        let spill_config = crate::tool_result_spill::SpillConfig {
            enabled: true,
            threshold: 10_000, // high spill threshold
            preview_lines: 5,
        };
        let spill_ctx = PruneSpillContext {
            session_id: "test-session",
            cwd: tmp.path(),
            config: &spill_config,
            seq: &seq,
        };

        // 500 chars > prune threshold (200) but < spill threshold (10_000)
        let messages = vec![Message::tool_result("id1", "shell_exec", &"x".repeat(500))];
        let pruned = prune_tool_outputs(&messages, Some(&spill_ctx));
        // Should fall back to generic placeholder since below spill threshold
        assert_eq!(pruned[0].text_content(), PRUNED_TOOL_PLACEHOLDER);
    }

    #[test]
    fn prune_tool_outputs_skips_spill_when_disabled() {
        use tempfile::TempDir;

        let tmp = TempDir::new().expect("tempdir");
        let seq = crate::tool_result_spill::SpillSequence::new();
        let spill_config = crate::tool_result_spill::SpillConfig {
            enabled: false, // disabled
            threshold: 100,
            preview_lines: 5,
        };
        let spill_ctx = PruneSpillContext {
            session_id: "test-session",
            cwd: tmp.path(),
            config: &spill_config,
            seq: &seq,
        };

        let messages = vec![Message::tool_result("id1", "shell_exec", &"x".repeat(500))];
        let pruned = prune_tool_outputs(&messages, Some(&spill_ctx));
        // Should use generic placeholder since spill is disabled
        assert_eq!(pruned[0].text_content(), PRUNED_TOOL_PLACEHOLDER);
    }

    #[test]
    fn extract_prior_summary_finds_prefixed_block() {
        let summary_text = "Prior summary content";
        let messages = vec![
            Message::system_summary(format!("{SUMMARY_PREFIX}{summary_text}")),
            Message::user("hello"),
        ];
        let extracted = extract_prior_summary(&messages);
        assert_eq!(extracted.as_deref(), Some(summary_text));
    }

    #[test]
    fn extract_prior_summary_returns_none_without_prefix() {
        let messages = vec![
            Message::system_summary("Regular context summary".to_string()),
            Message::user("hello"),
        ];
        let extracted = extract_prior_summary(&messages);
        assert!(extracted.is_none());
    }

    // ── compress_structural_only tests (FP12 circuit breaker) ─────

    #[test]
    fn structural_only_returns_original_when_too_few_messages() {
        let msgs = make_messages(5);
        let params = CompressionParams {
            context_window: 128_000,
            threshold: 0.50,
            target_ratio: 0.20,
            protect_last_n: 20,
        };
        let result = compress_structural_only(&msgs, &params, None);
        assert_eq!(
            result.len(),
            msgs.len(),
            "should return original when below protect threshold"
        );
    }

    #[test]
    fn structural_only_compresses_large_history() {
        // 200 messages, tiny context window → must compress
        let msgs: Vec<Message> = (0..200)
            .map(|i| {
                if i % 2 == 0 {
                    Message::user(&format!("question {i} {}", "x".repeat(100)))
                } else {
                    Message::assistant(&format!("answer {i} {}", "y".repeat(100)))
                }
            })
            .collect();
        let params = CompressionParams {
            context_window: 500,
            threshold: 0.10,
            target_ratio: 0.20,
            protect_last_n: 5,
        };
        let result = compress_structural_only(&msgs, &params, None);
        assert!(result.len() < msgs.len(), "should produce fewer messages");
        // Should contain a summary message with SUMMARY_PREFIX
        let has_summary = result
            .iter()
            .any(|m| m.text_content().contains(SUMMARY_PREFIX));
        assert!(has_summary, "should contain a structural summary message");
    }

    #[test]
    fn structural_only_preserves_recent_messages() {
        let msgs: Vec<Message> = (0..100)
            .map(|i| {
                if i % 2 == 0 {
                    Message::user(&format!("q{i}"))
                } else {
                    Message::assistant(&format!("a{i}"))
                }
            })
            .collect();
        let params = CompressionParams {
            context_window: 200,
            threshold: 0.05,
            target_ratio: 0.20,
            protect_last_n: 5,
        };
        let result = compress_structural_only(&msgs, &params, None);
        // Last message should be preserved
        let last_original = msgs
            .last()
            .expect("test messages should contain a last item")
            .text_content();
        let last_result = result
            .last()
            .expect("compressed result should preserve the last item")
            .text_content();
        assert_eq!(last_original, last_result, "last message must be preserved");
    }

    // ── FP29 / FP30 — compress_with_llm bool return + role-collision ──

    /// FP29: compress_with_llm falls back to structural when LLM fails and
    /// returns `llm_succeeded = false`.
    #[tokio::test]
    async fn compress_with_llm_returns_false_on_llm_failure() {
        use async_trait::async_trait;
        use edgequake_llm::error::LlmError;
        use edgequake_llm::traits::{
            ChatMessage, CompletionOptions, LLMProvider, LLMResponse, ToolChoice, ToolDefinition,
        };
        // futures::stream::BoxStream not needed

        struct FailingProvider;

        #[async_trait]
        impl LLMProvider for FailingProvider {
            fn name(&self) -> &str {
                "failing"
            }
            fn model(&self) -> &str {
                "test-model"
            }
            fn max_context_length(&self) -> usize {
                128_000
            }

            async fn complete(&self, _: &str) -> edgequake_llm::Result<LLMResponse> {
                Err(LlmError::ApiError("simulated failure".to_string()))
            }
            async fn complete_with_options(
                &self,
                _: &str,
                _: &CompletionOptions,
            ) -> edgequake_llm::Result<LLMResponse> {
                Err(LlmError::ApiError("simulated failure".to_string()))
            }
            async fn chat(
                &self,
                _: &[ChatMessage],
                _: Option<&CompletionOptions>,
            ) -> edgequake_llm::Result<LLMResponse> {
                Err(LlmError::ApiError("simulated failure".to_string()))
            }
            async fn chat_with_tools(
                &self,
                _: &[ChatMessage],
                _: &[ToolDefinition],
                _: Option<ToolChoice>,
                _: Option<&CompletionOptions>,
            ) -> edgequake_llm::Result<LLMResponse> {
                Err(LlmError::ApiError("simulated failure".to_string()))
            }
        }

        // Build enough messages to trigger compression.
        let msgs: Vec<Message> = (0..40)
            .map(|i| {
                if i % 2 == 0 {
                    Message::user(&format!("question {i} {}", "x".repeat(200)))
                } else {
                    Message::assistant(&format!("answer {i} {}", "y".repeat(200)))
                }
            })
            .collect();

        let params = CompressionParams {
            context_window: 1_000,
            threshold: 0.10,
            target_ratio: 0.20,
            protect_last_n: 5,
        };

        let provider: Arc<dyn LLMProvider> = Arc::new(FailingProvider);
        let (compressed, llm_succeeded) = compress_with_llm(&msgs, &params, &provider, None).await;

        // LLM failed → bool must be false.
        assert!(
            !llm_succeeded,
            "expected llm_succeeded=false when provider errors"
        );
        // But structural fallback should have reduced message count.
        assert!(
            compressed.len() < msgs.len(),
            "structural fallback must still compress"
        );
    }

    /// FP29: compress_with_llm returns `true` when LLM succeeds.
    #[tokio::test]
    async fn compress_with_llm_returns_true_on_llm_success() {
        use async_trait::async_trait;
        use edgequake_llm::traits::{
            ChatMessage, CompletionOptions, LLMProvider, LLMResponse, ToolChoice, ToolDefinition,
        };

        struct SuccessProvider;

        #[async_trait]
        impl LLMProvider for SuccessProvider {
            fn name(&self) -> &str {
                "success"
            }
            fn model(&self) -> &str {
                "test-model"
            }
            fn max_context_length(&self) -> usize {
                128_000
            }

            async fn complete(&self, _: &str) -> edgequake_llm::Result<LLMResponse> {
                Ok(LLMResponse::new("summary text", "test-model"))
            }
            async fn complete_with_options(
                &self,
                _: &str,
                _: &CompletionOptions,
            ) -> edgequake_llm::Result<LLMResponse> {
                Ok(LLMResponse::new("summary text", "test-model"))
            }
            async fn chat(
                &self,
                _: &[ChatMessage],
                _: Option<&CompletionOptions>,
            ) -> edgequake_llm::Result<LLMResponse> {
                Ok(LLMResponse::new("summary text", "test-model"))
            }
            async fn chat_with_tools(
                &self,
                _: &[ChatMessage],
                _: &[ToolDefinition],
                _: Option<ToolChoice>,
                _: Option<&CompletionOptions>,
            ) -> edgequake_llm::Result<LLMResponse> {
                Ok(LLMResponse::new("summary text", "test-model"))
            }
        }

        let msgs: Vec<Message> = (0..40)
            .map(|i| {
                if i % 2 == 0 {
                    Message::user(&format!("question {i} {}", "x".repeat(200)))
                } else {
                    Message::assistant(&format!("answer {i} {}", "y".repeat(200)))
                }
            })
            .collect();

        let params = CompressionParams {
            context_window: 1_000,
            threshold: 0.10,
            target_ratio: 0.20,
            protect_last_n: 5,
        };

        let provider: Arc<dyn LLMProvider> = Arc::new(SuccessProvider);
        let (_compressed, llm_succeeded) = compress_with_llm(&msgs, &params, &provider, None).await;

        assert!(
            llm_succeeded,
            "expected llm_succeeded=true when provider succeeds"
        );
    }

    /// FP30: Role-collision guard — when the compressible head ends with a
    /// System message, the summary must use User role to avoid adjacent system+system.
    #[tokio::test]
    async fn compress_with_llm_fp30_avoids_adjacent_system_messages() {
        use async_trait::async_trait;
        use edgequake_llm::traits::{
            ChatMessage, CompletionOptions, LLMProvider, LLMResponse, ToolChoice, ToolDefinition,
        };

        struct SuccessProvider;

        #[async_trait]
        impl LLMProvider for SuccessProvider {
            fn name(&self) -> &str {
                "success"
            }
            fn model(&self) -> &str {
                "test-model"
            }
            fn max_context_length(&self) -> usize {
                128_000
            }

            async fn complete(&self, _: &str) -> edgequake_llm::Result<LLMResponse> {
                Ok(LLMResponse::new("summary text", "test-model"))
            }
            async fn complete_with_options(
                &self,
                _: &str,
                _: &CompletionOptions,
            ) -> edgequake_llm::Result<LLMResponse> {
                Ok(LLMResponse::new("summary text", "test-model"))
            }
            async fn chat(
                &self,
                _: &[ChatMessage],
                _: Option<&CompletionOptions>,
            ) -> edgequake_llm::Result<LLMResponse> {
                Ok(LLMResponse::new("summary text", "test-model"))
            }
            async fn chat_with_tools(
                &self,
                _: &[ChatMessage],
                _: &[ToolDefinition],
                _: Option<ToolChoice>,
                _: Option<&CompletionOptions>,
            ) -> edgequake_llm::Result<LLMResponse> {
                Ok(LLMResponse::new("summary text", "test-model"))
            }
        }

        // Build a message list where position PROTECT_FIRST_N is a System message.
        // PROTECT_FIRST_N = 4 (head is always kept)
        // So the first 4 messages form the protected head.
        // We make message[3] (0-indexed) a System message → head_end boundary
        // will land just after it → last_head_role == System.
        let mut msgs = vec![
            Message::user("system context"), // 0
            Message::assistant("ok"),        // 1
            Message::user("follow up"),      // 2
            Message::system("extra system"), // 3  ← PROTECT_FIRST_N-1 (last in head)
        ];
        // Add enough body messages to trigger compression.
        for i in 0..30 {
            if i % 2 == 0 {
                msgs.push(Message::user(&format!("body q{i} {}", "x".repeat(300))));
            } else {
                msgs.push(Message::assistant(&format!(
                    "body a{i} {}",
                    "y".repeat(300)
                )));
            }
        }

        let params = CompressionParams {
            context_window: 2_000,
            threshold: 0.10,
            target_ratio: 0.20,
            protect_last_n: 3,
        };

        let provider: Arc<dyn LLMProvider> = Arc::new(SuccessProvider);
        let (compressed, _) = compress_with_llm(&msgs, &params, &provider, None).await;

        // Find the summary message (contains SUMMARY_PREFIX).
        let summary_msg = compressed
            .iter()
            .find(|m| m.text_content().contains(SUMMARY_PREFIX));

        if let Some(_s) = summary_msg {
            // When head ends with System, the summary must NOT be System.
            // Check adjacent pairs: no two consecutive System messages.
            for window in compressed.windows(2) {
                assert!(
                    !(window[0].role == lingshu_types::Role::System
                        && window[1].role == lingshu_types::Role::System),
                    "FP30: adjacent system+system messages found in compressed output"
                );
            }
        }
        // Whether or not compression triggered, the output must be non-empty.
        assert!(!compressed.is_empty());
    }
}