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

//! Content Router — Magika ML + regex fallback detection

use anyhow::Result;
use magika::Session;
use std::sync::Mutex;

use super::ContentType;

pub struct ContentRouter {
    magika: Mutex<Session>,
    confidence_threshold: f32,
}

impl ContentRouter {
    /// Create router with default Magika session
    pub fn new() -> Result<Self> {
        let magika = Session::new()?;
        Ok(Self {
            magika: Mutex::new(magika),
            confidence_threshold: 0.7,
        })
    }

    /// Detect content type using Magika ML first, then regex fallback
    pub fn detect(&self, content: &str) -> Result<ContentType> {
        // Try Magika ML classification
        let mut magika = self.magika.lock().map_err(|e| anyhow::anyhow!("mutex lock failed: {}", e))?;
        if let Ok(result) = magika.identify_content_sync(content.as_bytes()) {
            let label = result.info().label;
            let score = result.score();

            if score >= self.confidence_threshold {
                if let Some(ct) = self.map_magika_label(label) {
                    return Ok(ct);
                }
            }
        }

        // Fallback to regex heuristics
        Ok(self.regex_fallback(content))
    }

    /// Map Magika label to our ContentType
    fn map_magika_label(&self, label: &str) -> Option<ContentType> {
        match label {
            "json" | "jsonl" => Some(ContentType::Json),
            "python"
            | "javascript"
            | "typescript"
            | "rust"
            | "go"
            | "shell"
            | "bash"
            | "java"
            | "cpp"
            | "csharp"
            | "sql" => Some(ContentType::Code),
            "diff" | "patch" => Some(ContentType::Diff),
            "yaml" | "toml" | "ini" | "xml" => Some(ContentType::Config),
            "markdown" | "txt" => None, // Fallback to regex for better detection
            _ => None,
        }
    }

    /// Regex-based fallback detection
    fn regex_fallback(&self, content: &str) -> ContentType {
        if is_json(content) {
            return ContentType::Json;
        }
        if is_log(content) {
            return ContentType::Log;
        }
        if is_diff(content) {
            return ContentType::Diff;
        }
        if is_code(content) {
            return ContentType::Code;
        }
        ContentType::Text
    }
}

/// Check if content is valid JSON
fn is_json(content: &str) -> bool {
    let trimmed = content.trim();
    if !((trimmed.starts_with('{') || trimmed.starts_with('['))) {
        return false;
    }
    serde_json::from_str::<serde_json::Value>(trimmed).is_ok()
}

/// Check if content looks like logs (timestamps, log levels, errors)
fn is_log(content: &str) -> bool {
    let content_lower = content.to_lowercase();

    // Log level markers
    let has_log_level = content_lower.contains("error:")
        || content_lower.contains("warn:")
        || content_lower.contains("info:")
        || content_lower.contains("debug:")
        || content_lower.contains("err!")
        || content.contains("ERROR")
        || content.contains("WARN")
        || content.contains("INFO")
        || content.contains("FAIL");

    // ISO timestamp pattern
    let has_timestamp = regex_contains(
        content,
        r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}",
    );

    // Common error markers
    let has_error_markers = content_lower.contains("exception")
        || content_lower.contains("stack trace")
        || content_lower.contains("at line")
        || content.contains("npm ERR!")
        || content.contains("error[")
        || content.contains("panic:");

    has_log_level || (has_timestamp && has_error_markers)
}

/// Check if content looks like a unified diff
fn is_diff(content: &str) -> bool {
    let has_diff_markers = content.contains("---") && content.contains("+++")
        || content.contains("@@");

    has_diff_markers && (content.contains("+") || content.contains("-"))
}

/// Check if content looks like source code
fn is_code(content: &str) -> bool {
    let lines: Vec<&str> = content.lines().collect();

    // Check for code-like patterns
    let import_use_pattern =
        regex_contains(content, r"^(import|use|require|include|from|package|class|def|fn|public|private|const|let|var|function)\b");
    let has_brackets =
        content.contains('{') && content.contains('}') || content.contains('[') && content.contains(']');
    let has_indentation = lines.iter().any(|line| line.starts_with('\t') || line.starts_with("  "));

    (import_use_pattern || has_brackets) && has_indentation
}

/// Helper to check if text contains a regex pattern
fn regex_contains(text: &str, pattern: &str) -> bool {
    if let Ok(re) = regex::Regex::new(pattern) {
        re.is_match(text)
    } else {
        false
    }
}

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

    #[test]
    fn test_json_detection() {
        let json = r#"{"key": "value", "number": 42}"#;
        assert!(is_json(json));
    }

    #[test]
    fn test_json_array_detection() {
        let json = r#"[1, 2, 3, {"nested": true}]"#;
        assert!(is_json(json));
    }

    #[test]
    fn test_invalid_json_rejected() {
        let not_json = r#"{"key": "value", invalid}"#;
        assert!(!is_json(not_json));
    }

    #[test]
    fn test_log_detection_with_error_level() {
        let log = "2026-08-28T09:15:00Z ERROR: connection timeout";
        assert!(is_log(log));
    }

    #[test]
    fn test_log_detection_with_npm_err() {
        let log = "npm ERR! 404 Not Found - [email protected]";
        assert!(is_log(log));
    }

    #[test]
    fn test_diff_detection() {
        let diff = "--- a/file.txt\n+++ b/file.txt\n@@ -1,3 +1,3 @@\n-old\n+new";
        assert!(is_diff(diff));
    }

    #[test]
    fn test_code_detection() {
        let code = r#"
fn main() {
    println!("hello");
}
        "#;
        assert!(is_code(code));
    }

    #[test]
    fn test_text_detection_fallback() {
        let text = "This is just plain English text without any special structure.";
        assert_eq!(regex_fallback_helper(text), ContentType::Text);
    }
}

// Helper for test fallback
fn regex_fallback_helper(content: &str) -> ContentType {
    if is_json(content) {
        return ContentType::Json;
    }
    if is_log(content) {
        return ContentType::Log;
    }
    if is_diff(content) {
        return ContentType::Diff;
    }
    if is_code(content) {
        return ContentType::Code;
    }
    ContentType::Text
}