//! Every model the app can run, across providers, in one catalog the pickers
//! read: Claude's from the CLI's `initialize` reply, Pi's from asking `pi` for
//! the models its configured providers offer. Also the user's favorites.

use std::{
    io::{BufRead, BufReader, Write},
    process::{Command, Stdio},
    sync::mpsc,
    time::Duration,
};

use gpui_kit::*;
use serde_json::{Value, json};

use crate::claude::{Effort, find_pi_binary};
use crate::settings;

/// The agent behind a thread.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
pub enum Provider {
    #[default]
    Claude,
    Pi,
}

impl Provider {
    pub fn label(self) -> &'static str {
        match self {
            Provider::Claude => "Claude",
            Provider::Pi => "Pi",
        }
    }

    pub fn key(self) -> &'static str {
        match self {
            Provider::Claude => "claude",
            Provider::Pi => "pi",
        }
    }

    pub fn from_key(key: &str) -> Option<Self> {
        [Provider::Claude, Provider::Pi].into_iter().find(|provider| provider.key() == key)
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct ModelEntry {
    pub provider: Provider,
    /// What the agent is launched with: a Claude alias or model id, or Pi's
    /// `provider/model`.
    pub id: String,
    /// The concrete model, e.g. `claude-opus-5-5` behind the `opus` alias.
    pub resolved: String,
    pub name: String,
    pub detail: String,
    /// Older Claude models, listed under "Legacy models".
    pub legacy: bool,
    /// Reasoning levels the model accepts.
    pub efforts: Vec<Effort>,
    pub fast_mode: bool,
    pub auto_mode: bool,
    pub context_window: Option<u64>,
}

impl ModelEntry {
    /// The key favorites are stored under.
    pub fn key(&self) -> String {
        format!("{}:{}", self.provider.key(), self.id)
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum PiStatus {
    Loading,
    Ready,
    /// `pi` isn't installed (or isn't on a path we look at).
    Missing,
    Failed(String),
}

pub struct Catalog {
    pub claude: Vec<ModelEntry>,
    /// Claude's `default` entry, which the picker shows as the model it resolves to.
    pub claude_default: Option<ModelEntry>,
    pub pi: Vec<ModelEntry>,
    pub pi_status: PiStatus,
    /// The signed-in Claude account, once a Claude session reported it.
    pub claude_account: Option<String>,
    favorites: Vec<String>,
}

impl Global for Catalog {}

impl Catalog {
    pub fn init(cx: &mut App) {
        cx.set_global(Catalog {
            claude: Vec::new(),
            claude_default: None,
            pi: Vec::new(),
            pi_status: PiStatus::Loading,
            claude_account: None,
            favorites: settings::load().favorites,
        });
        Self::refresh_pi(cx);
    }

    /// Asks Pi (again) for the models its providers offer.
    pub fn refresh_pi(cx: &mut App) {
        cx.global_mut::<Catalog>().pi_status = PiStatus::Loading;
        cx.refresh_windows();
        let probe = cx.background_executor().spawn(async { probe_pi_models() });
        cx.spawn(async move |cx| {
            let result = probe.await;
            cx.update(|cx| {
                let catalog = cx.global_mut::<Catalog>();
                match result {
                    Ok(models) => {
                        catalog.pi = models;
                        catalog.pi_status = PiStatus::Ready;
                    }
                    Err(status) => {
                        catalog.pi.clear();
                        catalog.pi_status = status;
                    }
                }
                cx.refresh_windows();
            })
        })
        .detach();
    }

    pub fn global(cx: &App) -> &Catalog {
        cx.global::<Catalog>()
    }

    /// Takes Claude's model list and account from a session's `initialize` reply.
    pub fn set_claude_models(info: &Value, cx: &mut App) {
        if let Some(account) = info.pointer("/account/email").and_then(Value::as_str) {
            cx.global_mut::<Catalog>().claude_account = Some(account.to_string());
        }
        let (default, models) = claude_models_from(info);
        if models.is_empty() {
            return;
        }
        let catalog = cx.global_mut::<Catalog>();
        if catalog.claude != models {
            catalog.claude = models;
            catalog.claude_default = default;
            cx.refresh_windows();
        }
    }

    pub fn models(&self, provider: Provider) -> &[ModelEntry] {
        match provider {
            Provider::Claude => &self.claude,
            Provider::Pi => &self.pi,
        }
    }

    /// The entry for a thread's model; `None` means the provider's default.
    pub fn find(&self, provider: Provider, id: Option<&str>) -> Option<&ModelEntry> {
        match (provider, id) {
            (Provider::Claude, None) => self.claude_default.as_ref(),
            (provider, Some(id)) => self.models(provider).iter().find(|model| model.id == id),
            (Provider::Pi, None) => None,
        }
    }

    /// The entry stored under `key` (see [`ModelEntry::key`]).
    pub fn find_key(&self, key: &str) -> Option<&ModelEntry> {
        self.claude.iter().chain(&self.pi).find(|model| model.key() == key)
    }

    pub fn is_favorite(&self, model: &ModelEntry) -> bool {
        self.favorites.contains(&model.key())
    }

    pub fn favorites(&self) -> Vec<&ModelEntry> {
        self.favorites
            .iter()
            .filter_map(|key| self.claude.iter().chain(&self.pi).find(|model| &model.key() == key))
            .collect()
    }

    pub fn toggle_favorite(model: &ModelEntry, cx: &mut App) {
        let key = model.key();
        let catalog = cx.global_mut::<Catalog>();
        match catalog.favorites.iter().position(|favorite| *favorite == key) {
            Some(ix) => {
                catalog.favorites.remove(ix);
            }
            None => catalog.favorites.push(key),
        }
        let favorites = catalog.favorites.clone();
        settings::update(|settings| settings.favorites = favorites);
        cx.refresh_windows();
    }
}

/// Parses the `models` of Claude's `initialize` reply into the default entry
/// and the pickable ones.
fn claude_models_from(info: &Value) -> (Option<ModelEntry>, Vec<ModelEntry>) {
    let mut default = None;
    let mut models = Vec::new();
    for model in info.get("models").and_then(Value::as_array).into_iter().flatten() {
        let Some(id) = model.get("value").and_then(Value::as_str) else { continue };
        let text = |key: &str| model.get(key).and_then(Value::as_str).unwrap_or_default().to_string();
        let flag = |key: &str| model.get(key).and_then(Value::as_bool).unwrap_or(false);
        let display = text("displayName");
        let entry = ModelEntry {
            provider: Provider::Claude,
            id: id.to_string(),
            resolved: Some(text("resolvedModel")).filter(|model| !model.is_empty()).unwrap_or_else(|| id.into()),
            name: if display.starts_with("Claude") { display.clone() } else { format!("Claude {display}") },
            detail: text("description"),
            // Aliases (`opus`, `sonnet`) track the current models; full ids are older ones.
            legacy: id.starts_with("claude-"),
            efforts: model
                .get("supportedEffortLevels")
                .and_then(Value::as_array)
                .map(|levels| levels.iter().filter_map(Value::as_str).filter_map(Effort::from_cli).collect())
                .unwrap_or_default(),
            fast_mode: flag("supportsFastMode"),
            auto_mode: flag("supportsAutoMode"),
            context_window: None,
        };
        if id == "default" {
            default = Some(entry);
        } else {
            models.push(entry);
        }
    }
    // Present the default as the model it resolves to.
    if let Some(default) = default.as_mut()
        && let Some(current) = models.iter().find(|model| !model.legacy && model.resolved == default.resolved)
    {
        default.name = current.name.clone();
    }
    (default, models)
}

/// Starts `pi` just long enough to list the models its providers offer.
fn probe_pi_models() -> Result<Vec<ModelEntry>, PiStatus> {
    let binary = find_pi_binary().ok_or(PiStatus::Missing)?;
    let mut command = Command::new(&binary);
    command
        .args(["--mode", "rpc", "--no-session"])
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::null());
    if let Some(dir) = binary.parent() {
        let path = std::env::var_os("PATH").unwrap_or_default();
        let paths = std::iter::once(dir.to_path_buf()).chain(std::env::split_paths(&path));
        if let Ok(joined) = std::env::join_paths(paths) {
            command.env("PATH", joined);
        }
    }
    let failed = |err: &dyn std::fmt::Display| PiStatus::Failed(err.to_string());
    let mut child = command.spawn().map_err(|err| failed(&err))?;
    let mut stdin = child.stdin.take().ok_or_else(|| failed(&"no stdin"))?;
    let stdout = child.stdout.take().ok_or_else(|| failed(&"no stdout"))?;
    let request = json!({"id": "models", "type": "get_available_models"});
    writeln!(stdin, "{request}").map_err(|err| failed(&err))?;

    let (tx, rx) = mpsc::channel();
    std::thread::spawn(move || {
        for line in BufReader::new(stdout).lines().map_while(Result::ok) {
            if let Ok(value) = serde_json::from_str::<Value>(&line)
                && value.get("id").and_then(Value::as_str) == Some("models")
            {
                let _ = tx.send(value);
                return;
            }
        }
    });
    let reply = rx.recv_timeout(Duration::from_secs(30));
    drop(stdin);
    let _ = child.kill();
    let _ = child.wait();
    let reply = reply.map_err(|_| failed(&"pi didn't list its models"))?;
    if reply.get("success").and_then(Value::as_bool) != Some(true) {
        let error = reply.get("error").and_then(Value::as_str).unwrap_or("unknown error");
        return Err(failed(&error));
    }
    Ok(reply
        .pointer("/data/models")
        .and_then(Value::as_array)
        .into_iter()
        .flatten()
        .filter_map(pi_model_from)
        .collect())
}

pub fn pi_model_from(model: &Value) -> Option<ModelEntry> {
    let id = model.get("id")?.as_str()?;
    let provider = model.get("provider")?.as_str()?;
    let reasoning = model.get("reasoning").and_then(Value::as_bool).unwrap_or(false);
    Some(ModelEntry {
        provider: Provider::Pi,
        id: format!("{provider}/{id}"),
        resolved: id.to_string(),
        name: model.get("name").and_then(Value::as_str).unwrap_or(id).to_string(),
        detail: provider_label(provider),
        legacy: false,
        efforts: if reasoning {
            vec![Effort::Off, Effort::Minimal, Effort::Low, Effort::Medium, Effort::High]
        } else {
            vec![Effort::Off]
        },
        fast_mode: false,
        auto_mode: false,
        context_window: model.get("contextWindow").and_then(Value::as_u64),
    })
}

/// `openai-codex` → `OpenAI Codex`.
fn provider_label(provider: &str) -> String {
    provider
        .split(['-', '_'])
        .map(|word| match word {
            "openai" => "OpenAI".to_string(),
            "openrouter" => "OpenRouter".to_string(),
            "xai" => "xAI".to_string(),
            "ai" => "AI".to_string(),
            word => {
                let mut chars = word.chars();
                chars.next().map(|first| first.to_uppercase().chain(chars).collect()).unwrap_or_default()
            }
        })
        .collect::<Vec<_>>()
        .join(" ")
}

#[cfg(test)]
mod tests {
    use super::{claude_models_from, provider_label};
    use serde_json::json;

    #[::core::prelude::v1::test]
    fn provider_labels() {
        assert_eq!(provider_label("openai-codex"), "OpenAI Codex");
        assert_eq!(provider_label("anthropic"), "Anthropic");
        assert_eq!(provider_label("openrouter"), "OpenRouter");
        assert_eq!(provider_label("google_vertex"), "Google Vertex");
    }

    #[::core::prelude::v1::test]
    fn claude_default_takes_the_current_name() {
        let info = json!({"models": [
            {"value": "default", "resolvedModel": "claude-opus-5-5", "displayName": "Default (recommended)"},
            {"value": "opus", "resolvedModel": "claude-opus-5-5", "displayName": "Opus 5.5",
             "supportedEffortLevels": ["low", "high"], "supportsFastMode": true},
            {"value": "claude-haiku-4-5", "displayName": "Haiku 4.5"},
        ]});
        let (default, models) = claude_models_from(&info);
        assert_eq!(default.unwrap().name, "Claude Opus 5.5");
        assert_eq!(models.len(), 2);
        assert!(models[1].legacy);
        assert_eq!(models[0].efforts.len(), 2);
    }
}