mod cloud;

pub use cloud::{
    generate_anthropic, generate_custom, generate_openai, review_anthropic, review_custom,
    review_openai, split_anthropic, split_custom, split_openai,
};

use std::{
    collections::hash_map::DefaultHasher,
    hash::{Hash, Hasher},
    sync::Arc,
};

use mistralrs::{GgufModelBuilder, Model, RequestBuilder, TextMessageRole};
use tokio::sync::RwLock;

/// One selectable local (on-device) model. Larger models produce better commit messages
/// but take longer to download (first run only, then cached) and run slower on CPU.
#[derive(Debug, Clone, serde::Serialize)]
pub struct LocalModelOption {
    pub id: &'static str,
    pub label: &'static str,
    pub approx_size_mb: u32,
    repo: &'static str,
    file: &'static str,
    tokenizer_repo: &'static str,
}

pub const DEFAULT_LOCAL_MODEL_ID: &str = "qwen2.5-0.5b";

pub const LOCAL_MODELS: &[LocalModelOption] = &[
    LocalModelOption {
        id: "qwen2.5-0.5b",
        label: "Qwen2.5 0.5B Instruct — fast, lower quality",
        approx_size_mb: 490,
        repo: "Qwen/Qwen2.5-0.5B-Instruct-GGUF",
        file: "qwen2.5-0.5b-instruct-q4_k_m.gguf",
        tokenizer_repo: "Qwen/Qwen2.5-0.5B-Instruct",
    },
    LocalModelOption {
        id: "qwen2.5-1.5b",
        label: "Qwen2.5 1.5B Instruct — recommended",
        approx_size_mb: 1050,
        repo: "Qwen/Qwen2.5-1.5B-Instruct-GGUF",
        file: "qwen2.5-1.5b-instruct-q4_k_m.gguf",
        tokenizer_repo: "Qwen/Qwen2.5-1.5B-Instruct",
    },
    LocalModelOption {
        id: "qwen2.5-3b",
        label: "Qwen2.5 3B Instruct — best quality, slower",
        approx_size_mb: 2100,
        repo: "Qwen/Qwen2.5-3B-Instruct-GGUF",
        file: "qwen2.5-3b-instruct-q4_k_m.gguf",
        tokenizer_repo: "Qwen/Qwen2.5-3B-Instruct",
    },
];

fn find_local_model(model_id: &str) -> Option<&'static LocalModelOption> {
    LOCAL_MODELS.iter().find(|option| option.id == model_id)
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum LocalGenerationProfile {
    Fast,
    Balanced,
    Detailed,
}

impl Default for LocalGenerationProfile {
    fn default() -> Self {
        Self::Fast
    }
}

impl LocalGenerationProfile {
    pub fn from_id(value: Option<&str>) -> Self {
        match value
            .unwrap_or_default()
            .trim()
            .to_ascii_lowercase()
            .as_str()
        {
            "balanced" => Self::Balanced,
            "detailed" => Self::Detailed,
            _ => Self::Fast,
        }
    }

    pub fn diff_unified_context(self) -> &'static str {
        match self {
            Self::Fast => "--unified=1",
            Self::Balanced => "--unified=2",
            Self::Detailed => "--unified=3",
        }
    }

    fn max_diff_chars(self) -> usize {
        match self {
            Self::Fast => 8_000,
            Self::Balanced => 12_000,
            Self::Detailed => 24_000,
        }
    }

    fn max_output_tokens(self) -> usize {
        match self {
            Self::Fast => 160,
            Self::Balanced => 360,
            Self::Detailed => 750,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "lowercase")]
pub enum CommitAiPhase {
    /// Nothing has been requested yet.
    Idle,
    /// Downloading (first run only, then cached by hf-hub) and/or loading into memory.
    Loading,
    Ready,
    Error,
}

#[derive(Debug, Clone, serde::Serialize)]
pub struct CommitAiStatus {
    pub phase: CommitAiPhase,
    pub model_id: Option<String>,
    pub error: Option<String>,
}

struct Inner {
    phase: CommitAiPhase,
    model_id: Option<String>,
    error: Option<String>,
    model: Option<Arc<Model>>,
    cache: Option<GenerationCache>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct GenerationCacheKey {
    model_id: String,
    profile: LocalGenerationProfile,
    input_hash: u64,
}

#[derive(Debug, Clone)]
struct GenerationCache {
    key: GenerationCacheKey,
    message: String,
}

/// Manages the local (on-device) model only. Cloud providers are stateless HTTP calls
/// (see [`cloud`]) and don't need this — there's nothing to download or keep loaded.
#[derive(Clone)]
pub struct CommitAiEngine {
    inner: Arc<RwLock<Inner>>,
}

impl Default for CommitAiEngine {
    fn default() -> Self {
        Self {
            inner: Arc::new(RwLock::new(Inner {
                phase: CommitAiPhase::Idle,
                model_id: None,
                error: None,
                model: None,
                cache: None,
            })),
        }
    }
}

impl CommitAiEngine {
    pub fn new() -> Self {
        Self::default()
    }

    pub async fn status(&self) -> CommitAiStatus {
        let guard = self.inner.read().await;
        CommitAiStatus {
            phase: guard.phase,
            model_id: guard.model_id.clone(),
            error: guard.error.clone(),
        }
    }

    /// Downloads (first run only; hf-hub caches the files afterwards) and loads the given
    /// local model. Safe to call repeatedly — a call for the model that's already
    /// ready/loading is a no-op; a call for a *different* model switches to it (the
    /// previous one is dropped once no generation is still using it).
    pub async fn ensure_loaded(&self, model_id: &str) {
        {
            let guard = self.inner.read().await;
            let same_model = guard.model_id.as_deref() == Some(model_id);
            if same_model && matches!(guard.phase, CommitAiPhase::Ready | CommitAiPhase::Loading) {
                return;
            }
        }

        let Some(option) = find_local_model(model_id) else {
            let mut guard = self.inner.write().await;
            guard.phase = CommitAiPhase::Error;
            guard.model_id = Some(model_id.to_string());
            guard.error = Some(format!("Unknown local model: {model_id}"));
            guard.cache = None;
            return;
        };

        {
            let mut guard = self.inner.write().await;
            guard.phase = CommitAiPhase::Loading;
            guard.model_id = Some(model_id.to_string());
            guard.error = None;
            guard.model = None;
            guard.cache = None;
        }

        let result = GgufModelBuilder::new(option.repo, vec![option.file])
            .with_tok_model_id(option.tokenizer_repo)
            .with_logging()
            .build()
            .await;

        let mut guard = self.inner.write().await;
        // If the user switched to yet another model while this one was loading, drop this
        // (now stale) result instead of overwriting the newer request's state.
        if guard.model_id.as_deref() != Some(model_id) {
            return;
        }
        match result {
            Ok(model) => {
                guard.model = Some(Arc::new(model));
                guard.phase = CommitAiPhase::Ready;
                guard.error = None;
                guard.cache = None;
            }
            Err(err) => {
                guard.phase = CommitAiPhase::Error;
                guard.error = Some(err.to_string());
                guard.cache = None;
            }
        }
    }

    pub async fn generate_commit_message(
        &self,
        diff: &str,
        notes: Option<&str>,
        profile: LocalGenerationProfile,
    ) -> Result<String, String> {
        let (model, cache_key) = {
            let guard = self.inner.read().await;
            match (guard.phase, &guard.model) {
                (CommitAiPhase::Ready, Some(model)) => {
                    let cache_key = GenerationCacheKey {
                        model_id: guard.model_id.clone().unwrap_or_default(),
                        profile,
                        input_hash: generation_input_hash(diff, notes),
                    };
                    if let Some(cache) = &guard.cache {
                        if cache.key == cache_key {
                            return Ok(cache.message.clone());
                        }
                    }
                    (model.clone(), cache_key)
                }
                _ => return Err("The local AI model is not ready yet.".to_string()),
            }
        };

        let (system, user) = build_local_messages(diff, notes, profile)?;
        let request = RequestBuilder::new()
            .set_sampler_max_len(profile.max_output_tokens())
            .add_message(TextMessageRole::System, system)
            .add_message(TextMessageRole::User, user);

        let response = model
            .send_chat_request(request)
            .await
            .map_err(|err| err.to_string())?;

        let content = response
            .choices
            .first()
            .and_then(|choice| choice.message.content.clone())
            .ok_or_else(|| "The model did not return a response.".to_string())?;

        let message = sanitize_message(&content);
        if message.is_empty() {
            return Err("The model did not return a response.".to_string());
        }
        if looks_like_diff_echo(&message) {
            return Err(
                "The local model returned the diff instead of a commit message. Try a larger local model (1.5B or 3B) or a cloud provider.".to_string(),
            );
        }
        {
            let mut guard = self.inner.write().await;
            guard.cache = Some(GenerationCache {
                key: cache_key,
                message: message.clone(),
            });
        }
        Ok(message)
    }
}

fn generation_input_hash(diff: &str, notes: Option<&str>) -> u64 {
    let mut hasher = DefaultHasher::new();
    diff.hash(&mut hasher);
    notes.unwrap_or("").hash(&mut hasher);
    hasher.finish()
}

/// Models occasionally ignore the "no code fences" instruction (small local models
/// especially) — strip a wrapping ``` fence and wrapping quotes so the result can go
/// straight into the commit-message box.
pub(crate) fn sanitize_message(raw: &str) -> String {
    let mut text = raw.trim().to_string();
    if text.starts_with("```") {
        text = match text.split_once('\n') {
            // Drop the opening fence line (which may carry a language tag) and the closing fence.
            Some((_fence, rest)) => rest.trim_end().trim_end_matches("```").trim().to_string(),
            None => text.trim_matches('`').trim().to_string(),
        };
    }
    let trimmed = text.trim();
    if trimmed.len() >= 2 && trimmed.starts_with('"') && trimmed.ends_with('"') {
        return trimmed[1..trimmed.len() - 1].trim().to_string();
    }
    trimmed.to_string()
}

/// Weak models (small local ones especially) sometimes just echo the prompt's diff
/// sections back instead of writing a commit message. Catch that so the UI can show a
/// clear error instead of dumping raw diff text into the commit-message box.
pub(crate) fn looks_like_diff_echo(message: &str) -> bool {
    let lower = message.to_ascii_lowercase();
    lower.contains("diff --git")
        || lower.contains("staged files:")
        || lower.contains("staged changes:")
        || lower.contains("diff stat:")
        || lower.contains("detailed diff:")
        || message.lines().any(|line| line.starts_with("@@ "))
}

fn truncate_at_char_boundary(input: &str, max_chars: usize) -> String {
    if input.len() <= max_chars {
        return input.to_string();
    }

    let mut cut = max_chars;
    while !input.is_char_boundary(cut) {
        cut -= 1;
    }
    format!("{}\n\n[... diff truncated ...]", &input[..cut])
}

pub(crate) fn build_local_messages(
    diff: &str,
    notes: Option<&str>,
    profile: LocalGenerationProfile,
) -> Result<(String, String), String> {
    if diff.trim().is_empty() {
        return Err("No staged changes available for a commit message.".to_string());
    }

    let diff = truncate_at_char_boundary(diff, profile.max_diff_chars());
    // Appended to every profile below: small local models occasionally just echo the input
    // (the "Staged files:" / "Diff stat:" / "Detailed diff:" sections built in git.rs's
    // `staged_diff_local`) instead of writing a new commit message. Naming those exact
    // section headers here makes the failure mode explicit enough for weak models to avoid.
    const ANTI_ECHO: &str = " Never repeat, quote, or paraphrase the diff or its headers — \
do not include 'diff --git', '@@', 'Staged files:', 'Diff stat:', or 'Detailed diff:' \
anywhere in your answer.";
    let system = match profile {
        LocalGenerationProfile::Fast => {
            format!(
                "You generate Git commit messages. Respond only with one Conventional Commits subject line: <type>(<scope>): <subject>. Max 72 characters. No body, bullets, preamble, explanation, code fences, or quotes. Answer in English.{ANTI_ECHO}"
            )
        }
        LocalGenerationProfile::Balanced => {
            format!(
                "You generate Git commit messages. Respond with a Conventional Commits subject line, then a blank line, then at most two short bullet points. Keep every line under 72 characters. No preamble, explanation, code fences, or quotes. Answer in English.{ANTI_ECHO}"
            )
        }
        LocalGenerationProfile::Detailed => {
            format!(
                "You generate Git commit messages. Respond with a Conventional Commits subject line, then a blank line, then a concise body and up to four short bullet points grouped by affected area/file. Keep every line under 72 characters. No preamble, explanation, code fences, or quotes. Answer in English.{ANTI_ECHO}"
            )
        }
    };

    let mut user = String::new();
    if let Some(n) = notes.filter(|n| !n.trim().is_empty()) {
        user.push_str(&format!("Developer notes:\n{n}\n\n"));
    }
    user.push_str(&format!("Staged changes:\n{diff}"));
    Ok((system, user))
}

pub(crate) fn build_messages(diff: &str, notes: Option<&str>) -> Result<(String, String), String> {
    if diff.trim().is_empty() {
        return Err("No staged changes available for a commit message.".to_string());
    }

    // Rough token estimate — small models often have an 8-32k context window.
    const MAX_CHARS: usize = 24_000;
    let diff = truncate_at_char_boundary(diff, MAX_CHARS);

    let system = "You are a tool that writes a Git commit message describing a staged diff. \
Output ONLY the commit message text. Never quote, restate, or paraphrase the diff itself: \
do not include lines starting with 'diff --git', '@@', '+', '-', 'index ', 'Staged files:', \
or 'Diff stat:' anywhere in your answer. \
Format: a Conventional Commits header (<type>(<scope>): <subject>) in imperative mood, \
max. 72 characters, then a blank line, then a body. \
The body is required: a short, general paragraph (2-4 sentences) summarizing what changed \
and why at a high level — do NOT enumerate every changed file individually. \
You may optionally add up to 3 bullet points (- ...) afterward, but only for the most \
significant changes overall, never one bullet or heading per file. \
Never use bold text, backticks, or markdown headings for file names. \
Lines in the body max. 72 characters. \
No preamble, no explanation, no code fences, answer in English.\n\n\
Example:\n\
Diff:\n\
diff --git a/src/auth.py b/src/auth.py\n\
+def hash_password(pw):\n\
+    return bcrypt.hash(pw)\n\
diff --git a/src/routes.py b/src/routes.py\n\
-if password == stored_password:\n\
+if bcrypt.check(password, stored_password):\n\n\
Commit message:\n\
feat(auth): hash and verify passwords with bcrypt\n\n\
Passwords were previously compared as plain text. This adds a bcrypt-based\n\
hashing helper and updates the login check to verify against the hash\n\
instead of a direct string comparison.\n\n\
- Hash passwords on write, verify with bcrypt on login"
        .to_string();

    let mut user = String::new();
    if let Some(n) = notes.filter(|n| !n.trim().is_empty()) {
        user.push_str(&format!("Developer notes:\n{n}\n\n"));
    }
    user.push_str(&format!("Staged diff:\n{diff}"));
    Ok((system, user))
}

pub(crate) fn build_review_messages(diff: &str) -> Result<(String, String), String> {
    if diff.trim().is_empty() {
        return Err("No staged changes available for review.".to_string());
    }

    const MAX_CHARS: usize = 36_000;
    let diff = truncate_at_char_boundary(diff, MAX_CHARS);
    let system = r#"You are a senior software engineer performing a focused pre-commit review.
Review only the supplied staged Git diff. Look for concrete correctness bugs, security issues,
data loss, regressions, broken edge cases, unsafe error handling, and meaningful performance or
maintainability risks. Do not report formatting preferences or speculative nitpicks.

Return ONLY valid JSON with this exact shape:
{"summary":"one concise overall assessment","risk":"low|medium|high","findings":[{"severity":"critical|warning|info","title":"short title","description":"clear evidence and impact","file":"path or null","line":123,"suggestion":"specific safe next step"}]}

Use the new-file line number from the diff when it is known; otherwise use null. Use null for file
when the issue is repository-wide. Maximum 12 findings, ordered critical then warning then info.
If no actionable issue exists, return an empty findings array and risk low. Never use markdown,
code fences, commentary outside the JSON, or claim that tests were executed."#
        .to_string();
    let user = format!("Staged diff to review:\n{diff}");
    Ok((system, user))
}
