granite-cli 0.2.0

use crate::models::huggingface::hf_repo_id;
use crate::models::{ModelFunction, ModelMetadata, ModelVariant};
use crate::providers::base::{
    ApiEndpoint, ApiType, AuthType, HasProviderMetadata, HealthStatus, ModelFormat, Provider,
    ProviderError, ProviderMetadata, ProviderType, http_health_check,
};
use crate::registry::{ConfigConstructable, Secret};
use crate::utils::ui::Ui;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::Duration;

/// Extract the LM Studio model reference from an LM Studio variant URL.
///
/// `https://lmstudio.ai/models/{org}/{model}` -> `{org}/{model}`
/// (e.g. `https://lmstudio.ai/models/ibm/granite-4.1-30b` -> `ibm/granite-4.1-30b`)
///
/// Returns `None` for non-LM Studio URLs.
fn lmstudio_model_ref(url: &str) -> Option<String> {
    Some(
        url.strip_prefix("https://lmstudio.ai/models/")
            .or_else(|| url.strip_prefix("lmstudio.ai/mosels/"))
            .filter(|s| !s.is_empty())?
            .to_string(),
    )
}

/// Response from `POST /api/v1/models/download`.
#[derive(Debug, Deserialize)]
struct LMStudioDownloadResponse {
    job_id: Option<String>,
    status: String,
    total_size_bytes: Option<u64>,
}

/// Response from `GET /api/v1/models/download/status/:job_id`.
#[derive(Debug, Deserialize)]
struct LMStudioJobStatus {
    status: String,
    downloaded_bytes: Option<u64>,
    total_size_bytes: Option<u64>,
    error: Option<String>,
}

/*-- LM Studio Provider Configuration ----------------------------------------*/

#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
pub struct LMStudioProviderConfig {
    /// Base URL for the LM Studio server
    #[serde(default = "default_lmstudio_url")]
    pub base_url: String,

    /// API key for authentication (optional)
    pub api_key: Option<Secret>,

    /// Timeout for health checks in seconds
    #[serde(default = "default_timeout")]
    pub timeout_secs: u64,

    /// Whether to verify SSL certificates
    #[serde(default = "default_verify_ssl")]
    pub verify_ssl: bool,

    /// Endpoint to use for health checks
    #[serde(default = "default_lmstudio_health_endpoint")]
    pub health_check_endpoint: String,
}

fn default_lmstudio_url() -> String {
    "http://localhost:1234".to_string()
}

fn default_timeout() -> u64 {
    10
}

fn default_verify_ssl() -> bool {
    true
}

fn default_lmstudio_health_endpoint() -> String {
    "/v1/models".to_string()
}

impl Default for LMStudioProviderConfig {
    fn default() -> Self {
        Self {
            base_url: default_lmstudio_url(),
            api_key: None,
            timeout_secs: default_timeout(),
            verify_ssl: default_verify_ssl(),
            health_check_endpoint: default_lmstudio_health_endpoint(),
        }
    }
}

/*-- LM Studio Provider Implementation ---------------------------------------*/

pub struct LMStudioProvider {
    instance_id: String,
    config: LMStudioProviderConfig,
    client: reqwest::Client,
}

impl LMStudioProvider {
    fn default_function_endpoints() -> HashMap<ModelFunction, Vec<ApiEndpoint>> {
        let mut map = HashMap::new();

        map.insert(
            ModelFunction::Chat,
            vec![ApiEndpoint::OpenAIChat, ApiEndpoint::AnthropicMessages],
        );

        map.insert(
            ModelFunction::ToolCalling,
            vec![ApiEndpoint::OpenAIChat, ApiEndpoint::AnthropicMessages],
        );

        map.insert(
            ModelFunction::Thinking,
            vec![ApiEndpoint::OpenAIChat, ApiEndpoint::AnthropicMessages],
        );

        map.insert(
            ModelFunction::ImageUnderstanding,
            vec![ApiEndpoint::OpenAIChat, ApiEndpoint::AnthropicMessages],
        );

        map.insert(
            ModelFunction::Guardian,
            vec![ApiEndpoint::OpenAIChat, ApiEndpoint::AnthropicMessages],
        );

        map.insert(
            ModelFunction::Embeddings,
            vec![ApiEndpoint::OpenAIEmbeddings],
        );

        map
    }

    fn default_formats() -> Vec<ModelFormat> {
        let mut formats = vec![ModelFormat::GGUF, ModelFormat::LMStudio];

        if cfg!(target_os = "macos") {
            formats.push(ModelFormat::MLX);
        }

        formats
    }
}

impl ConfigConstructable for LMStudioProvider {
    type Config = LMStudioProviderConfig;

    fn new(
        instance_id: &str,
        cfg: &serde_json::Value,
        _global_config: &crate::config::Config,
    ) -> Self {
        let config: LMStudioProviderConfig =
            serde_json::from_value(cfg.clone()).unwrap_or_default();

        let client = reqwest::Client::builder()
            .timeout(Duration::from_secs(config.timeout_secs))
            .danger_accept_invalid_certs(!config.verify_ssl)
            .build()
            .expect("Failed to create HTTP client");

        Self {
            instance_id: instance_id.to_string(),
            config,
            client,
        }
    }
}

impl crate::registry::Named for LMStudioProvider {
    fn instance_id(&self) -> &str {
        &self.instance_id
    }
}

#[async_trait]
impl Provider for LMStudioProvider {
    fn name(&self) -> &str {
        "LM Studio"
    }

    fn function_endpoints(&self) -> HashMap<ModelFunction, Vec<ApiEndpoint>> {
        Self::default_function_endpoints()
    }

    fn supported_api_types(&self) -> Vec<ApiType> {
        vec![ApiType::OpenAI, ApiType::Anthropic]
    }

    fn base_url(&self) -> &str {
        &self.config.base_url
    }

    fn api_key(&self) -> Option<&Secret> {
        self.config.api_key.as_ref()
    }

    fn verify_ssl(&self) -> bool {
        self.config.verify_ssl
    }

    fn supported_formats(&self) -> Vec<ModelFormat> {
        Self::default_formats()
    }

    fn can_run_model(&self, variant_format: &str, _variant_precision: &str) -> bool {
        let format = variant_format.to_lowercase();
        matches!(format.as_str(), "gguf" | "lmstudio" | "mlx" if format != "mlx" || cfg!(target_os = "macos"))
    }

    async fn health_check(&self) -> Result<HealthStatus, ProviderError> {
        http_health_check(
            &self.client,
            &self.config.base_url,
            &self.config.health_check_endpoint,
            self.config.api_key.as_ref(),
        )
        .await
    }

    async fn pull_model(
        &self,
        model: &ModelMetadata,
        variant: &ModelVariant,
        ui: &dyn Ui,
    ) -> Result<crate::providers::PullResult, ProviderError> {
        let model_ref = if let Some(ref_str) = lmstudio_model_ref(&variant.url) {
            ref_str
        } else if let Some(repo) = hf_repo_id(&variant.url) {
            format!("https://huggingface.co/{repo}")
        } else {
            return Err(ProviderError::Other(format!(
                "cannot determine a model reference for {} variant {}/{}",
                model.family, variant.format, variant.precision
            )));
        };
        let label = format!(
            "{} ({} {})",
            model.family, variant.format, variant.precision
        );

        let url = format!("{}/api/v1/models/download", self.config.base_url);
        let mut request = self.client.post(&url).json(&serde_json::json!({
            "model": model_ref,
            "quantization": variant.precision,
        }));
        if let Some(key) = &self.config.api_key {
            request = request.bearer_auth(&key.0);
        }

        let response = request.send().await?;
        if !response.status().is_success() {
            let status = response.status();
            let body = response.text().await.unwrap_or_default();
            return Err(ProviderError::Other(format!(
                "LM Studio download request failed ({status}): {body}"
            )));
        }
        let started: LMStudioDownloadResponse = response.json().await?;

        let handle = ui.pull_start(&label, started.total_size_bytes);

        if started.status == "already_downloaded" {
            ui.pull_finish(handle, &label, None);
            return Ok(crate::providers::PullResult::Success);
        }

        let job_id = match started.job_id {
            Some(id) => id,
            None => {
                ui.pull_finish(handle, &label, None);
                return Ok(crate::providers::PullResult::Success);
            }
        };
        let status_url = format!(
            "{}/api/v1/models/download/status/{}",
            self.config.base_url, job_id
        );

        loop {
            tokio::time::sleep(Duration::from_secs(1)).await;

            let mut status_request = self.client.get(&status_url);
            if let Some(key) = &self.config.api_key {
                status_request = status_request.bearer_auth(&key.0);
            }

            let status_response = status_request.send().await?;
            if !status_response.status().is_success() {
                let status = status_response.status();
                let body = status_response.text().await.unwrap_or_default();
                let err = format!("LM Studio status check failed ({status}): {body}");
                ui.pull_finish(handle, &label, Some(&err));
                return Err(ProviderError::Other(err));
            }
            let job: LMStudioJobStatus = status_response.json().await?;

            ui.pull_progress(
                handle,
                job.downloaded_bytes.unwrap_or(0),
                job.total_size_bytes.or(started.total_size_bytes),
            );

            match job.status.as_str() {
                "completed" => {
                    ui.pull_finish(handle, &label, None);
                    return Ok(crate::providers::PullResult::Success);
                }
                "failed" => {
                    let err = job.error.unwrap_or_else(|| "download failed".to_string());
                    ui.pull_finish(handle, &label, Some(&err));
                    return Err(ProviderError::Other(err));
                }
                _ => continue,
            }
        }
    }
}

impl HasProviderMetadata for LMStudioProvider {
    fn metadata() -> ProviderMetadata {
        let mut formats = vec![ModelFormat::GGUF, ModelFormat::LMStudio];
        let mut tags = vec![
            "lm-studio".to_string(),
            "local".to_string(),
            "gguf".to_string(),
        ];

        if cfg!(target_os = "macos") {
            formats.push(ModelFormat::MLX);
            tags.push("mlx".to_string());
            tags.push("apple-silicon".to_string());
        }

        ProviderMetadata {
            name: "LM Studio".to_string(),
            description:
                "User-friendly local inference server with GUI, supporting GGUF and MLX models"
                    .to_string(),
            provider_type: ProviderType::Local,
            default_endpoint: "http://localhost:1234".to_string(),
            supported_api_types: vec![ApiType::OpenAI, ApiType::Anthropic],
            default_function_endpoints: Self::default_function_endpoints(),
            supported_formats: formats,
            authentication: vec![AuthType::None, AuthType::BearerToken],
            tags,
        }
    }
}

/*-- tests -------------------------------------------------------------------*/

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

    #[test]
    fn test_default_config() {
        let config = LMStudioProviderConfig::default();
        assert_eq!(config.base_url, "http://localhost:1234");
        assert!(config.api_key.is_none());
        assert_eq!(config.timeout_secs, 10);
        assert!(config.verify_ssl);
        assert_eq!(config.health_check_endpoint, "/v1/models");
    }

    #[test]
    fn test_provider_metadata() {
        let meta = LMStudioProvider::metadata();
        assert_eq!(meta.name, "LM Studio");
        assert!(meta.supported_api_types.contains(&ApiType::OpenAI));
        assert!(meta.supported_api_types.contains(&ApiType::Anthropic));
        assert!(
            meta.default_function_endpoints
                .contains_key(&ModelFunction::Chat)
        );
    }

    #[test]
    fn test_provider_constructs_from_json() {
        let cfg = serde_json::json!({
            "base_url": "http://example.com:5678",
            "timeout_secs": 30
        });
        let provider =
            LMStudioProvider::new("my-lmstudio", &cfg, &crate::config::Config::default());
        assert_eq!(provider.config.base_url, "http://example.com:5678");
        assert_eq!(provider.config.timeout_secs, 30);
    }

    #[test]
    fn test_can_run_model_accepts_gguf() {
        let provider = LMStudioProvider::new(
            "my-lmstudio",
            &serde_json::json!({}),
            &crate::config::Config::default(),
        );
        assert!(provider.can_run_model("gguf", "Q4_K_M"));
        assert!(provider.can_run_model("GGUF", "fp16"));
    }

    #[test]
    fn test_can_run_model_rejects_non_supported() {
        let provider = LMStudioProvider::new(
            "my-lmstudio",
            &serde_json::json!({}),
            &crate::config::Config::default(),
        );
        assert!(!provider.can_run_model("safetensors", "fp16"));
        assert!(!provider.can_run_model("onnx", "fp32"));
    }

    #[test]
    fn test_mlx_formats_on_macos() {
        let provider = LMStudioProvider::new(
            "my-lmstudio",
            &serde_json::json!({}),
            &crate::config::Config::default(),
        );
        let formats = provider.supported_formats();

        #[cfg(target_os = "macos")]
        assert!(formats.contains(&ModelFormat::MLX));

        #[cfg(not(target_os = "macos"))]
        assert!(!formats.contains(&ModelFormat::MLX));
    }

    #[test]
    fn test_can_run_mlx_model() {
        let provider = LMStudioProvider::new(
            "my-lmstudio",
            &serde_json::json!({}),
            &crate::config::Config::default(),
        );

        #[cfg(target_os = "macos")]
        assert!(provider.can_run_model("mlx", "fp16"));

        #[cfg(not(target_os = "macos"))]
        assert!(!provider.can_run_model("mlx", "fp16"));
    }

    #[test]
    fn test_lmstudio_download_response_parses() {
        let body = r#"{"job_id":"abc123","status":"downloading","total_size_bytes":1000}"#;
        let resp: LMStudioDownloadResponse = serde_json::from_str(body).unwrap();
        assert_eq!(resp.job_id, Some("abc123".to_string()));
        assert_eq!(resp.total_size_bytes, Some(1000));
    }

    #[test]
    fn test_lmstudio_job_status_parses() {
        let body = r#"{"status":"completed","downloaded_bytes":1000,"total_size_bytes":1000}"#;
        let job: LMStudioJobStatus = serde_json::from_str(body).unwrap();
        assert_eq!(job.status, "completed");
        assert_eq!(job.downloaded_bytes, Some(1000));
    }

    #[test]
    fn test_lmstudio_model_ref_from_simple_url() {
        assert_eq!(
            lmstudio_model_ref("https://lmstudio.ai/models/ibm/granite-4.1-30b"),
            Some("ibm/granite-4.1-30b".to_string())
        );
    }

    #[test]
    fn test_lmstudio_model_ref_rejects_non_lmstudio_url() {
        assert_eq!(
            lmstudio_model_ref("https://huggingface.co/ibm-granite/granite-speech-4.1-2b"),
            None
        );
        assert_eq!(
            lmstudio_model_ref("https://ollama.com/library/granite4:1b"),
            None
        );
        assert_eq!(lmstudio_model_ref("https://lmstudio.ai/models/"), None);
    }

    #[test]
    fn test_pull_model_uses_lmstudio_ref() {
        // Tests that lmstudio_model_ref extracts the correct model reference
        // from LM Studio variant URLs, and that hf_repo_id falls back for
        // non-LMStudio URLs — the model_ref selection logic in pull_model.
        assert_eq!(
            lmstudio_model_ref("https://lmstudio.ai/models/ibm/granite-4.1-30b"),
            Some("ibm/granite-4.1-30b".to_string())
        );
        assert_eq!(
            lmstudio_model_ref("https://huggingface.co/ibm-granite/granite-speech-4.1-2b"),
            None
        );
        assert_eq!(
            hf_repo_id("https://huggingface.co/ibm-granite/granite-speech-4.1-2b"),
            Some("ibm-granite/granite-speech-4.1-2b")
        );
    }
}