sqz_engine/

pipeline.rs

use sha2::{Digest, Sha256};

use crate::ansi_strip::AnsiStripper;
use crate::dict_compressor::DictCompressor;
use crate::entropy_truncator::EntropyTruncator;
use crate::error::{Result, SqzError};
use crate::preset::Preset;
use crate::prompt_cache::PromptCacheDetector;
use crate::stages::{
    CollapseArraysStage, CondenseStage, CustomTransformsStage, FlattenStage, GitDiffFoldStage,
    KeepFieldsStage, StripFieldsStage, StripNullsStage, TruncateStringsStage,
};
use crate::token_counter::TokenCounter;
use crate::token_pruner::TokenPruner;
use crate::toon::ToonEncoder;
use crate::types::{CompressedContent, Content, ContentType, ModelFamily, StageConfig};

/// Minimal session context passed to the pipeline. Just carries the session
/// ID so stages can log or cache per-session.
pub struct SessionContext {
    pub session_id: String,
}

/// The compression pipeline orchestrator.
///
/// Runs content through a sequence of configurable stages (condense, strip
/// nulls, collapse arrays, etc.), then applies post-processing: RLE
/// compression, sliding-window dedup, entropy truncation, token pruning,
/// dictionary compression, and TOON encoding for JSON.
///
/// Stages are sorted by priority — lower numbers run first. You can insert
/// custom stages with [`insert_stage`](CompressionPipeline::insert_stage).
pub struct CompressionPipeline {
    stages: Vec<Box<dyn crate::stages::CompressionStage>>,
    toon_encoder: ToonEncoder,
    token_counter: TokenCounter,
    token_counts: std::cell::RefCell<CountCache>,
    token_pruner: TokenPruner,
    dict_compressor: std::cell::RefCell<DictCompressor>,
    entropy_truncator: EntropyTruncator,
    #[allow(dead_code)]
    prompt_cache_detector: PromptCacheDetector,
}

impl CompressionPipeline {
    /// Construct the pipeline from a preset, creating all built-in stages
    /// sorted by priority, plus the token pruner, dictionary compressor,
    /// and entropy truncator.
    pub fn new(_preset: &Preset) -> Self {
        let mut stages: Vec<Box<dyn crate::stages::CompressionStage>> = vec![
            Box::new(AnsiStripper),
            Box::new(KeepFieldsStage),
            Box::new(StripFieldsStage),
            Box::new(CondenseStage),
            Box::new(GitDiffFoldStage),
            Box::new(StripNullsStage),
            Box::new(FlattenStage),
            Box::new(TruncateStringsStage),
            Box::new(CollapseArraysStage),
            Box::new(CustomTransformsStage),
        ];
        stages.sort_by_key(|s| s.priority());

        Self {
            stages,
            toon_encoder: ToonEncoder,
            token_counter: TokenCounter::new(),
            token_counts: std::cell::RefCell::new(CountCache::default()),
            token_pruner: TokenPruner::new(),
            dict_compressor: std::cell::RefCell::new(DictCompressor::new()),
            entropy_truncator: EntropyTruncator::new(),
            prompt_cache_detector: PromptCacheDetector,
        }
    }

    /// Run content through all enabled stages, then apply:
    /// 1. Dictionary compression + TOON encoding for JSON
    /// 2. Token pruning for prose/plain text
    /// 3. Entropy-weighted truncation for long content
    pub fn compress(
        &self,
        input: &str,
        ctx: &SessionContext,
        preset: &Preset,
    ) -> Result<CompressedContent> {
        let mut result = self.compress_inner(input, ctx, preset, false)?;
        self.count_output(&mut result, preset);
        Ok(result)
    }

    pub(crate) fn compress_uncounted(
        &self,
        input: &str,
        ctx: &SessionContext,
        preset: &Preset,
    ) -> Result<CompressedContent> {
        self.compress_inner(input, ctx, preset, false)
    }

    pub(crate) fn count_output(&self, result: &mut CompressedContent, preset: &Preset) {
        result.tokens_compressed = self.count_tokens(&result.data, preset);
        result.compression_ratio = if result.tokens_original == 0 {
            1.0
        } else {
            result.tokens_compressed as f64 / result.tokens_original as f64
        };
    }

    /// Byte-faithful compression for file reads (`sqz_read_file`, `sqz_grep`,
    /// `sqz_list_dir`). Returns the input verbatim apart from ANSI-escape
    /// stripping: every content-altering stage is skipped — both the droppers
    /// (entropy truncation, token pruning, JSON field projection, and the lossy
    /// configured stages: condense, strip_fields, strip_nulls, flatten,
    /// truncate_strings, collapse_arrays, git_diff_fold) and the
    /// reversible-but-reformatting ones (RLE, sliding-window dedup, TOON/dict),
    /// which would shift bytes or line numbers. No line, field, or segment is
    /// ever dropped or moved. Token savings on this path come from the dedup
    /// `§ref` cache on repeat reads, not from re-encoding.
    /// See <https://github.com/ojuschugh1/sqz/issues/32>
    pub fn compress_lossless(
        &self,
        input: &str,
        ctx: &SessionContext,
        preset: &Preset,
    ) -> Result<CompressedContent> {
        let mut result = self.compress_inner(input, ctx, preset, true)?;
        self.count_output(&mut result, preset);
        Ok(result)
    }

    /// Token count with the same tokenizer `compress` used for
    /// `tokens_original`, so callers that post-process the output can report
    /// comparable numbers.
    pub fn count_tokens(&self, text: &str, preset: &Preset) -> u32 {
        let family = model_family_from_preset(preset);
        if text.len() < COUNT_CACHE_MIN_BYTES || matches!(family, ModelFamily::Local(_)) {
            return self.token_counter.count(text, &family);
        }
        let key: CountKey = (
            std::mem::discriminant(&family),
            text.len(),
            Sha256::digest(text.as_bytes()).into(),
        );
        if let Some(n) = self.token_counts.borrow().get(&key) {
            return n;
        }
        let n = self.token_counter.count(text, &family);
        self.token_counts.borrow_mut().insert(key, n);
        n
    }

    fn compress_inner(
        &self,
        input: &str,
        _ctx: &SessionContext,
        preset: &Preset,
        lossless: bool,
    ) -> Result<CompressedContent> {
        let tokens_original = self.count_tokens(input, preset);

        let mut content = Content {
            raw: input.to_owned(),
            content_type: ContentType::PlainText,
            metadata: crate::types::ContentMetadata {
                source: None,
                path: None,
                language: None,
            },
            tokens_original,
        };

        let mut stages_applied: Vec<String> = Vec::new();
        let mut condensed: Option<(String, String)> = None;

        for stage in &self.stages {
            // On the lossless (file-read) path, run only the reversible stages.
            if lossless && !stage_is_lossless(stage.name()) {
                continue;
            }
            let config = stage_config_from_preset(stage.name(), preset);
            if config.enabled {
                let before = (stage.name() == "condense").then(|| content.raw.clone());
                stage.process(&mut content, &config)?;
                if let Some(before) = before.filter(|before| *before != content.raw) {
                    condensed = Some((before, content.raw.clone()));
                }
                stages_applied.push(stage.name().to_owned());
            }
        }

        // Post-stage processing: apply advanced compression techniques

        let is_json = ToonEncoder::is_json(&content.raw);
        // Source code gets no line-dropping stages: an agent that cats a
        // file is about to edit it, and imports or a class header score as
        // "low information" to the truncator. Dedup still covers re-reads.
        let is_code = !is_json && looks_like_code(&content.raw);
        if is_code {
            stages_applied.push("source_code".to_owned());
        }

        // JSON projection: strip internal/debug fields, empty collections,
        // deep nesting, and redundant timestamps before other JSON processing
        if !lossless && is_json && content.raw.len() > 100 {
            let proj_config = crate::json_projection::ProjectionConfig::default();
            if let Ok(proj_result) = crate::json_projection::project_json(&content.raw, &proj_config) {
                if proj_result.fields_removed > 0 {
                    content.raw = proj_result.data;
                    stages_applied.push("json_projection".to_owned());
                }
            }
        }

        // RLE: collapse repeated patterns in non-JSON content (generalizes condense).
        // Skipped on the lossless path — even reversible `[×N]` markers make a
        // file read non-verbatim. Only apply when content is long enough.
        if !lossless && !is_json && !is_code && content.raw.len() > 200 {
            let original = condensed
                .as_ref()
                .filter(|(_, after)| *after == content.raw)
                .map(|(before, _)| before.as_str());
            if let Ok(rle_result) =
                crate::rle_compressor::rle_compress_with_original(&content.raw, original, 3)
            {
                if rle_result.runs_collapsed > 0 {
                    // Safety check: verify critical markers are preserved
                    let has_error = content.raw.contains("ERROR") || content.raw.contains("error:");
                    let rle_has_error = rle_result.text.contains("ERROR") || rle_result.text.contains("error:");
                    if !has_error || rle_has_error {
                        content.raw = rle_result.text;
                        stages_applied.push("rle".to_owned());
                    }
                }
            }
        }

        // Log-template compaction: collapse consecutive lines identical
        // except for timestamps/durations/percentages — what exact RLE
        // misses in real logs. Identifier-bearing lines never collapse.
        // Skipped on the lossless path.
        if !lossless && !is_json && !is_code && content.raw.len() > 200 {
            if let Ok(lt_result) = crate::rle_compressor::log_template_collapse(&content.raw, 3) {
                if lt_result.runs_collapsed > 0 {
                    let has_error = content.raw.contains("ERROR") || content.raw.contains("error:");
                    let lt_has_error = lt_result.text.contains("ERROR") || lt_result.text.contains("error:");
                    if !has_error || lt_has_error {
                        content.raw = lt_result.text;
                        stages_applied.push("log_template".to_owned());
                    }
                }
            }
        }

        // Sliding window dedup: catch repeated substrings across non-adjacent lines.
        // Skipped on the lossless path — its `[→Ln]` back-references have no
        // agent-facing expand path and make a file read non-faithful (#32).
        if !lossless && !is_json && !is_code && content.raw.len() > 300 {
            if let Ok(sw_result) = crate::rle_compressor::sliding_window_dedup(&content.raw, 4) {
                if sw_result.dedup_count > 0 {
                    // Safety check: verify critical markers are preserved
                    let has_error = content.raw.contains("ERROR") || content.raw.contains("error:");
                    let sw_has_error = sw_result.text.contains("ERROR") || sw_result.text.contains("error:");
                    if !has_error || sw_has_error {
                        content.raw = sw_result.text;
                        stages_applied.push("sliding_window_dedup".to_owned());
                    }
                }
            }
        }

        // Aligned-column table compaction: collapse alignment padding in
        // tabular CLI output (docker ps, kubectl get, ps aux, ...) to a
        // two-space separator. Cells survive byte-exact; the detector
        // refuses anything with leading indentation, so code, YAML,
        // diffs, and pretty JSON are never touched. Runs before entropy
        // truncation so the truncator scores the compact form.
        if !lossless && !is_json && !is_code && content.raw.len() > 200 {
            if let Some(table_result) =
                crate::table_compactor::compact_aligned_table(&content.raw)
            {
                content.raw = table_result.text;
                stages_applied.push("table_compact".to_owned());
            }
        }

        // Entropy-weighted truncation for long non-JSON content.
        // LOSSY — skipped on the lossless file-read path (issue #32).
        // Benign-aware threshold: output with no error markers (clean
        // builds, passing test runs) gets 2x headroom before truncation
        // kicks in, since dropping detail from a success costs more than
        // it saves.
        let entropy_threshold = if looks_benign(&content.raw) { 1000 } else { 500 };
        let record_list = content.raw.len() <= RECORD_LIST_MAX_CHARS
            && (stages_applied.iter().any(|s| s == "table_compact")
                || looks_like_record_list(&content.raw));
        if !lossless && !is_json && !is_code && !record_list && content.raw.len() > entropy_threshold {
            if let Ok(trunc_result) = self.entropy_truncator.truncate_string(&content.raw) {
                if trunc_result.segments_dropped > 0 {
                    content.raw = trunc_result.text;
                    stages_applied.push("entropy_truncate".to_owned());
                }
            }
        }

        // Token pruning for prose content (non-JSON, non-code).
        // LOSSY — skipped on the lossless file-read path (issue #32).
        if !lossless && !is_json && content.raw.len() > 100 && looks_like_prose(&content.raw) {
            if let Ok(prune_result) = self.token_pruner.prune(&content.raw) {
                if prune_result.tokens_removed > 0 {
                    content.raw = prune_result.text;
                    stages_applied.push("token_prune".to_owned());
                }
            }
        }

        // Apply dictionary compression + TOON encoding for JSON.
        // Skipped on the lossless path so a JSON file read returns its bytes
        // verbatim — TOON re-encoding is data-lossless but reformats and
        // shifts line numbers, which breaks "read this file faithfully".
        let data = if !lossless && ToonEncoder::is_json(&content.raw) {
            let json: serde_json::Value = serde_json::from_str(&content.raw)
                .map_err(|e| SqzError::Other(format!("pipeline: JSON parse error: {e}")))?;
            let encoded = self.toon_encoder.encode(&json)?;

            // Short codes for repeated keys. Tables already print each key
            // once, so the dictionary is kept only when it beats plain TOON
            // with its header counted.
            self.dict_compressor.borrow_mut().observe(&content.raw);
            let dict_result = self.dict_compressor.borrow().compress(&content.raw)?;
            let with_dict = if dict_result.substitutions > 0 {
                dict_result.data.find("\n{").and_then(|pos| {
                    let (header, body) = dict_result.data.split_at(pos + 1);
                    let json = serde_json::from_str::<serde_json::Value>(body).ok()?;
                    let body = self.toon_encoder.encode(&json).ok()?;
                    Some(format!("{header}{body}"))
                })
            } else {
                None
            };
            match with_dict {
                Some(dict) if dict.len() < encoded.len() => {
                    stages_applied.push("dict_compress".to_owned());
                    stages_applied.push("toon_encode".to_owned());
                    dict
                }
                _ => {
                    stages_applied.push("toon_encode".to_owned());
                    encoded
                }
            }
        } else {
            content.raw
        };

        Ok(CompressedContent {
            data,
            tokens_compressed: 0,
            tokens_original,
            stages_applied,
            compression_ratio: 1.0,
            provenance: crate::types::Provenance::default(),
            verify: None,
        })
    }

    /// Insert a plugin stage and re-sort by priority.
    pub fn insert_stage(&mut self, stage: Box<dyn crate::stages::CompressionStage>) {
        self.stages.push(stage);
        self.stages.sort_by_key(|s| s.priority());
    }

    /// Rebuild stage list from a new preset (hot-reload support).
    /// Built-in stages are recreated; plugin stages are dropped and must be
    /// re-inserted by the caller.
    pub fn reload_preset(&mut self, _preset: &Preset) -> Result<()> {
        let mut stages: Vec<Box<dyn crate::stages::CompressionStage>> = vec![
            Box::new(AnsiStripper),
            Box::new(KeepFieldsStage),
            Box::new(StripFieldsStage),
            Box::new(CondenseStage),
            Box::new(GitDiffFoldStage),
            Box::new(StripNullsStage),
            Box::new(FlattenStage),
            Box::new(TruncateStringsStage),
            Box::new(CollapseArraysStage),
            Box::new(CustomTransformsStage),
        ];
        stages.sort_by_key(|s| s.priority());
        self.stages = stages;
        // Reset session-level state in dict compressor
        self.dict_compressor.borrow_mut().reset();
        Ok(())
    }
}

const COUNT_CACHE_MIN_BYTES: usize = 4 * 1024;
const COUNT_CACHE_MAX_ENTRIES: usize = 64;

type CountKey = (std::mem::Discriminant<ModelFamily>, usize, [u8; 32]);

#[derive(Default)]
struct CountCache {
    entries: std::collections::VecDeque<(CountKey, u32)>,
}

impl CountCache {
    fn get(&self, key: &CountKey) -> Option<u32> {
        self.entries.iter().find(|(k, _)| k == key).map(|(_, n)| *n)
    }

    fn insert(&mut self, key: CountKey, tokens: u32) {
        if self.entries.len() == COUNT_CACHE_MAX_ENTRIES {
            self.entries.pop_front();
        }
        self.entries.push_back((key, tokens));
    }
}

/// Derive the `ModelFamily` from the preset's model configuration.
fn model_family_from_preset(preset: &Preset) -> ModelFamily {
    match preset.model.family.to_lowercase().as_str() {
        "anthropic" | "claude" => ModelFamily::AnthropicClaude,
        "openai" | "gpt" => ModelFamily::OpenAiGpt,
        "google" | "gemini" => ModelFamily::GoogleGemini,
        other => ModelFamily::Local(other.to_string()),
    }
}

/// Build a `StageConfig` for a named stage from the preset's compression config.
fn stage_config_from_preset(name: &str, preset: &Preset) -> StageConfig {
    let c = &preset.compression;
    match name {
        "ansi_strip" => StageConfig {
            enabled: true,
            options: serde_json::Value::Object(Default::default()),
        },
        "keep_fields" => {
            if let Some(cfg) = &c.keep_fields {
                StageConfig {
                    enabled: cfg.enabled,
                    options: serde_json::json!({ "fields": cfg.fields }),
                }
            } else {
                StageConfig::default()
            }
        }
        "strip_fields" => {
            if let Some(cfg) = &c.strip_fields {
                StageConfig {
                    enabled: cfg.enabled,
                    options: serde_json::json!({ "fields": cfg.fields }),
                }
            } else {
                StageConfig::default()
            }
        }
        "condense" => {
            if let Some(cfg) = &c.condense {
                StageConfig {
                    enabled: cfg.enabled,
                    options: serde_json::json!({
                        "max_repeated_lines": cfg.max_repeated_lines
                    }),
                }
            } else {
                StageConfig::default()
            }
        }
        "git_diff_fold" => {
            if let Some(cfg) = &c.git_diff_fold {
                StageConfig {
                    enabled: cfg.enabled,
                    options: serde_json::json!({
                        "max_context_lines": cfg.max_context_lines
                    }),
                }
            } else {
                // Default: enabled with 2 context lines
                StageConfig {
                    enabled: true,
                    options: serde_json::json!({ "max_context_lines": 2 }),
                }
            }
        }
        "strip_nulls" => {
            if let Some(cfg) = &c.strip_nulls {
                StageConfig {
                    enabled: cfg.enabled,
                    options: serde_json::Value::Object(Default::default()),
                }
            } else {
                StageConfig::default()
            }
        }
        "flatten" => {
            if let Some(cfg) = &c.flatten {
                StageConfig {
                    enabled: cfg.enabled,
                    options: serde_json::json!({ "max_depth": cfg.max_depth }),
                }
            } else {
                StageConfig::default()
            }
        }
        "truncate_strings" => {
            if let Some(cfg) = &c.truncate_strings {
                StageConfig {
                    enabled: cfg.enabled,
                    options: serde_json::json!({ "max_length": cfg.max_length }),
                }
            } else {
                StageConfig::default()
            }
        }
        "collapse_arrays" => {
            if let Some(cfg) = &c.collapse_arrays {
                StageConfig {
                    enabled: cfg.enabled,
                    options: serde_json::json!({
                        "max_items": cfg.max_items,
                        "max_chars": cfg.max_chars,
                        "summary_template": cfg.summary_template
                    }),
                }
            } else {
                StageConfig::default()
            }
        }
        "custom_transforms" => {
            if let Some(cfg) = &c.custom_transforms {
                StageConfig {
                    enabled: cfg.enabled,
                    options: serde_json::Value::Object(Default::default()),
                }
            } else {
                StageConfig::default()
            }
        }
        _ => StageConfig::default(),
    }
}

/// Whether a configured stage is safe to run on the lossless file-read path.
/// Only stages that never drop or irreversibly alter content qualify; every
/// other stage (condense, strip_fields, strip_nulls, flatten, truncate_strings,
/// collapse_arrays, git_diff_fold, keep_fields) is skipped when `lossless`.
fn stage_is_lossless(name: &str) -> bool {
    matches!(name, "ansi_strip" | "custom_transforms")
}

/// Record lists at or under this size are never entropy-truncated.
const RECORD_LIST_MAX_CHARS: usize = 16_000;

/// One fact per line: grep matches, file lists, `ls -l`, commit lists.
fn looks_like_record_list(text: &str) -> bool {
    let lines: Vec<&str> = text.lines().map(str::trim_end).filter(|l| !l.is_empty()).collect();
    if lines.len() < 8 {
        return false;
    }
    let unique: std::collections::HashSet<&str> = lines.iter().copied().collect();
    if unique.len() * 10 < lines.len() * 9 {
        return false;
    }
    let records = lines.iter().filter(|l| is_record_line(l)).count();
    records * 10 >= lines.len() * 8
}

fn is_record_line(line: &str) -> bool {
    let t = line.trim_start();
    let pathish = |s: &str| {
        (s.contains('/') || s.contains('.')) && s.bytes().any(|b| b.is_ascii_alphanumeric())
    };
    // `path:12:text` (grep -n, compiler diagnostics)
    if let Some((path, rest)) = t.split_once(':') {
        let digits = rest.bytes().take_while(|b| b.is_ascii_digit()).count();
        if digits > 0 && rest.as_bytes().get(digits) == Some(&b':')
            && !path.contains(char::is_whitespace) && pathish(path)
        {
            return true;
        }
    }
    // A bare path (find, ls -1, git ls-files, diff --name-only)
    if !t.contains(char::is_whitespace) && pathish(t) {
        return true;
    }
    // `ls -l`: drwxr-xr-x ...
    let b = t.as_bytes();
    if b.len() > 10
        && b"-dlcbps".contains(&b[0])
        && b[1..10].iter().all(|c| b"rwxsStT-".contains(c))
    {
        return true;
    }
    // `git log --oneline`: abc1234 subject
    let sha = t.bytes().take_while(|b| b.is_ascii_hexdigit()).count();
    sha >= 7 && sha <= 40 && t.as_bytes().get(sha).map_or(true, |b| *b == b' ')
}

/// True when the output carries no error markers — a clean build, passing
/// tests, an uneventful log. Benign output earns a higher truncation
/// threshold.
fn looks_benign(text: &str) -> bool {
    let lower = text.to_lowercase();
    const MARKERS: &[&str] = &[
        "error", "panic", "traceback", "failed", "failure", "exception", "fatal",
    ];
    !MARKERS.iter().any(|m| lower.contains(m))
}

/// Heuristic: does this text look like prose (documentation, error messages,
/// README content) rather than code or structured data?
fn looks_like_prose(text: &str) -> bool {
    let lines: Vec<&str> = text.lines().take(20).collect();
    if lines.is_empty() {
        return false;
    }

    let mut prose_lines = 0;
    let mut code_lines = 0;

    for line in &lines {
        let trimmed = line.trim();
        if trimmed.is_empty() {
            continue;
        }
        // Code indicators
        if trimmed.ends_with('{')
            || trimmed.ends_with('}')
            || trimmed.ends_with(';')
            || trimmed.starts_with("fn ")
            || trimmed.starts_with("pub ")
            || trimmed.starts_with("let ")
            || trimmed.starts_with("const ")
            || trimmed.starts_with("import ")
            || trimmed.starts_with("def ")
            || trimmed.starts_with("class ")
            || trimmed.contains("::")
            || trimmed.contains("->")
            || trimmed.contains("=>")
        {
            code_lines += 1;
        } else {
            prose_lines += 1;
        }
    }

    // Consider it prose if more than half the lines look like prose
    prose_lines > code_lines
}

/// Heuristic: is this a source file rather than command output? Needs a
/// couple of declaration lines and a clear majority of code-shaped lines.
/// Timestamped lines are ignored so logs never qualify. False positives
/// only cost compression; false negatives cost imports and class headers.
fn looks_like_code(text: &str) -> bool {
    const DECL: &[&str] = &[
        "def ", "class ", "fn ", "pub ", "func ", "function ", "async ", "import ", "from ",
        "use ", "package ", "struct ", "impl ", "enum ", "trait ", "interface ", "type ",
        "export ", "const ", "let ", "var ", "static ", "public ", "private ", "protected ",
        "return ", "#include", "@",
    ];
    let lines: Vec<&str> = text.lines().filter(|l| !l.trim().is_empty()).take(80).collect();
    if lines.len() < 3 {
        return false;
    }
    let mut decl = 0usize;
    let mut codeish = 0usize;
    for line in &lines {
        let t = line.trim_start();
        if t.starts_with(|c: char| c.is_ascii_digit()) {
            continue;
        }
        let is_decl = DECL.iter().any(|k| t.starts_with(k));
        if is_decl {
            decl += 1;
        }
        let ends_like_code = t.trim_end().ends_with(['{', '}', ';', ':', ')', ',']);
        if is_decl
            || ends_like_code
            || t.contains(" = ")
            || t.starts_with("//")
            || t.starts_with('#')
            || t.starts_with("/*")
            || t.starts_with('*')
            || line.starts_with("    ")
            || line.starts_with('\t')
        {
            codeish += 1;
        }
    }
    decl >= 2 && codeish * 10 >= lines.len() * 7
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use crate::preset::{
        BudgetConfig, CollapseArraysConfig, CompressionConfig, CondenseConfig,
        CustomTransformsConfig, ModelConfig, PresetMeta,
        StripNullsConfig, ToolSelectionConfig, TruncateStringsConfig,
        TerseModeConfig,
    };

    fn default_preset() -> Preset {
        Preset {
            preset: PresetMeta {
                name: "test".into(),
                version: "1.0".into(),
                description: String::new(),
            },
            compression: CompressionConfig {
                stages: vec![],
                keep_fields: None,
                strip_fields: None,
                condense: Some(CondenseConfig {
                    enabled: true,
                    max_repeated_lines: 3,
                }),
                git_diff_fold: None,
                strip_nulls: Some(StripNullsConfig { enabled: true }),
                flatten: None,
                truncate_strings: Some(TruncateStringsConfig {
                    enabled: true,
                    max_length: 500,
                }),
                collapse_arrays: Some(CollapseArraysConfig {
                    enabled: true,
                    max_items: 5,
                    max_chars: 32_000,
                    summary_template: "... and {remaining} more items".into(),
                }),
                custom_transforms: Some(CustomTransformsConfig { enabled: true }),
            },
            tool_selection: ToolSelectionConfig {
                max_tools: 5,
                similarity_threshold: 0.7,
                default_tools: vec![],
            },
            budget: BudgetConfig {
                warning_threshold: 0.70,
                ceiling_threshold: 0.85,
                default_window_size: 200_000,
                agents: Default::default(),
            },
            terse_mode: TerseModeConfig {
                enabled: false,
                level: crate::preset::TerseLevel::Moderate,
            },
            model: ModelConfig {
                family: "anthropic".into(),
                primary: "claude-sonnet-4-20250514".into(),
                local: String::new(),
                complexity_threshold: 0.4,
                pricing: None,
            },
        }
    }

    fn ctx() -> SessionContext {
        SessionContext {
            session_id: "test-session".into(),
        }
    }

    /// Regression for https://github.com/ojuschugh1/sqz/issues/32 — the
    /// file-read path (sqz_read_file / sqz_grep / sqz_list_dir) must never
    /// silently drop content. `compress_lossless` keeps every segment; the
    /// default `compress` truncates below-median-entropy segments.
    #[test]
    fn source_code_skips_line_dropping_stages() {
        let pipeline = CompressionPipeline::new(&default_preset());
        let preset = default_preset();
        let mut code = String::from(
            "from dataclasses import dataclass\n\nfrom .errors import RefundError\nfrom .gateway import Gateway\n\n\n@dataclass\nclass RefundService:\n    gateway: Gateway\n\n",
        );
        for i in 0..12 {
            code.push_str(&format!(
                "    def method_{i}(self, charge: Charge, amount_cents: int) -> Refund:\n        if amount_cents <= 0:\n            raise RefundError(\"invalid refund amount\")\n        refund = self.gateway.refund(charge.id, amount_cents)\n        self.ledger.record(refund)\n        return refund\n\n"
            ));
        }
        assert!(code.len() > 1000);
        assert!(looks_like_code(&code));
        let result = pipeline.compress(&code, &ctx(), &preset).unwrap();
        for stage in ["entropy_truncate", "sliding_window_dedup", "rle", "log_template", "table_compact"] {
            assert!(!result.stages_applied.iter().any(|s| s == stage), "{stage} ran on source code: {:?}", result.stages_applied);
        }
        assert!(result.stages_applied.iter().any(|s| s == "source_code"), "{:?}", result.stages_applied);
        assert!(result.data.contains("from .errors import RefundError"));
        assert!(result.data.contains("class RefundService:"));

        let mut log = String::new();
        for i in 0..60 {
            log.push_str(&format!("2026-09-14T10:22:{:02}Z INFO  GET /healthz 200 1ms\n", i % 60));
        }
        log.push_str("2026-09-14T10:23:41Z ERROR refund failed order=ord_8812\n");
        assert!(!looks_like_code(&log));
        let result = pipeline.compress(&log, &ctx(), &preset).unwrap();
        assert!(result.data.len() < log.len() / 2, "log should still collapse: {}", result.data);
        assert!(result.data.contains("ERROR refund failed order=ord_8812"));

        let status = "On branch main\n\nChanges not staged for commit:\n\tmodified:   README.md\n\tmodified:   src/lib.rs\n\nUntracked files:\n\tnotes.txt\n";
        assert!(!looks_like_code(status));
    }

    #[test]
    fn table_compact_stage_fires_on_tabular_output_only() {
        let pipeline = CompressionPipeline::new(&default_preset());

        let table = "\
NAME                          READY   STATUS    RESTARTS   AGE
api-6d4cf56db6-2m5xk          1/1     Running   0          4d20h
api-6d4cf56db6-9lq2v          1/1     Running   0          4d20h
worker-7c9f8b5d44-abcde       1/1     Running   2          12h
worker-7c9f8b5d44-fghij       0/1     Pending   0          3m
ingress-nginx-controller-x    1/1     Running   0          30d
";
        let preset = default_preset();
        let result = pipeline.compress(table, &ctx(), &preset).unwrap();
        assert!(
            result.stages_applied.iter().any(|s| s == "table_compact"),
            "stages: {:?}",
            result.stages_applied
        );
        assert!(result.data.contains("api-6d4cf56db6-2m5xk"), "{}", result.data);

        let code = "\
fn main() {
    let alignment_one   = 1;
    let alignment_two   = 2;
    let alignment_three = 3;
    let alignment_four  = 4;
    println!(\"{}\", alignment_one + alignment_two);
}
";
        let result = pipeline.compress(code, &ctx(), &preset).unwrap();
        assert!(
            !result.stages_applied.iter().any(|s| s == "table_compact"),
            "code must never hit the table compactor: {:?}",
            result.stages_applied
        );

        // The lossless path must never table-compact either.
        let result = pipeline.compress_lossless(table, &ctx(), &preset).unwrap();
        assert!(
            !result.stages_applied.iter().any(|s| s == "table_compact"),
            "lossless path must stay byte-faithful: {:?}",
            result.stages_applied
        );
    }

    #[test]
    fn compress_lossless_never_drops_segments() {
        let preset = default_preset();
        let pipeline = CompressionPipeline::new(&preset);

        // 10 distinct, blank-line-separated segments with bimodal entropy,
        // well over the 500-byte entropy-truncation threshold.
        let mut segments = Vec::new();
        // 24 segments so the fixture clears the benign-aware 1000-byte
        // threshold (this content has no error markers).
        for i in 0..24 {
            if i % 2 == 0 {
                segments.push(format!("SEG{i} llllllllllllllllllllllllllllllllllllllll"));
            } else {
                segments.push(format!(
                    "SEG{i} the quick brown fox jumps over {i} lazy dogs by silver rivers"
                ));
            }
        }
        let input = segments.join("\n\n");
        assert!(input.len() > 1000, "input must exceed the benign truncation threshold");

        // The default path is lossy — it drops below-median-entropy segments.
        let lossy = pipeline.compress(&input, &ctx(), &preset).unwrap();
        assert!(
            lossy.data.contains("low-information segments omitted"),
            "expected the default path to truncate; got: {}",
            lossy.data
        );

        // The lossless path must preserve EVERY segment, byte-for-byte.
        let lossless = pipeline.compress_lossless(&input, &ctx(), &preset).unwrap();
        assert!(
            !lossless.data.contains("omitted"),
            "lossless path must not truncate; got: {}",
            lossless.data
        );
        assert_eq!(
            lossless.data, input,
            "lossless path must return content unchanged"
        );
        assert!(!lossless.stages_applied.contains(&"entropy_truncate".to_owned()));
        assert!(!lossless.stages_applied.contains(&"token_prune".to_owned()));
    }

    #[test]
    fn repeated_lines_report_the_real_count() {
        let preset = default_preset();
        let pipeline = CompressionPipeline::new(&preset);
        let head = "2024-01-01 09:59:59 [INFO] Server started on 0.0.0.0:8080\n";
        let line = "2024-01-01 10:00:00 [INFO] Connected to database";
        let tail = "2024-01-01 10:00:11 [ERROR] Connection timeout after 30s\n";
        for n in [2, 3, 4, 10, 100] {
            let input = format!("{head}{}{tail}", format!("{line}\n").repeat(n));
            assert!(input.len() > 200, "n={n}");
            let result = pipeline.compress(&input, &ctx(), &preset).unwrap();
            let expected = if n < 3 { input.clone() } else { format!("{head}{line} [×{n}]\n{tail}") };
            assert_eq!(result.data, expected, "n={n}");
        }
    }

    #[test]
    fn folded_and_collapsed_counts_are_exact() {
        let preset = default_preset();
        let pipeline = CompressionPipeline::new(&preset);
        let context: String = (0..6).map(|i| format!(" context_{i:02}()\n")).collect();
        let diff = format!(
            "diff --git a/a.py b/a.py\nindex 1111111..2222222 100644\n--- a/a.py\n+++ b/a.py\n@@ -1,30 +1,30 @@\n-old_call()\n+new_call()\n{context}{}{context}-old_tail()\n+new_tail()\n",
            " \n".repeat(10)
        );
        let result = pipeline.compress(&diff, &ctx(), &preset).unwrap();
        assert!(result.data.contains("[18 unchanged lines]"), "{}", result.data);

        let mut samples = vec![0; 20_000];
        samples.push(1);
        let pretty = serde_json::to_string_pretty(&serde_json::json!({ "samples": samples })).unwrap();
        let result = pipeline.compress(&pretty, &ctx(), &preset).unwrap();
        let decoded = ToonEncoder.decode(&result.data).unwrap();
        let arr = decoded["samples"].as_array().unwrap();
        let (summary, kept) = arr.split_last().unwrap();
        let remaining: usize = summary
            .as_str()
            .and_then(|s| s.strip_prefix("... and "))
            .and_then(|s| s.strip_suffix(" more items"))
            .and_then(|n| n.parse().ok())
            .unwrap();
        assert_eq!(kept.len() + remaining, 20_001);
        assert!(kept.iter().all(|v| v == 0));
    }

    /// The lossless path returns JSON byte-for-byte — no stripped nulls,
    /// collapsed arrays, truncated strings, projected-away fields, or TOON
    /// re-encoding (which would shift line numbers).
    #[test]
    fn compress_lossless_preserves_json_fields() {
        let preset = default_preset();
        let pipeline = CompressionPipeline::new(&preset);
        let json = r#"{"keep":1,"gone":null,"items":[1,2,3,4,5,6,7,8],"note":"x"}"#;

        // Sanity: the default path transforms the JSON (strips the null,
        // TOON-encodes), so it differs from the input.
        let lossy = pipeline.compress(json, &ctx(), &preset).unwrap();
        assert_ne!(lossy.data, json, "default path should transform JSON");

        // The lossless path returns the JSON verbatim.
        let result = pipeline.compress_lossless(json, &ctx(), &preset).unwrap();
        assert_eq!(result.data, json, "lossless must return JSON byte-for-byte");
    }

    /// Even highly repetitive content is returned verbatim on the lossless
    /// path — RLE `[×N]` collapsing is skipped so a file read stays faithful.
    #[test]
    fn compress_lossless_returns_repeated_lines_verbatim() {
        let preset = default_preset();
        let pipeline = CompressionPipeline::new(&preset);
        let line = "this is a repeated diagnostic line that is quite long\n";
        let input = format!(
            "{}a unique trailing line of sufficient length here\n",
            line.repeat(6)
        );
        assert!(input.len() > 300);

        // Default path collapses the repeated run; lossless returns it untouched.
        let lossy = pipeline.compress(&input, &ctx(), &preset).unwrap();
        assert!(
            lossy.data.len() < input.len(),
            "default path should collapse repeats"
        );

        let result = pipeline.compress_lossless(&input, &ctx(), &preset).unwrap();
        assert_eq!(
            result.data, input,
            "lossless must return repeated content verbatim"
        );
        assert!(!result.stages_applied.contains(&"rle".to_owned()));
    }

    /// Benign-aware threshold: the same sub-1000-byte shape truncates when
    /// it carries an error marker and stays whole when it doesn't.
    #[test]
    fn benign_output_gets_double_truncation_headroom() {
        let preset = default_preset();
        let pipeline = CompressionPipeline::new(&preset);

        let mut segments = Vec::new();
        for i in 0..12 {
            if i % 2 == 0 {
                segments.push(format!("SEG{i} llllllllllllllllllllllllllllllllllllllll"));
            } else {
                segments.push(format!(
                    "SEG{i} the quick brown fox jumps over {i} lazy dogs by rivers"
                ));
            }
        }
        let benign = segments.join("\n\n");
        assert!(benign.len() > 500 && benign.len() < 1000, "fixture must sit between thresholds: {}", benign.len());

        let result = pipeline.compress(&benign, &ctx(), &preset).unwrap();
        assert!(
            !result.stages_applied.contains(&"entropy_truncate".to_owned()),
            "benign content under 1000 bytes must not be truncated"
        );

        let with_error = format!("error: something broke\n\n{benign}");
        let result = pipeline.compress(&with_error, &ctx(), &preset).unwrap();
        assert!(
            result.stages_applied.contains(&"entropy_truncate".to_owned()),
            "error-bearing content over 500 bytes should truncate"
        );
    }

    #[test]
    fn new_creates_pipeline_with_sorted_stages() {
        let preset = default_preset();
        let pipeline = CompressionPipeline::new(&preset);
        // Verify stages are sorted by priority
        let priorities: Vec<u32> = pipeline.stages.iter().map(|s| s.priority()).collect();
        let mut sorted = priorities.clone();
        sorted.sort();
        assert_eq!(priorities, sorted);
    }

    #[test]
    fn compress_plain_text_passthrough() {
        let preset = default_preset();
        let pipeline = CompressionPipeline::new(&preset);
        let result = pipeline.compress("hello world", &ctx(), &preset).unwrap();
        assert_eq!(result.data, "hello world");
        assert!(!result.stages_applied.contains(&"toon_encode".to_owned()));
    }

    #[test]
    fn compress_json_applies_toon() {
        let preset = default_preset();
        let pipeline = CompressionPipeline::new(&preset);
        let json = r#"{"name":"Alice","age":30}"#;
        let result = pipeline.compress(json, &ctx(), &preset).unwrap();
        assert!(result.data.starts_with("TOON:"), "data: {}", result.data);
        assert!(result.stages_applied.contains(&"toon_encode".to_owned()));
    }

    #[test]
    fn compress_strips_nulls_from_json() {
        let preset = default_preset();
        let pipeline = CompressionPipeline::new(&preset);
        let json = r#"{"a":1,"b":null}"#;
        let result = pipeline.compress(json, &ctx(), &preset).unwrap();
        // After strip_nulls, "b" is gone; TOON encodes the result
        assert!(result.data.starts_with("TOON:"));
        // Decode and verify null is gone
        let decoded = ToonEncoder.decode(&result.data).unwrap();
        assert!(decoded.get("b").is_none());
        assert_eq!(decoded["a"], serde_json::json!(1));
    }

    #[test]
    fn compress_returns_token_counts() {
        let preset = default_preset();
        let pipeline = CompressionPipeline::new(&preset);
        let input = "a".repeat(100);
        let result = pipeline.compress(&input, &ctx(), &preset).unwrap();
        assert!(result.tokens_original > 0);
        assert!(result.tokens_compressed > 0);
    }

    #[test]
    fn count_cache_holds_no_text() {
        let preset = default_preset();
        let pipeline = CompressionPipeline::new(&preset);
        let texts: Vec<String> = (0..COUNT_CACHE_MAX_ENTRIES + 8)
            .map(|i| format!("token=secret-{i:03} in a tool result\n").repeat(150))
            .collect();
        let family = model_family_from_preset(&preset);
        for (i, text) in texts.iter().enumerate() {
            let n = pipeline.token_counter.count(text, &family);
            assert_eq!(pipeline.count_tokens(text, &preset), n);
            assert_eq!(pipeline.count_tokens(text, &preset), n);
            let held = pipeline.token_counts.borrow().entries.len();
            assert_eq!(held, (i + 1).min(COUNT_CACHE_MAX_ENTRIES));
        }
        assert!(!std::mem::needs_drop::<(CountKey, u32)>());
        let cache = pipeline.token_counts.borrow();
        let bytes = cache.entries.capacity() * std::mem::size_of::<(CountKey, u32)>();
        assert!(bytes < texts[0].len(), "cache holds {bytes} bytes");
    }

    #[test]
    fn compress_ratio_is_reasonable() {
        let preset = default_preset();
        let pipeline = CompressionPipeline::new(&preset);
        let result = pipeline.compress("hello", &ctx(), &preset).unwrap();
        assert!(result.compression_ratio > 0.0);
    }

    #[test]
    fn insert_stage_re_sorts_by_priority() {
        use crate::stages::CompressionStage;
        use crate::types::StageConfig;

        struct LowPriorityStage;
        impl CompressionStage for LowPriorityStage {
            fn name(&self) -> &str {
                "low_priority"
            }
            fn priority(&self) -> u32 {
                5 // lower than all built-in stages
            }
            fn process(
                &self,
                _content: &mut Content,
                _config: &StageConfig,
            ) -> crate::error::Result<()> {
                Ok(())
            }
        }

        let preset = default_preset();
        let mut pipeline = CompressionPipeline::new(&preset);
        pipeline.insert_stage(Box::new(LowPriorityStage));

        let priorities: Vec<u32> = pipeline.stages.iter().map(|s| s.priority()).collect();
        let mut sorted = priorities.clone();
        sorted.sort();
        assert_eq!(priorities, sorted);
        assert_eq!(pipeline.stages[0].name(), "ansi_strip");
        assert_eq!(pipeline.stages[1].name(), "low_priority");
    }

    #[test]
    fn reload_preset_rebuilds_stages() {
        let preset = default_preset();
        let mut pipeline = CompressionPipeline::new(&preset);
        let original_count = pipeline.stages.len();
        pipeline.reload_preset(&preset).unwrap();
        assert_eq!(pipeline.stages.len(), original_count);
    }

    #[test]
    fn compress_keep_fields_filters_json() {
        use crate::preset::KeepFieldsConfig;
        let mut preset = default_preset();
        preset.compression.keep_fields = Some(KeepFieldsConfig {
            enabled: true,
            fields: vec!["id".into(), "name".into()],
        });
        let pipeline = CompressionPipeline::new(&preset);
        let json = r#"{"id":1,"name":"Bob","debug":"x"}"#;
        let result = pipeline.compress(json, &ctx(), &preset).unwrap();
        let decoded = ToonEncoder.decode(&result.data).unwrap();
        assert!(decoded.get("debug").is_none());
        assert_eq!(decoded["id"], serde_json::json!(1));
    }

    #[test]
    fn compress_empty_string() {
        let preset = default_preset();
        let pipeline = CompressionPipeline::new(&preset);
        let result = pipeline.compress("", &ctx(), &preset).unwrap();
        assert_eq!(result.data, "");
        assert_eq!(result.tokens_original, 0);
    }

    #[test]
    fn stage_config_from_preset_unknown_stage() {
        let preset = default_preset();
        let config = stage_config_from_preset("nonexistent", &preset);
        assert!(!config.enabled);
    }

    // ---------------------------------------------------------------------------
    // Property tests
    // ---------------------------------------------------------------------------

    use proptest::prelude::*;

    /// Generate a significant line from a fixed set of meaningful tokens.
    fn significant_line_strategy() -> impl Strategy<Value = String> {
        prop_oneof![
            Just("error: connection refused".to_owned()),
            Just("warning: deprecated API usage".to_owned()),
            Just("failed: build step exited with code 1".to_owned()),
            Just("success: deployment complete".to_owned()),
            Just("status: all checks passed".to_owned()),
            Just("error: file not found".to_owned()),
            Just("warning: unused variable detected".to_owned()),
        ]
    }

    /// Generate a noise line (repeated decorative content).
    fn noise_line_strategy() -> impl Strategy<Value = String> {
        prop_oneof![
            Just("---".to_owned()),
            Just("Loading...".to_owned()),
            Just("================".to_owned()),
            Just("...".to_owned()),
        ]
    }

    /// Recursive strategy that generates arbitrary serde_json::Value instances.
    /// Mirrors the strategy in toon.rs tests.
    fn arb_json_value() -> impl Strategy<Value = serde_json::Value> {
        let leaf = prop_oneof![
            Just(serde_json::Value::Null),
            any::<bool>().prop_map(serde_json::Value::Bool),
            any::<i64>().prop_map(|n| serde_json::json!(n)),
            any::<f64>()
                .prop_filter("must be finite", |f| f.is_finite())
                .prop_map(|f| serde_json::json!(f)),
            ".*".prop_map(serde_json::Value::String),
        ];

        leaf.prop_recursive(4, 64, 8, |inner| {
            prop_oneof![
                prop::collection::vec(inner.clone(), 0..8)
                    .prop_map(serde_json::Value::Array),
                prop::collection::hash_map(".*", inner, 0..8).prop_map(|m| {
                    serde_json::Value::Object(m.into_iter().collect())
                }),
            ]
        })
    }

    proptest! {
        /// **Validates: Requirements 17.1, 17.2, 13.2**
        ///
        /// Property 22: no raw control characters.
        ///
        /// For any JSON input, the Compression_Pipeline SHALL produce output
        /// whose only control characters are standard whitespace (\t, \n,
        /// \r). Text outside the control range is written as-is, the same as
        /// non-JSON output.
        #[test]
        fn prop_pipeline_ascii_safe_output(v in arb_json_value()) {
            let preset = default_preset();
            let pipeline = CompressionPipeline::new(&preset);

            let json_input = serde_json::to_string(&v).expect("serialize should not fail");
            let result = pipeline.compress(&json_input, &ctx(), &preset)
                .expect("compress should not fail");

            for ch in result.data.chars() {
                let is_standard_whitespace = matches!(ch, '\t' | '\n' | '\r');
                prop_assert!(
                    !ch.is_control() || is_standard_whitespace,
                    "raw control character in output: U+{:04X}\noutput: {:?}",
                    ch as u32, result.data
                );
            }
        }
    }

    proptest! {
        /// **Validates: Requirements 1.3**
        ///
        /// Property 1: Compression preserves semantically significant content.
        ///
        /// For any CLI output containing significant tokens (errors, warnings,
        /// status messages) mixed with noise (repeated identical lines), the
        /// Compression_Pipeline SHALL produce output that:
        ///   1. Contains all significant lines.
        ///   2. Contains each noise line at most `max_repeated_lines` times.
        #[test]
        fn prop_compression_preserves_significant_content(
            significant_lines in prop::collection::vec(significant_line_strategy(), 1..=5),
            noise_line in noise_line_strategy(),
            noise_repeat in 5u32..=10u32,
        ) {
            let preset = default_preset(); // condense enabled, max_repeated_lines=3
            let pipeline = CompressionPipeline::new(&preset);

            // Build interleaved input: noise, significant, noise, significant, ...
            let mut lines: Vec<String> = Vec::new();
            for sig in &significant_lines {
                for _ in 0..noise_repeat {
                    lines.push(noise_line.clone());
                }
                lines.push(sig.clone());
            }
            // Trailing noise block
            for _ in 0..noise_repeat {
                lines.push(noise_line.clone());
            }

            let input = lines.join("\n");
            let result = pipeline.compress(&input, &ctx(), &preset).unwrap();
            let output = &result.data;

            // 1. All significant lines must appear in the output.
            for sig in &significant_lines {
                prop_assert!(
                    output.contains(sig.as_str()),
                    "significant line missing from output: {:?}\noutput: {:?}",
                    sig,
                    output
                );
            }

            // 2. No consecutive run of the noise line exceeds max_repeated_lines (3).
            let mut max_run = 0usize;
            let mut current_run = 0usize;
            for line in output.lines() {
                if line == noise_line.as_str() {
                    current_run += 1;
                    max_run = max_run.max(current_run);
                } else {
                    current_run = 0;
                }
            }
            prop_assert!(
                max_run <= 3,
                "noise line {:?} has a consecutive run of {} (max 3)\noutput: {:?}",
                noise_line,
                max_run,
                output
            );
        }
    }
}