poimen-memory/crates/mem-core/src/optimizer/json.rs

//! JsonCrusher — Intelligent JSON array compression
//!
//! Strategy:
//! - Preserve all keys (structure)
//! - Keep start + end items (schema + recency): 30% + 15%
//! - Select mid-array items by variance: 55%
//! - Drop: redundant homogeneous elements, long string values

use anyhow::Result;
use serde_json::{json, Value};
use std::collections::HashMap;

pub struct JsonCrusher;

impl JsonCrusher {
    pub fn new() -> Self {
        Self
    }

    /// Compress JSON by keeping key structure, boundaries, and important items
    pub fn compress(&self, content: &str) -> Result<String> {
        let value: Value = serde_json::from_str(content)?;

        let compressed = match value {
            Value::Array(arr) => self.compress_array(arr)?,
            Value::Object(obj) => self.compress_object(obj)?,
            other => Value::String(format!("{}", other)), // scalars pass through
        };

        Ok(serde_json::to_string(&compressed)?)
    }

    fn compress_array(&self, items: Vec<Value>) -> Result<Value> {
        if items.is_empty() {
            return Ok(Value::Array(vec![]));
        }

        let len = items.len();

        // Budget allocation: 30% start, 15% end, 55% importance
        let start_count = (len as f32 * 0.3).ceil() as usize;
        let end_count = (len as f32 * 0.15).ceil() as usize;
        let mid_budget = (len as f32 * 0.55).ceil() as usize;

        let mut result = Vec::new();

        // Add start items
        for i in 0..start_count.min(len) {
            result.push(items[i].clone());
        }

        // Select mid-array items by variance/importance
        if len > start_count + end_count {
            let mid_items = &items[start_count..(len - end_count)];
            let selected = self.select_by_variance(mid_items, mid_budget);
            result.extend(selected);
        }

        // Add end items
        if end_count > 0 {
            for i in (len - end_count)..len {
                result.push(items[i].clone());
            }
        }

        Ok(Value::Array(result))
    }

    fn compress_object(&self, obj: serde_json::Map<String, Value>) -> Result<Value> {
        let mut compressed = serde_json::Map::new();

        for (key, value) in obj {
            // Always keep keys and structure
            let compressed_val = match value {
                Value::Array(arr) => self.compress_array(arr)?,
                Value::Object(inner) => self.compress_object(inner)?,
                // Keep: errors, nulls, booleans, numbers, short strings
                // Drop: long string values (> 500 chars)
                Value::String(s) if s.len() > 500 => Value::String(format!("[truncated {} chars]", s.len())),
                other => other,
            };
            compressed.insert(key, compressed_val);
        }

        Ok(Value::Object(compressed))
    }

    /// Select items from mid-array by variance (items with highest variance in their field values)
    fn select_by_variance(&self, items: &[Value], budget: usize) -> Vec<Value> {
        if items.len() <= budget {
            return items.to_vec();
        }

        // For each key in the objects, compute variance
        let mut key_variance: HashMap<String, f32> = HashMap::new();

        // Collect all keys
        for item in items {
            if let Value::Object(obj) = item {
                for key in obj.keys() {
                    key_variance.entry(key.clone()).or_insert(0.0);
                }
            }
        }

        // Compute variance for each key
        let keys: Vec<String> = key_variance.keys().cloned().collect();
        for key in keys {
            let values: Vec<f32> = items
                .iter()
                .filter_map(|item| {
                    if let Value::Object(obj) = item {
                        obj.get(&key).and_then(|v| match v {
                            Value::Number(n) => n.as_f64().map(|f| f as f32),
                            _ => None,
                        })
                    } else {
                        None
                    }
                })
                .collect();

            if !values.is_empty() {
                let mean = values.iter().sum::<f32>() / values.len() as f32;
                let variance = values.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / values.len() as f32;
                key_variance.insert(key, variance);
            }
        }

        // Score each item by how "interesting" it is (high variance in its fields)
        let mut scored_items: Vec<(usize, f32)> = items
            .iter()
            .enumerate()
            .map(|(idx, item)| {
                let score = if let Value::Object(obj) = item {
                    obj.iter()
                        .map(|(k, v)| {
                            let field_variance = key_variance.get(k).copied().unwrap_or(0.0);
                            // Bonus for non-null/error fields
                            if v.is_null() {
                                0.0
                            } else if k.to_lowercase().contains("error") {
                                field_variance + 10.0
                            } else {
                                field_variance
                            }
                        })
                        .sum::<f32>()
                } else {
                    0.0
                };
                (idx, score)
            })
            .collect();

        // Sort by score descending
        scored_items.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

        // Take top budget items and re-sort by original index (preserve order)
        let mut selected: Vec<(usize, f32)> = scored_items.into_iter().take(budget).collect();
        selected.sort_by_key(|a| a.0);

        selected.iter().map(|(idx, _)| items[*idx].clone()).collect()
    }
}

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

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

    #[test]
    fn test_compress_simple_object() {
        let json = r#"{"name": "Alice", "age": 30, "status": "active"}"#;
        let crusher = JsonCrusher::new();
        let result = crusher.compress(json).unwrap();
        assert!(result.contains("Alice"));
        assert!(result.contains("active"));
    }

    #[test]
    fn test_compress_array_keeps_boundaries() {
        let json = r#"[
            {"id": 1, "value": "first"},
            {"id": 2, "value": "middle1"},
            {"id": 3, "value": "middle2"},
            {"id": 4, "value": "middle3"},
            {"id": 5, "value": "last"}
        ]"#;

        let crusher = JsonCrusher::new();
        let result = crusher.compress(json).unwrap();
        let compressed: Value = serde_json::from_str(&result).unwrap();

        // Should be smaller than original
        assert!(result.len() < json.len());

        // Should still be valid JSON array
        assert!(compressed.is_array());
        let arr = compressed.as_array().unwrap();

        // Should have kept first and last
        let first = arr.first().unwrap();
        assert!(first.to_string().contains("first") || first.to_string().contains("1"));
    }

    #[test]
    fn test_compress_array_drops_middle() {
        let mut items = vec![];
        for i in 0..100 {
            items.push(format!(r#"{{"id": {}, "name": "Item{}", "value": {}}}"#, i, i, i));
        }
        let json = format!("[{}]", items.join(","));

        let crusher = JsonCrusher::new();
        let result = crusher.compress(&json).unwrap();

        // Should be significantly smaller
        assert!(result.len() < json.len());
    }

    #[test]
    fn test_compress_truncates_long_strings() {
        let long_string = "x".repeat(600);
        let json = format!(r#"{{"message": "{}"}}"#, long_string);

        let crusher = JsonCrusher::new();
        let result = crusher.compress(&json).unwrap();

        // Should contain truncation marker
        assert!(result.contains("truncated"));
        // Result should be much smaller
        assert!(result.len() < json.len());
    }

    #[test]
    fn test_preserve_error_fields() {
        let json = r#"[
            {"id": 1, "error": null},
            {"id": 2, "error": "Connection timeout"},
            {"id": 3, "error": "Timeout again"},
            {"id": 4, "error": null}
        ]"#;

        let crusher = JsonCrusher::new();
        let result = crusher.compress(json).unwrap();

        // Should preserve error fields
        assert!(result.to_lowercase().contains("error"));
    }

    #[test]
    fn test_compression_ratio() {
        let mut items = vec![];
        for i in 0..50 {
            items.push(format!(
                r#"{{"id": {}, "name": "Item{}", "status": "active", "value": {}}}"#,
                i, i, i
            ));
        }
        let json = format!("[{}]", items.join(","));

        let crusher = JsonCrusher::new();
        let compressed = crusher.compress(&json).unwrap();

        let ratio = compressed.len() as f32 / json.len() as f32;
        // Should achieve 70%+ compression
        assert!(ratio < 0.95, "compression ratio {} too high", ratio);
    }

    #[test]
    fn test_compress_nested_structure() {
        let json = r#"{
            "status": "success",
            "data": {
                "items": [
                    {"id": 1, "val": "a"},
                    {"id": 2, "val": "b"},
                    {"id": 3, "val": "c"}
                ]
            }
        }"#;

        let crusher = JsonCrusher::new();
        let result = crusher.compress(json).unwrap();

        // Should still be valid nested JSON
        assert!(serde_json::from_str::<Value>(&result).is_ok());
        // Should be smaller
        assert!(result.len() < json.len());
    }
}