lingshu-core 0.10.0

//! # Model Router — smart routing & provider fallback
//!
//! WHY: Not every turn needs a $10/M-token frontier model. Simple
//! messages ("thanks", "what time is it?") can be routed to a cheaper
//! model, saving cost while preserving quality for complex turns.
//!
//! ```text
//!   user_message
//!       │
//!       ▼
//!   ┌──────────────────────┐
//!   │   classify(message)  │
//!   │   simple? complex?   │
//!   └─────────┬────────────┘
//!        simple│       │complex
//!             ▼        ▼
//!       cheap_model  primary_model
//!             │        │
//!             ▼        ▼
//!       ┌──────────────────┐
//!       │  try provider    │
//!       │  on error →      │
//!       │  try fallback    │
//!       └──────────────────┘
//! ```

use std::collections::HashSet;

use crate::config::{FallbackConfig, ModelConfig};
use lingshu_types::ApiMode;

// ─── Complexity classification ────────────────────────────────────────

/// Keywords that signal the message involves code, debugging, or tool use.
/// If any of these appear as a word in the message, we keep the primary model.
const COMPLEX_KEYWORDS: &[&str] = &[
    "debug",
    "debugging",
    "implement",
    "implementation",
    "refactor",
    "patch",
    "traceback",
    "stacktrace",
    "exception",
    "error",
    "analyze",
    "analysis",
    "investigate",
    "architecture",
    "design",
    "compare",
    "benchmark",
    "optimize",
    "optimise",
    "review",
    "terminal",
    "shell",
    "tool",
    "tools",
    "pytest",
    "test",
    "tests",
    "plan",
    "planning",
    "delegate",
    "subagent",
    "cron",
    "docker",
    "kubernetes",
    "code",
    "function",
    "class",
    "struct",
    "enum",
    "compile",
    "build",
    "deploy",
    "fix",
    "bug",
];

/// Short edit requests are often where cheap routing hurts most:
/// "replace labels", "change the UI", "add a file", etc.
/// We only escalate when an edit verb is paired with a likely implementation target.
const EDIT_INTENT_VERBS: &[&str] = &[
    "add",
    "change",
    "create",
    "delete",
    "edit",
    "implement",
    "modify",
    "move",
    "patch",
    "refactor",
    "remove",
    "rename",
    "replace",
    "rewrite",
    "update",
];

const IMPLEMENTATION_TARGETS: &[&str] = &[
    "agent",
    "api",
    "app",
    "backend",
    "button",
    "buttons",
    "cache",
    "class",
    "cli",
    "code",
    "component",
    "config",
    "css",
    "dialog",
    "endpoint",
    "file",
    "files",
    "flow",
    "footer",
    "form",
    "frontend",
    "function",
    "header",
    "html",
    "javascript",
    "js",
    "json",
    "label",
    "labels",
    "layout",
    "loop",
    "markdown",
    "md",
    "message",
    "messages",
    "model",
    "module",
    "page",
    "pages",
    "prompt",
    "python",
    "react",
    "repo",
    "repository",
    "route",
    "router",
    "rs",
    "rust",
    "screen",
    "session",
    "sidebar",
    "sql",
    "state",
    "style",
    "styles",
    "styling",
    "table",
    "test",
    "tests",
    "theme",
    "toml",
    "tsx",
    "typescript",
    "ts",
    "ui",
    "ux",
    "workflow",
    "yaml",
    "yml",
];

const CODEISH_SUFFIXES: &[&str] = &[
    ".c", ".cc", ".cpp", ".cs", ".css", ".go", ".h", ".hpp", ".html", ".java", ".js", ".json",
    ".jsx", ".kt", ".md", ".mjs", ".php", ".py", ".rb", ".rs", ".sh", ".sql", ".swift", ".toml",
    ".ts", ".tsx", ".txt", ".yaml", ".yml",
];

/// Thresholds for classifying a message as "simple".
#[derive(Debug, Clone)]
pub struct RoutingThresholds {
    pub max_chars: usize,
    pub max_words: usize,
    pub max_newlines: usize,
}

impl Default for RoutingThresholds {
    fn default() -> Self {
        Self {
            max_chars: 160,
            max_words: 28,
            max_newlines: 1,
        }
    }
}

/// WHY the message looks complex.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ComplexityReason {
    TooLong,
    TooManyWords,
    MultiLine,
    ContainsCodeFence,
    ContainsInlineCode,
    ContainsUrl,
    ContainsComplexKeyword(String),
    ContainsEditIntent(String),
    ContainsPathLikeToken(String),
    Empty,
}

/// Classify a user message as simple or complex.
///
/// Conservative: if in doubt, keep the primary (strong) model.
/// Only short, keyword-free, single-line plain-text messages qualify
/// as "simple".
pub fn classify_message(text: &str, thresholds: &RoutingThresholds) -> Option<ComplexityReason> {
    let text = text.trim();

    if text.is_empty() {
        return Some(ComplexityReason::Empty);
    }
    if text.len() > thresholds.max_chars {
        return Some(ComplexityReason::TooLong);
    }
    if text.split_whitespace().count() > thresholds.max_words {
        return Some(ComplexityReason::TooManyWords);
    }
    if text.chars().filter(|&c| c == '\n').count() > thresholds.max_newlines {
        return Some(ComplexityReason::MultiLine);
    }
    if text.contains("```") {
        return Some(ComplexityReason::ContainsCodeFence);
    }
    if text.contains('`') {
        return Some(ComplexityReason::ContainsInlineCode);
    }
    // URL detection: "http://" or "https://" or "www."
    if text.contains("http://") || text.contains("https://") || text.contains("www.") {
        return Some(ComplexityReason::ContainsUrl);
    }

    let lowered = text.to_lowercase();
    let words: HashSet<&str> = lowered
        .split_whitespace()
        .map(|w| w.trim_matches(|c: char| c.is_ascii_punctuation()))
        .filter(|w| !w.is_empty())
        .collect();

    for &kw in COMPLEX_KEYWORDS {
        if words.contains(kw) {
            return Some(ComplexityReason::ContainsComplexKeyword(kw.to_string()));
        }
    }

    if let Some(reason) = detect_edit_intent(&words) {
        return Some(reason);
    }

    if let Some(reason) = detect_path_like_token(text) {
        return Some(reason);
    }

    None // simple!
}

fn detect_edit_intent(words: &HashSet<&str>) -> Option<ComplexityReason> {
    let verb = EDIT_INTENT_VERBS
        .iter()
        .find(|verb| words.contains(**verb))
        .copied()?;
    let target = IMPLEMENTATION_TARGETS
        .iter()
        .find(|target| words.contains(**target))
        .copied()?;
    Some(ComplexityReason::ContainsEditIntent(format!(
        "{verb}+{target}"
    )))
}

fn detect_path_like_token(text: &str) -> Option<ComplexityReason> {
    for token in text.split_whitespace() {
        let token = token.trim_matches(|c: char| {
            c.is_ascii_punctuation() && c != '.' && c != '_' && c != '-' && c != '/' && c != '\\'
        });
        if token.is_empty() {
            continue;
        }

        let lowered = token.to_ascii_lowercase();
        if lowered.contains('/') || lowered.contains('\\') {
            return Some(ComplexityReason::ContainsPathLikeToken(token.to_string()));
        }
        if CODEISH_SUFFIXES
            .iter()
            .any(|suffix| lowered.ends_with(suffix))
        {
            return Some(ComplexityReason::ContainsPathLikeToken(token.to_string()));
        }
    }
    None
}

// ─── Route resolution ─────────────────────────────────────────────────

/// Describes which model + provider to use for a given turn.
#[derive(Debug, Clone)]
pub struct TurnRoute {
    pub model: String,
    pub api_mode: ApiMode,
    pub base_url: Option<String>,
    pub api_key_env: String,
    /// Human-readable label, e.g. "smart route → gpt-4.1-mini (openai)"
    pub label: Option<String>,
    /// Whether this is the primary model or a routed/fallback model.
    pub is_primary: bool,
}

/// Configuration for smart/cheap model routing.
#[derive(Debug, Clone, Default)]
pub struct SmartRoutingConfig {
    pub enabled: bool,
    pub cheap_model: String,
    pub cheap_base_url: Option<String>,
    pub cheap_api_key_env: Option<String>,
    pub thresholds: RoutingThresholds,
}

/// Resolve which model to use for a given user message.
///
/// Steps:
///   1. If smart routing is enabled and the message is simple → cheap model
///   2. Otherwise → primary model
pub fn resolve_turn_route(
    user_message: &str,
    model_config: &ModelConfig,
    smart_routing: &SmartRoutingConfig,
) -> TurnRoute {
    // Try cheap model for simple messages
    if smart_routing.enabled
        && !smart_routing.cheap_model.is_empty()
        && classify_message(user_message, &smart_routing.thresholds).is_none()
    {
        let base_url = smart_routing
            .cheap_base_url
            .clone()
            .or_else(|| model_config.base_url.clone());
        let api_key_env = smart_routing
            .cheap_api_key_env
            .clone()
            .unwrap_or_else(|| model_config.api_key_env.clone());

        let api_mode = ApiMode::detect(
            base_url.as_deref().unwrap_or(""),
            &smart_routing.cheap_model,
        );

        return TurnRoute {
            model: smart_routing.cheap_model.clone(),
            api_mode,
            base_url,
            api_key_env,
            label: Some(format!("smart route → {}", smart_routing.cheap_model)),
            is_primary: false,
        };
    }

    // Primary model
    let api_mode = ApiMode::detect(
        model_config.base_url.as_deref().unwrap_or(""),
        &model_config.default_model,
    );

    TurnRoute {
        model: model_config.default_model.clone(),
        api_mode,
        base_url: model_config.base_url.clone(),
        api_key_env: model_config.api_key_env.clone(),
        label: None,
        is_primary: true,
    }
}

/// Build a fallback route from config, for use when the primary provider errors.
pub fn fallback_route(fallback: &FallbackConfig) -> TurnRoute {
    let api_mode = ApiMode::detect(fallback.base_url.as_deref().unwrap_or(""), &fallback.model);

    TurnRoute {
        model: fallback.model.clone(),
        api_mode,
        base_url: fallback.base_url.clone(),
        api_key_env: fallback
            .api_key_env
            .clone()
            .unwrap_or_else(|| "OPENROUTER_API_KEY".into()),
        label: Some(format!("fallback → {}", fallback.model)),
        is_primary: false,
    }
}

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

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

    #[test]
    fn simple_message_classified_as_simple() {
        let thresholds = RoutingThresholds::default();
        assert!(classify_message("hello there!", &thresholds).is_none());
        assert!(classify_message("what's 2+2?", &thresholds).is_none());
        assert!(classify_message("thanks!", &thresholds).is_none());
    }

    #[test]
    fn complex_keyword_detected() {
        let thresholds = RoutingThresholds::default();
        let result = classify_message("please debug this", &thresholds);
        assert_eq!(
            result,
            Some(ComplexityReason::ContainsComplexKeyword("debug".into()))
        );
    }

    #[test]
    fn edit_intent_with_ui_target_is_complex() {
        let thresholds = RoutingThresholds::default();
        assert!(matches!(
            classify_message("replace the labels in the UI", &thresholds),
            Some(ComplexityReason::ContainsEditIntent(_))
        ));
    }

    #[test]
    fn edit_intent_with_code_target_is_complex() {
        let thresholds = RoutingThresholds::default();
        assert!(matches!(
            classify_message("add a Rust file for the session cache", &thresholds),
            Some(ComplexityReason::ContainsEditIntent(_))
        ));
    }

    #[test]
    fn long_message_is_complex() {
        let thresholds = RoutingThresholds::default();
        let long = "a ".repeat(100); // 200 chars
        assert_eq!(
            classify_message(&long, &thresholds),
            Some(ComplexityReason::TooLong)
        );
    }

    #[test]
    fn code_fence_is_complex() {
        let thresholds = RoutingThresholds::default();
        assert_eq!(
            classify_message("look at this ```code```", &thresholds),
            Some(ComplexityReason::ContainsCodeFence)
        );
    }

    #[test]
    fn inline_code_is_complex() {
        let thresholds = RoutingThresholds::default();
        assert_eq!(
            classify_message("what does `foo()` do?", &thresholds),
            Some(ComplexityReason::ContainsInlineCode)
        );
    }

    #[test]
    fn url_is_complex() {
        let thresholds = RoutingThresholds::default();
        assert!(matches!(
            classify_message("check https://example.com", &thresholds),
            Some(ComplexityReason::ContainsUrl)
        ));
    }

    #[test]
    fn multiline_is_complex() {
        let thresholds = RoutingThresholds::default();
        assert_eq!(
            classify_message("line1\nline2\nline3", &thresholds),
            Some(ComplexityReason::MultiLine)
        );
    }

    #[test]
    fn path_like_token_is_complex() {
        let thresholds = RoutingThresholds::default();
        assert!(matches!(
            classify_message("update src/agent.rs", &thresholds),
            Some(ComplexityReason::ContainsPathLikeToken(_))
        ));
        assert!(matches!(
            classify_message("change Cargo.toml", &thresholds),
            Some(ComplexityReason::ContainsPathLikeToken(_))
        ));
    }

    #[test]
    fn empty_message_is_complex() {
        let thresholds = RoutingThresholds::default();
        assert_eq!(
            classify_message("", &thresholds),
            Some(ComplexityReason::Empty)
        );
    }

    #[test]
    fn smart_routing_disabled_uses_primary() {
        let model_config = ModelConfig::default();
        let smart = SmartRoutingConfig::default(); // disabled
        let route = resolve_turn_route("hi", &model_config, &smart);
        assert!(route.is_primary);
        assert_eq!(route.model, model_config.default_model);
    }

    #[test]
    fn smart_routing_enabled_simple_message() {
        let model_config = ModelConfig::default();
        let smart = SmartRoutingConfig {
            enabled: true,
            cheap_model: "gpt-4.1-mini".into(),
            cheap_base_url: Some("https://api.openai.com/v1".into()),
            cheap_api_key_env: Some("OPENAI_API_KEY".into()),
            thresholds: RoutingThresholds::default(),
        };
        let route = resolve_turn_route("hi there!", &model_config, &smart);
        assert!(!route.is_primary);
        assert_eq!(route.model, "gpt-4.1-mini");
        assert!(route.label.is_some());
    }

    #[test]
    fn smart_routing_enabled_complex_message_uses_primary() {
        let model_config = ModelConfig::default();
        let smart = SmartRoutingConfig {
            enabled: true,
            cheap_model: "gpt-4.1-mini".into(),
            cheap_base_url: None,
            cheap_api_key_env: None,
            thresholds: RoutingThresholds::default(),
        };
        let route = resolve_turn_route("please debug this error", &model_config, &smart);
        assert!(route.is_primary);
    }

    #[test]
    fn smart_routing_keeps_primary_for_short_implementation_request() {
        let model_config = ModelConfig::default();
        let smart = SmartRoutingConfig {
            enabled: true,
            cheap_model: "gpt-4.1-mini".into(),
            cheap_base_url: None,
            cheap_api_key_env: None,
            thresholds: RoutingThresholds::default(),
        };
        let route = resolve_turn_route(
            "replace the labels in the React component",
            &model_config,
            &smart,
        );
        assert!(route.is_primary);
    }

    #[test]
    fn fallback_route_construction() {
        let fb = FallbackConfig {
            model: "anthropic/claude-haiku".into(),
            provider: "anthropic".into(),
            base_url: Some("https://api.anthropic.com/v1".into()),
            api_key_env: Some("ANTHROPIC_API_KEY".into()),
        };
        let route = fallback_route(&fb);
        assert!(!route.is_primary);
        assert_eq!(route.model, "anthropic/claude-haiku");
        assert_eq!(route.api_key_env, "ANTHROPIC_API_KEY");
        assert!(route.label.as_ref().is_some_and(|l| l.contains("fallback")));
    }
}