collet 0.1.1

//! Model capability profiles for strategy selection.
//!
//! Different LLM models have varying strengths — some excel at precise
//! text reproduction (needed for surgical edits), others work better with
//! diff-based approaches. This module provides model-aware strategy selection.

/// Preferred editing strategy based on model capabilities.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EditStrategy {
    /// High-accuracy models: prefer surgical file_edit.
    SurgicalOnly,
    /// Mid-tier models: prefer surgical edit but include file content
    /// hints on failure to help the model correct itself.
    SurgicalWithHint,
    /// Lower-accuracy models or large changes: prefer git_patch.
    PatchPreferred,
}

/// Cost tier for routing decisions.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CostTier {
    /// Economy models (flash variants).
    Cheap,
    /// Standard models.
    Standard,
    /// Premium models.
    Premium,
}

/// Capability profile for a specific LLM model.
#[derive(Debug, Clone)]
pub struct ModelProfile {
    /// Model identifier pattern (matched by prefix).
    pub name_prefix: &'static str,
    /// Maximum output tokens per request.
    pub max_output_tokens: u32,
    /// Whether the model supports function/tool calling.
    pub supports_tool_use: bool,
    /// Whether the model supports reasoning/thinking tokens.
    pub supports_reasoning: bool,
    /// Whether the model supports vision (image) inputs.
    pub supports_vision: bool,
    /// Preferred code editing strategy.
    pub edit_strategy: EditStrategy,
    /// Context window size.
    pub context_window: usize,
    /// Cost classification.
    pub cost_tier: CostTier,
}

/// Built-in model profiles for known models.
static PROFILES: &[ModelProfile] = &[
    // ── Z.ai / GLM ──────────────────────────────────────────────────────────
    ModelProfile {
        name_prefix: "glm-5",
        max_output_tokens: 8192,
        supports_tool_use: true,
        supports_reasoning: true,
        supports_vision: true,
        edit_strategy: EditStrategy::SurgicalOnly,
        context_window: 128_000,
        cost_tier: CostTier::Premium,
    },
    ModelProfile {
        name_prefix: "glm-4.7-flash",
        max_output_tokens: 4096,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::SurgicalWithHint,
        context_window: 128_000,
        cost_tier: CostTier::Cheap,
    },
    ModelProfile {
        name_prefix: "glm-4.7",
        max_output_tokens: 4096,
        supports_tool_use: true,
        supports_reasoning: true,
        supports_vision: false,
        edit_strategy: EditStrategy::SurgicalWithHint,
        context_window: 128_000,
        cost_tier: CostTier::Standard,
    },
    ModelProfile {
        name_prefix: "glm-4.6",
        max_output_tokens: 4096,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::SurgicalWithHint,
        context_window: 128_000,
        cost_tier: CostTier::Standard,
    },
    ModelProfile {
        name_prefix: "glm-4.5-air",
        max_output_tokens: 4096,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::PatchPreferred,
        context_window: 128_000,
        cost_tier: CostTier::Cheap,
    },
    ModelProfile {
        name_prefix: "glm-4.5",
        max_output_tokens: 4096,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::SurgicalWithHint,
        context_window: 128_000,
        cost_tier: CostTier::Standard,
    },
    // ── OpenAI ───────────────────────────────────────────────────────────────
    ModelProfile {
        name_prefix: "o4-mini",
        max_output_tokens: 16384,
        supports_tool_use: true,
        supports_reasoning: true,
        supports_vision: true,
        edit_strategy: EditStrategy::SurgicalOnly,
        context_window: 200_000,
        cost_tier: CostTier::Cheap,
    },
    ModelProfile {
        name_prefix: "o3-mini",
        max_output_tokens: 16384,
        supports_tool_use: true,
        supports_reasoning: true,
        supports_vision: false,
        edit_strategy: EditStrategy::SurgicalOnly,
        context_window: 200_000,
        cost_tier: CostTier::Cheap,
    },
    ModelProfile {
        name_prefix: "o3",
        max_output_tokens: 16384,
        supports_tool_use: true,
        supports_reasoning: true,
        supports_vision: true,
        edit_strategy: EditStrategy::SurgicalOnly,
        context_window: 200_000,
        cost_tier: CostTier::Premium,
    },
    ModelProfile {
        name_prefix: "o1-mini",
        max_output_tokens: 16384,
        supports_tool_use: true,
        supports_reasoning: true,
        supports_vision: false,
        edit_strategy: EditStrategy::SurgicalOnly,
        context_window: 128_000,
        cost_tier: CostTier::Cheap,
    },
    ModelProfile {
        name_prefix: "o1",
        max_output_tokens: 16384,
        supports_tool_use: true,
        supports_reasoning: true,
        supports_vision: true,
        edit_strategy: EditStrategy::SurgicalOnly,
        context_window: 128_000,
        cost_tier: CostTier::Premium,
    },
    ModelProfile {
        name_prefix: "gpt-4o-mini",
        max_output_tokens: 16384,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: true,
        edit_strategy: EditStrategy::SurgicalWithHint,
        context_window: 128_000,
        cost_tier: CostTier::Cheap,
    },
    ModelProfile {
        name_prefix: "gpt-4o",
        max_output_tokens: 16384,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: true,
        edit_strategy: EditStrategy::SurgicalOnly,
        context_window: 128_000,
        cost_tier: CostTier::Premium,
    },
    ModelProfile {
        name_prefix: "gpt-4-turbo",
        max_output_tokens: 4096,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: true,
        edit_strategy: EditStrategy::SurgicalWithHint,
        context_window: 128_000,
        cost_tier: CostTier::Standard,
    },
    ModelProfile {
        name_prefix: "gpt-4",
        max_output_tokens: 4096,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::SurgicalWithHint,
        context_window: 8_192,
        cost_tier: CostTier::Standard,
    },
    // ── Anthropic ────────────────────────────────────────────────────────────
    ModelProfile {
        name_prefix: "claude-opus",
        max_output_tokens: 16384,
        supports_tool_use: true,
        supports_reasoning: true,
        supports_vision: true,
        edit_strategy: EditStrategy::SurgicalOnly,
        context_window: 200_000,
        cost_tier: CostTier::Premium,
    },
    ModelProfile {
        name_prefix: "claude-sonnet",
        max_output_tokens: 16384,
        supports_tool_use: true,
        supports_reasoning: true,
        supports_vision: true,
        edit_strategy: EditStrategy::SurgicalOnly,
        context_window: 200_000,
        cost_tier: CostTier::Standard,
    },
    ModelProfile {
        name_prefix: "claude-haiku",
        max_output_tokens: 8192,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: true,
        edit_strategy: EditStrategy::SurgicalWithHint,
        context_window: 200_000,
        cost_tier: CostTier::Cheap,
    },
    // ── Google Gemini ────────────────────────────────────────────────────────
    ModelProfile {
        name_prefix: "gemini-2.5-pro",
        max_output_tokens: 16384,
        supports_tool_use: true,
        supports_reasoning: true,
        supports_vision: true,
        edit_strategy: EditStrategy::SurgicalOnly,
        context_window: 1_000_000,
        cost_tier: CostTier::Premium,
    },
    ModelProfile {
        name_prefix: "gemini-2.5-flash",
        max_output_tokens: 16384,
        supports_tool_use: true,
        supports_reasoning: true,
        supports_vision: true,
        edit_strategy: EditStrategy::SurgicalOnly,
        context_window: 1_000_000,
        cost_tier: CostTier::Cheap,
    },
    ModelProfile {
        name_prefix: "gemini-2.0-flash",
        max_output_tokens: 8192,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: true,
        edit_strategy: EditStrategy::SurgicalWithHint,
        context_window: 1_000_000,
        cost_tier: CostTier::Cheap,
    },
    ModelProfile {
        name_prefix: "gemini-1.5-pro",
        max_output_tokens: 8192,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: true,
        edit_strategy: EditStrategy::SurgicalWithHint,
        context_window: 1_000_000,
        cost_tier: CostTier::Standard,
    },
    // ── DeepSeek ─────────────────────────────────────────────────────────────
    ModelProfile {
        name_prefix: "deepseek-reasoner",
        max_output_tokens: 8192,
        supports_tool_use: true,
        supports_reasoning: true,
        supports_vision: false,
        edit_strategy: EditStrategy::SurgicalWithHint,
        context_window: 64_000,
        cost_tier: CostTier::Standard,
    },
    ModelProfile {
        name_prefix: "deepseek-chat",
        max_output_tokens: 8192,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::SurgicalWithHint,
        context_window: 64_000,
        cost_tier: CostTier::Cheap,
    },
    ModelProfile {
        name_prefix: "deepseek-coder",
        max_output_tokens: 8192,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::SurgicalWithHint,
        context_window: 64_000,
        cost_tier: CostTier::Cheap,
    },
    // ── Mistral ──────────────────────────────────────────────────────────────
    ModelProfile {
        name_prefix: "mistral-large",
        max_output_tokens: 8192,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::SurgicalWithHint,
        context_window: 128_000,
        cost_tier: CostTier::Standard,
    },
    ModelProfile {
        name_prefix: "codestral",
        max_output_tokens: 8192,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::SurgicalWithHint,
        context_window: 256_000,
        cost_tier: CostTier::Standard,
    },
    ModelProfile {
        name_prefix: "devstral",
        max_output_tokens: 8192,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::SurgicalWithHint,
        context_window: 128_000,
        cost_tier: CostTier::Standard,
    },
    ModelProfile {
        name_prefix: "mistral-small",
        max_output_tokens: 4096,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::PatchPreferred,
        context_window: 128_000,
        cost_tier: CostTier::Cheap,
    },
    // ── Groq (hosted) ────────────────────────────────────────────────────────
    ModelProfile {
        name_prefix: "llama-3.3",
        max_output_tokens: 8192,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::PatchPreferred,
        context_window: 131_072,
        cost_tier: CostTier::Cheap,
    },
    ModelProfile {
        name_prefix: "llama3-",
        max_output_tokens: 8192,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::PatchPreferred,
        context_window: 131_072,
        cost_tier: CostTier::Cheap,
    },
    // ── Ollama local models ──────────────────────────────────────────────────
    // Ollama models typically use "name:tag" format; these prefixes catch the
    // most common pulls.  Local models generally lack tool-call support.
    ModelProfile {
        name_prefix: "llama3:",
        max_output_tokens: 4096,
        supports_tool_use: false,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::PatchPreferred,
        context_window: 8_192,
        cost_tier: CostTier::Cheap,
    },
    ModelProfile {
        name_prefix: "qwen2.5-coder:",
        max_output_tokens: 4096,
        supports_tool_use: false,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::PatchPreferred,
        context_window: 32_768,
        cost_tier: CostTier::Cheap,
    },
    ModelProfile {
        name_prefix: "qwen2.5:",
        max_output_tokens: 4096,
        supports_tool_use: false,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::PatchPreferred,
        context_window: 32_768,
        cost_tier: CostTier::Cheap,
    },
    ModelProfile {
        name_prefix: "codellama:",
        max_output_tokens: 4096,
        supports_tool_use: false,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::PatchPreferred,
        context_window: 16_384,
        cost_tier: CostTier::Cheap,
    },
    ModelProfile {
        name_prefix: "deepseek-coder:",
        max_output_tokens: 4096,
        supports_tool_use: false,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::PatchPreferred,
        context_window: 16_384,
        cost_tier: CostTier::Cheap,
    },
    ModelProfile {
        name_prefix: "mistral:",
        max_output_tokens: 4096,
        supports_tool_use: false,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::PatchPreferred,
        context_window: 8_192,
        cost_tier: CostTier::Cheap,
    },
    ModelProfile {
        name_prefix: "phi3:",
        max_output_tokens: 4096,
        supports_tool_use: false,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::PatchPreferred,
        context_window: 8_192,
        cost_tier: CostTier::Cheap,
    },
    ModelProfile {
        name_prefix: "gemma:",
        max_output_tokens: 4096,
        supports_tool_use: false,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::PatchPreferred,
        context_window: 8_192,
        cost_tier: CostTier::Cheap,
    },
    // ── LM Studio local models (community GGUF format) ─────────────────────
    // LM Studio uses "org/model-GGUF" naming. Tool calling is supported
    // per LM Studio docs (localhost:1234/v1).
    ModelProfile {
        name_prefix: "lmstudio-community/",
        max_output_tokens: 4096,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::PatchPreferred,
        context_window: 32_768,
        cost_tier: CostTier::Cheap,
    },
    // ── OpenRouter pass-through prefixes ─────────────────────────────────────
    ModelProfile {
        name_prefix: "anthropic/claude-",
        max_output_tokens: 16384,
        supports_tool_use: true,
        supports_reasoning: true,
        supports_vision: true,
        edit_strategy: EditStrategy::SurgicalOnly,
        context_window: 200_000,
        cost_tier: CostTier::Standard,
    },
    ModelProfile {
        name_prefix: "openai/gpt-4o",
        max_output_tokens: 16384,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: true,
        edit_strategy: EditStrategy::SurgicalOnly,
        context_window: 128_000,
        cost_tier: CostTier::Premium,
    },
    ModelProfile {
        name_prefix: "openai/o3",
        max_output_tokens: 16384,
        supports_tool_use: true,
        supports_reasoning: true,
        supports_vision: true,
        edit_strategy: EditStrategy::SurgicalOnly,
        context_window: 200_000,
        cost_tier: CostTier::Premium,
    },
    ModelProfile {
        name_prefix: "google/gemini-2.5",
        max_output_tokens: 16384,
        supports_tool_use: true,
        supports_reasoning: true,
        supports_vision: true,
        edit_strategy: EditStrategy::SurgicalOnly,
        context_window: 1_000_000,
        cost_tier: CostTier::Standard,
    },
    // ── Together AI (meta-llama prefixed) ────────────────────────────────────
    ModelProfile {
        name_prefix: "meta-llama/",
        max_output_tokens: 8192,
        supports_tool_use: true,
        supports_reasoning: false,
        supports_vision: false,
        edit_strategy: EditStrategy::PatchPreferred,
        context_window: 131_072,
        cost_tier: CostTier::Cheap,
    },
];

/// Default profile for unknown models.
/// `supports_tool_use` is **false** to avoid sending tool JSON to models that
/// cannot parse it (the original bug trigger).
static DEFAULT_PROFILE: ModelProfile = ModelProfile {
    name_prefix: "",
    max_output_tokens: 4096,
    supports_tool_use: false,
    supports_reasoning: false,
    supports_vision: false,
    edit_strategy: EditStrategy::SurgicalWithHint,
    context_window: 120_000,
    cost_tier: CostTier::Standard,
};

/// Look up the capability profile for a model name.
///
/// Matches by longest prefix. Falls back to a sensible default.
pub fn profile_for(model: &str) -> &'static ModelProfile {
    PROFILES
        .iter()
        .filter(|p| model.starts_with(p.name_prefix))
        .max_by_key(|p| p.name_prefix.len())
        .unwrap_or(&DEFAULT_PROFILE)
}

/// Get the recommended context window size for a model.
pub fn context_window_for(model: &str) -> usize {
    profile_for(model).context_window
}

/// Get the recommended edit strategy for a model.
pub fn edit_strategy_for(model: &str) -> EditStrategy {
    profile_for(model).edit_strategy
}

/// Check whether a model supports vision (image) inputs.
pub fn supports_vision_for(model: &str) -> bool {
    profile_for(model).supports_vision
}

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

    #[test]
    fn test_profile_glm5() {
        let p = profile_for("glm-5");
        assert_eq!(p.edit_strategy, EditStrategy::SurgicalOnly);
        assert_eq!(p.cost_tier, CostTier::Premium);
        assert!(p.supports_reasoning);
    }

    #[test]
    fn test_profile_glm47() {
        let p = profile_for("glm-4.7");
        assert_eq!(p.edit_strategy, EditStrategy::SurgicalWithHint);
        assert_eq!(p.cost_tier, CostTier::Standard);
    }

    #[test]
    fn test_profile_glm47_flash() {
        let p = profile_for("glm-4.7-flash");
        assert_eq!(p.edit_strategy, EditStrategy::SurgicalWithHint);
        assert_eq!(p.cost_tier, CostTier::Cheap);
    }

    #[test]
    fn test_profile_glm45_flash() {
        let p = profile_for("glm-4.5-air");
        assert_eq!(p.edit_strategy, EditStrategy::PatchPreferred);
        assert_eq!(p.cost_tier, CostTier::Cheap);
    }

    #[test]
    fn test_profile_unknown_model() {
        let p = profile_for("some-unknown-model");
        assert_eq!(p.edit_strategy, EditStrategy::SurgicalWithHint);
        assert_eq!(p.cost_tier, CostTier::Standard);
        assert!(
            !p.supports_tool_use,
            "unknown models must default to no tool support"
        );
    }

    #[test]
    fn test_longest_prefix_match() {
        let p = profile_for("glm-4.7-flash");
        assert!(!p.supports_reasoning);
        assert_eq!(p.cost_tier, CostTier::Cheap);
    }

    #[test]
    fn test_context_window_for() {
        assert_eq!(context_window_for("glm-5"), 128_000);
    }

    #[test]
    fn test_edit_strategy_for() {
        assert_eq!(edit_strategy_for("glm-5"), EditStrategy::SurgicalOnly);
        assert_eq!(
            edit_strategy_for("glm-4.5-air"),
            EditStrategy::PatchPreferred
        );
    }

    #[test]
    fn test_openai_profiles() {
        let p = profile_for("gpt-4o");
        assert!(p.supports_tool_use);
        assert_eq!(p.edit_strategy, EditStrategy::SurgicalOnly);

        let p = profile_for("o3-mini");
        assert!(p.supports_tool_use);
        assert!(p.supports_reasoning);
    }

    #[test]
    fn test_anthropic_profiles() {
        let p = profile_for("claude-sonnet-4-6");
        assert!(p.supports_tool_use);
        assert!(p.supports_reasoning);
        assert_eq!(p.context_window, 200_000);
    }

    #[test]
    fn test_ollama_local_no_tools() {
        let p = profile_for("llama3:8b");
        assert!(
            !p.supports_tool_use,
            "Ollama local models should not support tools"
        );
        assert_eq!(p.edit_strategy, EditStrategy::PatchPreferred);

        let p = profile_for("qwen2.5-coder:7b");
        assert!(!p.supports_tool_use);

        let p = profile_for("codellama:13b");
        assert!(!p.supports_tool_use);
    }

    #[test]
    fn test_gemini_profiles() {
        let p = profile_for("gemini-2.5-pro");
        assert!(p.supports_tool_use);
        assert!(p.supports_reasoning);
        assert_eq!(p.context_window, 1_000_000);
    }

    #[test]
    fn test_deepseek_profiles() {
        let p = profile_for("deepseek-reasoner");
        assert!(p.supports_reasoning);
        assert!(p.supports_tool_use);

        let p = profile_for("deepseek-chat");
        assert!(!p.supports_reasoning);
    }

    #[test]
    fn test_openrouter_passthrough() {
        let p = profile_for("anthropic/claude-sonnet-4");
        assert!(p.supports_tool_use);
        assert_eq!(p.context_window, 200_000);
    }

    #[test]
    fn test_together_ai() {
        let p = profile_for("meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo");
        assert!(p.supports_tool_use);
        assert_eq!(p.edit_strategy, EditStrategy::PatchPreferred);
    }

    #[test]
    fn test_supports_vision() {
        // Vision-capable models
        assert!(supports_vision_for("gpt-4o"));
        assert!(supports_vision_for("gpt-4o-mini"));
        assert!(supports_vision_for("gpt-4-turbo"));
        assert!(supports_vision_for("claude-sonnet-4-6"));
        assert!(supports_vision_for("claude-opus-4"));
        assert!(supports_vision_for("gemini-2.5-pro"));
        assert!(supports_vision_for("o3"));
        assert!(supports_vision_for("o4-mini"));
        assert!(supports_vision_for("glm-5"));

        // Non-vision models
        assert!(!supports_vision_for("gpt-4"));
        assert!(!supports_vision_for("o3-mini"));
        assert!(!supports_vision_for("o1-mini"));
        assert!(!supports_vision_for("deepseek-chat"));
        assert!(!supports_vision_for("codestral"));
        assert!(!supports_vision_for("llama3:8b"));
        assert!(!supports_vision_for("some-unknown-model"));
    }
}