src/local_ai.rs

use anyhow::{bail, Context, Result};
use base64::{engine::general_purpose::STANDARD as B64, Engine as _};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::{
    env, fs,
    io::Write,
    path::{Path, PathBuf},
    process::{Command, Stdio},
};
use uuid::Uuid;

#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct AiMessage {
    pub role: String,
    pub content: Value,
}

#[derive(Debug, Deserialize, Clone)]
pub struct ChatInput {
    pub model: Option<String>,
    pub messages: Vec<AiMessage>,
    pub temperature: Option<f32>,
    pub max_tokens: Option<u32>,
    pub think: Option<bool>,
}

#[derive(Debug, Deserialize)]
pub struct VisionInput {
    pub model: Option<String>,
    pub prompt: Option<String>,
    pub image_base64: String,
}

#[derive(Debug, Deserialize)]
pub struct TranscriptionInput {
    pub audio_base64: String,
    pub format: Option<String>,
    pub language: Option<String>,
}

#[derive(Debug, Deserialize)]
pub struct SpeechInput {
    pub text: String,
}

#[derive(Debug, Clone)]
pub struct BackendConfig {
    pub ollama_enabled: bool,
    pub ollama_url: String,
    pub llama_cpp_enabled: bool,
    pub llama_cpp_url: String,
}

impl BackendConfig {
    pub fn new(
        ollama_enabled: bool,
        ollama_url: Option<String>,
        llama_cpp_enabled: bool,
        llama_cpp_url: Option<String>,
    ) -> Self {
        Self {
            ollama_enabled,
            ollama_url: normalize_url(ollama_url.unwrap_or_else(default_ollama_url)),
            llama_cpp_enabled,
            llama_cpp_url: normalize_url(llama_cpp_url.unwrap_or_else(default_llama_cpp_url)),
        }
    }
}

pub async fn chat(
    client: &reqwest::Client,
    models: &Value,
    root: &Path,
    backend: &BackendConfig,
    input: ChatInput,
) -> Result<Value> {
    let hardware = hardware(root);
    let installed = installed_ollama_models(client, backend).await;
    let available_catalog = catalog_with_availability(models, &installed, root);
    let selected = select_model(
        &available_catalog,
        input.model.as_deref(),
        &hardware,
        "chat",
    )
    .context("no compatible chat model is configured")?;
    let profile_id = selected
        .get("id")
        .and_then(|v| v.as_str())
        .unwrap_or("auto");
    let temperature = input.temperature.unwrap_or(0.4);

    if backend.ollama_enabled {
        if let Some(ollama_model) = selected.get("ollama").and_then(|v| v.as_str()) {
            let mut options = json!({
                "temperature": temperature,
                "num_predict": input.max_tokens.unwrap_or(512),
                "num_ctx": effective_num_ctx(selected, &hardware)
            });
            merge_model_options(&mut options, selected);
            let mut body = json!({
                "model": ollama_model,
                "messages": input.messages,
                "stream": false,
                "keep_alive": model_keep_alive(selected),
                "options": options
            });
            if let Some(think) = input.think {
                body["think"] = json!(think);
            }
            if let Ok(response) = client
                .post(format!("{}/api/chat", backend.ollama_url))
                .json(&body)
                .send()
                .await
            {
                if response.status().is_success() {
                    let payload: Value =
                        response.json().await.context("invalid Ollama response")?;
                    let content = payload
                        .pointer("/message/content")
                        .and_then(|v| v.as_str())
                        .unwrap_or("");
                    return Ok(openai_response(
                        content,
                        profile_id,
                        "ollama",
                        payload.get("eval_count").cloned(),
                    ));
                }
            }
        }
    }

    if !backend.llama_cpp_enabled {
        bail!("Ollama is unavailable and llama.cpp is disabled by runtime configuration");
    }

    let llama_body = json!({
        "messages": input.messages,
        "stream": false,
        "temperature": temperature,
        "max_tokens": input.max_tokens.unwrap_or(512)
    });
    let response = client
        .post(format!("{}/v1/chat/completions", backend.llama_cpp_url))
        .json(&llama_body)
        .send()
        .await
        .context("Ollama is unavailable and llama.cpp is not reachable")?;
    if !response.status().is_success() {
        bail!("llama.cpp returned HTTP {}", response.status());
    }
    let payload: Value = response
        .json()
        .await
        .context("invalid llama.cpp response")?;
    Ok(payload)
}

pub async fn analyze_image(
    client: &reqwest::Client,
    models: &Value,
    root: &Path,
    backend: &BackendConfig,
    input: VisionInput,
) -> Result<Value> {
    if !backend.ollama_enabled {
        bail!("Ollama vision backend is disabled by runtime configuration");
    }
    let hardware = hardware(root);
    let installed = installed_ollama_models(client, backend).await;
    let available_catalog = catalog_with_availability(models, &installed, root);
    let selected = select_model(
        &available_catalog,
        input.model.as_deref(),
        &hardware,
        "vision",
    )
    .context("no compatible vision model is configured")?;
    let model_id = selected
        .get("id")
        .and_then(|v| v.as_str())
        .unwrap_or("vision");
    let ollama_model = selected
        .get("ollama")
        .and_then(|v| v.as_str())
        .context("the selected vision model has no Ollama model reference")?;
    let image = strip_data_url(&input.image_base64);
    B64.decode(image).context("image is not valid base64")?;
    let prompt = input.prompt.unwrap_or_else(|| {
        "Describe this image carefully and extract any readable text. State uncertainty explicitly.".to_string()
    });
    let mut options = json!({ "num_ctx": effective_num_ctx(selected, &hardware) });
    merge_model_options(&mut options, selected);
    let body = json!({
        "model": ollama_model,
        "messages": [{
            "role": "user",
            "content": prompt,
            "images": [image]
        }],
        "stream": false,
        "keep_alive": model_keep_alive(selected),
        "options": options
    });
    let response = client
        .post(format!("{}/api/chat", backend.ollama_url))
        .json(&body)
        .send()
        .await
        .context("Ollama is not reachable for vision analysis")?;
    if !response.status().is_success() {
        let status = response.status();
        let detail = response.text().await.unwrap_or_default();
        bail!(
            "Ollama vision request returned HTTP {status}: {}",
            detail.trim()
        );
    }
    let payload: Value = response
        .json()
        .await
        .context("invalid Ollama vision response")?;
    let content = payload
        .pointer("/message/content")
        .and_then(|v| v.as_str())
        .unwrap_or("");
    Ok(json!({
        "ok": true,
        "model": model_id,
        "backend_model": ollama_model,
        "content": content,
        "backend": "ollama"
    }))
}

pub async fn backend_status(client: &reqwest::Client, backend: &BackendConfig) -> Value {
    let ollama = if backend.ollama_enabled {
        match client
            .get(format!("{}/api/tags", backend.ollama_url))
            .send()
            .await
        {
            Ok(response) if response.status().is_success() => {
                json!({"available": true, "enabled": true, "url": backend.ollama_url})
            }
            Ok(response) => {
                json!({"available": false, "enabled": true, "url": backend.ollama_url, "error": format!("HTTP {}", response.status())})
            }
            Err(error) => {
                json!({"available": false, "enabled": true, "url": backend.ollama_url, "error": error.to_string()})
            }
        }
    } else {
        json!({"available": false, "enabled": false, "url": backend.ollama_url})
    };
    let llamacpp = if backend.llama_cpp_enabled {
        match client
            .get(format!("{}/v1/models", backend.llama_cpp_url))
            .send()
            .await
        {
            Ok(response) if response.status().is_success() => {
                json!({"available": true, "enabled": true, "url": backend.llama_cpp_url})
            }
            Ok(response) => {
                json!({"available": false, "enabled": true, "url": backend.llama_cpp_url, "error": format!("HTTP {}", response.status())})
            }
            Err(error) => {
                json!({"available": false, "enabled": true, "url": backend.llama_cpp_url, "error": error.to_string()})
            }
        }
    } else {
        json!({"available": false, "enabled": false, "url": backend.llama_cpp_url})
    };
    json!({"ollama": ollama, "llama_cpp": llamacpp})
}

pub async fn installed_ollama_models(
    client: &reqwest::Client,
    backend: &BackendConfig,
) -> Vec<String> {
    if !backend.ollama_enabled {
        return Vec::new();
    }
    let response = match client
        .get(format!("{}/api/tags", backend.ollama_url))
        .send()
        .await
    {
        Ok(value) if value.status().is_success() => value,
        _ => return Vec::new(),
    };
    let payload: Value = match response.json().await {
        Ok(value) => value,
        Err(_) => return Vec::new(),
    };
    payload
        .get("models")
        .and_then(|v| v.as_array())
        .into_iter()
        .flatten()
        .filter_map(|model| {
            model
                .get("name")
                .and_then(|v| v.as_str())
                .map(str::to_string)
        })
        .collect()
}

pub fn catalog_with_availability(models: &Value, installed: &[String], root: &Path) -> Value {
    let mut output = models.clone();
    if let Some(entries) = output.get_mut("models").and_then(|v| v.as_array_mut()) {
        for model in entries {
            let ollama = model.get("ollama").and_then(|v| v.as_str());
            let installed_ollama = ollama
                .map(|name| model_name_matches(installed, name))
                .unwrap_or(false);
            let gguf = model
                .get("gguf")
                .and_then(|v| v.as_str())
                .map(|path| root.join(path).is_file())
                .unwrap_or(false);
            model["available"] = json!(installed_ollama || gguf);
            model["installed_ollama"] = json!(installed_ollama);
            model["installed_gguf"] = json!(gguf);
        }
    }
    output
}

pub fn voice_status(root: &Path) -> Value {
    let whisper_binary = find_first(root, whisper_binary_candidates());
    let whisper_model = find_by_extension(&root.join("models/whisper"), &["bin"]);
    let piper_binary = find_first(root, piper_binary_candidates());
    let piper_model = find_by_extension(&root.join("models/piper"), &["onnx"]);
    json!({
        "stt": {
            "available": whisper_binary.is_some() && whisper_model.is_some(),
            "engine": "whisper.cpp",
            "binary": display_option(&whisper_binary),
            "model": display_option(&whisper_model)
        },
        "tts": {
            "available": piper_binary.is_some() && piper_model.is_some(),
            "engine": "piper",
            "binary": display_option(&piper_binary),
            "model": display_option(&piper_model)
        }
    })
}

pub fn transcribe(root: &Path, input: &TranscriptionInput) -> Result<Value> {
    let whisper_binary = find_first(root, whisper_binary_candidates())
        .context("whisper.cpp binary is not installed")?;
    let whisper_model = find_by_extension(&root.join("models/whisper"), &["bin"])
        .context("whisper.cpp model is not installed")?;
    let audio = B64
        .decode(strip_data_url(&input.audio_base64))
        .context("audio is not valid base64")?;
    let extension = sanitize_audio_extension(input.format.as_deref().unwrap_or("wav"));
    let id = Uuid::new_v4().to_string();
    let input_path = root
        .join("workspace/voice")
        .join(format!("{id}.{extension}"));
    let output_prefix = root
        .join("workspace/voice")
        .join(format!("{id}-transcript"));
    fs::write(&input_path, audio).context("cannot write temporary audio")?;

    let mut command = Command::new(&whisper_binary);
    command
        .arg("-m")
        .arg(&whisper_model)
        .arg("-f")
        .arg(&input_path)
        .arg("-otxt")
        .arg("-of")
        .arg(&output_prefix)
        .arg("-nt");
    if let Some(language) = input.language.as_deref() {
        if !language.trim().is_empty() {
            command.arg("-l").arg(language.trim());
        }
    }
    let output = command.output().context("cannot run whisper.cpp")?;
    let _ = fs::remove_file(&input_path);
    if !output.status.success() {
        bail!(
            "whisper.cpp failed: {}",
            String::from_utf8_lossy(&output.stderr).trim()
        );
    }
    let transcript_path = output_prefix.with_extension("txt");
    let text =
        fs::read_to_string(&transcript_path).context("whisper.cpp did not produce a transcript")?;
    let _ = fs::remove_file(transcript_path);
    Ok(json!({"ok": true, "text": text.trim(), "engine": "whisper.cpp"}))
}

pub fn synthesize(root: &Path, input: &SpeechInput) -> Result<Value> {
    if input.text.trim().is_empty() {
        bail!("text cannot be empty");
    }
    let piper_binary =
        find_first(root, piper_binary_candidates()).context("Piper binary is not installed")?;
    let piper_model = find_by_extension(&root.join("models/piper"), &["onnx"])
        .context("Piper voice model is not installed")?;
    let output_path = root
        .join("workspace/voice")
        .join(format!("{}.wav", Uuid::new_v4()));
    let mut child = Command::new(&piper_binary)
        .arg("--model")
        .arg(&piper_model)
        .arg("--output_file")
        .arg(&output_path)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
        .context("cannot run Piper")?;
    child
        .stdin
        .as_mut()
        .context("cannot open Piper stdin")?
        .write_all(input.text.as_bytes())
        .context("cannot send text to Piper")?;
    let output = child.wait_with_output().context("cannot wait for Piper")?;
    if !output.status.success() {
        bail!(
            "Piper failed: {}",
            String::from_utf8_lossy(&output.stderr).trim()
        );
    }
    let audio = fs::read(&output_path).context("Piper did not produce audio")?;
    let _ = fs::remove_file(output_path);
    Ok(json!({
        "ok": true,
        "engine": "piper",
        "mime_type": "audio/wav",
        "audio_base64": B64.encode(audio)
    }))
}

pub fn select_model<'a>(
    models: &'a Value,
    requested: Option<&str>,
    hardware: &Value,
    feature: &str,
) -> Option<&'a Value> {
    let entries = models.get("models")?.as_array()?;
    if let Some(requested) = requested.filter(|value| *value != "auto") {
        if let Some(model) = entries.iter().find(|model| {
            model.get("id").and_then(|v| v.as_str()) == Some(requested)
                || model.get("ollama").and_then(|v| v.as_str()) == Some(requested)
        }) {
            return if has_feature(model, feature)
                && model
                    .get("available")
                    .and_then(|v| v.as_bool())
                    .unwrap_or(false)
            {
                Some(model)
            } else {
                None
            };
        }
        if let Some(default_id) = models
            .get("profiles")?
            .as_array()?
            .iter()
            .find(|profile| profile.get("id").and_then(|v| v.as_str()) == Some(requested))
            .and_then(|profile| profile.get("recommended_default").and_then(|v| v.as_str()))
        {
            if let Some(model) = entries
                .iter()
                .find(|model| model.get("id").and_then(|v| v.as_str()) == Some(default_id))
            {
                return if has_feature(model, feature)
                    && model
                        .get("available")
                        .and_then(|v| v.as_bool())
                        .unwrap_or(false)
                {
                    Some(model)
                } else {
                    None
                };
            }
        }
        return None;
    }
    let ram = hardware
        .get("ram_gb")
        .and_then(|v| v.as_f64())
        .unwrap_or(8.0);
    let vram = hardware
        .get("vram_gb")
        .and_then(|v| v.as_f64())
        .unwrap_or(0.0);
    entries
        .iter()
        .filter(|model| has_feature(model, feature))
        .filter(|model| {
            model
                .get("available")
                .and_then(|v| v.as_bool())
                .unwrap_or(false)
        })
        .filter(|model| {
            // 3% slack: physical "32 GB" machines report ~31.7 GB usable.
            ram >= model
                .get("min_ram_gb")
                .and_then(|v| v.as_f64())
                .unwrap_or(0.0)
                * 0.97
        })
        .filter(|model| {
            vram >= model
                .get("min_vram_gb")
                .and_then(|v| v.as_f64())
                .unwrap_or(0.0)
                * 0.97
        })
        .last()
}

/// Context window for the request: the catalog's `default_context`, scaled to
/// the customer hardware so weak machines never thrash and strong machines
/// get the room they paid for. A per-model `options.num_ctx` overrides this.
pub fn effective_num_ctx(model: &Value, hardware: &Value) -> u64 {
    let base = model
        .get("default_context")
        .and_then(|v| v.as_u64())
        .unwrap_or(4096);
    let ram = hardware.get("ram_gb").and_then(|v| v.as_f64()).unwrap_or(8.0);
    let vram = hardware.get("vram_gb").and_then(|v| v.as_f64()).unwrap_or(0.0);
    if vram >= 16.0 || ram >= 32.0 {
        base.max(8192)
    } else if ram < 10.0 && vram < 4.0 {
        base.min(4096)
    } else {
        base
    }
}

/// How long Ollama keeps the model loaded after a request. Reloading a model
/// between messages is the single biggest perceived slowdown on USB targets.
pub fn model_keep_alive(model: &Value) -> String {
    model
        .get("keep_alive")
        .and_then(|v| v.as_str())
        .unwrap_or("30m")
        .to_string()
}

/// Merges the catalog model's optional `options` object (verbatim Ollama
/// options such as num_gpu, num_thread, top_p) over the computed defaults.
pub fn merge_model_options(options: &mut Value, model: &Value) {
    if let Some(extra) = model.get("options").and_then(|v| v.as_object()) {
        if let Some(target) = options.as_object_mut() {
            for (key, value) in extra {
                target.insert(key.clone(), value.clone());
            }
        }
    }
}

pub fn hardware(root: &Path) -> Value {
    let path = root.join("diagnostics/hardware.json");
    fs::read_to_string(path)
        .ok()
        .and_then(|text| serde_json::from_str(&text).ok())
        .unwrap_or_else(|| json!({"ram_gb": 8, "vram_gb": 0}))
}

fn has_feature(model: &Value, feature: &str) -> bool {
    model
        .get("features")
        .and_then(|v| v.as_array())
        .into_iter()
        .flatten()
        .any(|entry| entry.as_str() == Some(feature))
}

fn openai_response(content: &str, model: &str, backend: &str, eval_count: Option<Value>) -> Value {
    json!({
        "id": format!("jackailocal-{}", Uuid::new_v4()),
        "object": "chat.completion",
        "model": model,
        "backend": backend,
        "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
        "usage": {"completion_tokens": eval_count}
    })
}

fn model_name_matches(installed: &[String], expected: &str) -> bool {
    installed.iter().any(|name| {
        name == expected || name.trim_end_matches(":latest") == expected.trim_end_matches(":latest")
    })
}

fn default_ollama_url() -> String {
    env::var("OLLAMA_HOST").unwrap_or_else(|_| "http://127.0.0.1:11434".to_string())
}

fn default_llama_cpp_url() -> String {
    env::var("LLAMA_CPP_URL").unwrap_or_else(|_| "http://127.0.0.1:8080".to_string())
}

fn normalize_url(value: String) -> String {
    let value = value.trim().trim_end_matches('/').to_string();
    if value.starts_with("http://") || value.starts_with("https://") {
        value
    } else {
        format!("http://{value}")
    }
}

fn strip_data_url(value: &str) -> &str {
    value.split_once(',').map(|(_, data)| data).unwrap_or(value)
}

fn sanitize_audio_extension(value: &str) -> String {
    let extension = value.trim().trim_start_matches('.').to_ascii_lowercase();
    if matches!(extension.as_str(), "wav" | "mp3" | "flac" | "ogg") {
        extension
    } else {
        "wav".to_string()
    }
}

fn find_first(root: &Path, candidates: &[&str]) -> Option<PathBuf> {
    candidates
        .iter()
        .map(|candidate| root.join(candidate))
        .find(|path| path.is_file())
}

fn find_by_extension(dir: &Path, extensions: &[&str]) -> Option<PathBuf> {
    fs::read_dir(dir)
        .ok()?
        .flatten()
        .map(|entry| entry.path())
        .find(|path| {
            path.is_file()
                && path
                    .extension()
                    .and_then(|v| v.to_str())
                    .map(|v| extensions.contains(&v))
                    .unwrap_or(false)
        })
}

fn display_option(path: &Option<PathBuf>) -> Option<String> {
    path.as_ref()
        .map(|value| value.to_string_lossy().to_string())
}

#[cfg(target_os = "windows")]
fn whisper_binary_candidates() -> &'static [&'static str] {
    &[
        "backends/whisper.cpp/windows/whisper-cli.exe",
        "backends/whisper.cpp/windows/main.exe",
    ]
}
#[cfg(target_os = "macos")]
fn whisper_binary_candidates() -> &'static [&'static str] {
    &[
        "backends/whisper.cpp/macos/whisper-cli",
        "backends/whisper.cpp/macos/main",
    ]
}
#[cfg(all(unix, not(target_os = "macos")))]
fn whisper_binary_candidates() -> &'static [&'static str] {
    &[
        "backends/whisper.cpp/linux/whisper-cli",
        "backends/whisper.cpp/linux/main",
    ]
}

#[cfg(target_os = "windows")]
fn piper_binary_candidates() -> &'static [&'static str] {
    &["backends/piper/windows/piper.exe"]
}
#[cfg(target_os = "macos")]
fn piper_binary_candidates() -> &'static [&'static str] {
    &["backends/piper/macos/piper"]
}
#[cfg(all(unix, not(target_os = "macos")))]
fn piper_binary_candidates() -> &'static [&'static str] {
    &["backends/piper/linux/piper"]
}