modelforge 0.1.4

//! Model provider and target parsing primitives for `thesa`.
//!
//! Made by Trevor Knott for Knott Dynamics.
//!
//! ModelForge normalizes CLI/user input for Hugging Face, Ollama, and CivitAI
//! archive targets. It deliberately does not call provider APIs or download
//! models; callers use it to turn user-facing strings into stable target enums.
//! The crate keeps parsing deterministic and side-effect free.
//!
//! ```
//! assert_eq!(
//!     modelforge::parse_model_target_for_provider(
//!         "https://huggingface.co/openai-community/gpt2",
//!         modelforge::ModelProvider::Hf,
//!     ),
//!     Ok(modelforge::ModelTarget::Model("openai-community/gpt2".to_string()))
//! );
//! assert_eq!(
//!     modelforge::parse_model_target_for_provider("newest", modelforge::ModelProvider::Civitai),
//!     Ok(modelforge::ModelTarget::Latest)
//! );
//! ```

use std::error::Error;
use std::fmt;

/// Result type used by ModelForge APIs.
pub type Result<T> = std::result::Result<T, ModelForgeError>;

/// Supported model/archive providers.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ModelProvider {
    /// Hugging Face Hub.
    Hf,
    /// Local Ollama registry/runtime.
    Ollama,
    /// CivitAI model catalog.
    Civitai,
}

impl ModelProvider {
    /// Human-readable provider label.
    pub fn label(self) -> &'static str {
        match self {
            Self::Hf => "Hugging Face",
            Self::Ollama => "Ollama",
            Self::Civitai => "CivitAI",
        }
    }

    /// Short UI label used by compact archive displays.
    pub fn short_label(self) -> &'static str {
        match self {
            Self::Hf => "HF",
            Self::Ollama => "OLM",
            Self::Civitai => "CIV",
        }
    }

    /// Stable lowercase provider slug.
    pub fn slug(self) -> &'static str {
        match self {
            Self::Hf => "hf",
            Self::Ollama => "ollama",
            Self::Civitai => "civitai",
        }
    }

    /// Parse a provider slug accepted by thesa-style CLIs.
    pub fn from_slug(value: &str) -> Option<Self> {
        match value.trim().to_ascii_lowercase().as_str() {
            "hf" | "huggingface" | "hugging-face" => Some(Self::Hf),
            "ollama" | "olm" => Some(Self::Ollama),
            "civitai" | "civit" | "civ" => Some(Self::Civitai),
            _ => None,
        }
    }
}

impl fmt::Display for ModelProvider {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Hf | Self::Ollama | Self::Civitai => write!(f, "{}", self.slug()),
        }
    }
}

/// Normalized archive target requested by a user or CLI.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ModelTarget {
    /// Provider namespace, organization, owner, or search keyword.
    Org(String),
    /// Concrete model id.
    Model(String),
    /// Provider-specific top/trending preset.
    Top,
    /// Provider-specific latest/newest preset.
    Latest,
}

impl ModelTarget {
    /// Return true for `Top` or `Latest` preset targets.
    pub fn is_preset(&self) -> bool {
        matches!(self, Self::Top | Self::Latest)
    }

    /// Stable target kind label.
    pub fn kind(&self) -> &'static str {
        match self {
            Self::Org(_) => "org",
            Self::Model(_) => "model",
            Self::Top => "top",
            Self::Latest => "latest",
        }
    }

    /// Borrow the contained target id for `Org` or `Model` targets.
    pub fn value(&self) -> Option<&str> {
        match self {
            Self::Org(value) | Self::Model(value) => Some(value),
            Self::Top | Self::Latest => None,
        }
    }
}

/// Errors returned while parsing model targets.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ModelForgeError {
    /// The provided target string could not be normalized.
    InvalidTarget(String),
}

impl ModelForgeError {
    /// The invalid target payload or reason string.
    pub fn target(&self) -> &str {
        match self {
            Self::InvalidTarget(target) => target,
        }
    }
}

impl fmt::Display for ModelForgeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InvalidTarget(target) => write!(f, "invalid model target '{target}'"),
        }
    }
}

impl Error for ModelForgeError {}

/// Parse a Hugging Face-style target into `Org` or `Model`.
///
/// This is provider-neutral only for simple `owner` and `owner/model` forms.
/// Prefer `parse_model_target_for_provider` when the provider is known.
pub fn parse_model_target(input: &str) -> Result<ModelTarget> {
    let cleaned = input.trim();
    if cleaned.is_empty() {
        return Err(ModelForgeError::InvalidTarget(
            "empty model target".to_string(),
        ));
    }

    if let Some(path) = cleaned.strip_prefix("https://huggingface.co/") {
        return parse_model_path(path);
    }
    if let Some(path) = cleaned.strip_prefix("http://huggingface.co/") {
        return parse_model_path(path);
    }

    let parts = cleaned
        .split('/')
        .filter(|part| !part.is_empty())
        .collect::<Vec<_>>();
    match parts.len() {
        0 => Err(ModelForgeError::InvalidTarget(cleaned.to_string())),
        1 => Ok(ModelTarget::Org(parts[0].to_string())),
        _ => Ok(ModelTarget::Model(format!("{}/{}", parts[0], parts[1]))),
    }
}

/// Parse a model target using provider-specific rules.
///
/// Accepted presets for all providers are `top`, `trending`, `popular`,
/// `latest`, `newest`, and `new`.
///
/// CivitAI additionally accepts numeric ids and `https://civitai.com/models/<id>`
/// URLs as concrete model targets.
pub fn parse_model_target_for_provider(
    input: &str,
    provider: ModelProvider,
) -> Result<ModelTarget> {
    let normalized = input.trim();
    let canonical = normalized.to_lowercase();

    match canonical.as_str() {
        "top" | "trending" | "popular" => return Ok(ModelTarget::Top),
        "latest" | "newest" | "new" => return Ok(ModelTarget::Latest),
        _ => {}
    }

    if normalized.is_empty() {
        return Err(ModelForgeError::InvalidTarget(
            "empty model target".to_string(),
        ));
    }

    if let ModelProvider::Civitai = provider {
        if let Some(path) = normalized
            .strip_prefix("https://civitai.com/models/")
            .or_else(|| normalized.strip_prefix("http://civitai.com/models/"))
        {
            let id = path.split(['/', '?', '#']).next().unwrap_or("").trim();
            if !id.is_empty() && id.chars().all(|c| c.is_ascii_digit()) {
                return Ok(ModelTarget::Model(id.to_string()));
            }
        }

        if normalized.chars().all(|c| c.is_ascii_digit()) {
            return Ok(ModelTarget::Model(normalized.to_string()));
        }
    }

    let model_target = parse_model_target(normalized)?;
    match provider {
        ModelProvider::Civitai => match model_target {
            ModelTarget::Model(model) if model.chars().all(|c| c.is_ascii_digit()) => {
                Ok(ModelTarget::Model(model))
            }
            other => Ok(other),
        },
        _ => Ok(model_target),
    }
}

fn parse_model_path(path: &str) -> Result<ModelTarget> {
    let parts = path
        .trim_matches('/')
        .split('/')
        .filter(|part| !part.is_empty())
        .collect::<Vec<_>>();
    match parts.len() {
        0 => Err(ModelForgeError::InvalidTarget(path.to_string())),
        1 => Ok(ModelTarget::Org(parts[0].to_string())),
        _ => Ok(ModelTarget::Model(format!("{}/{}", parts[0], parts[1]))),
    }
}

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

    #[test]
    fn parses_hugging_face_targets() {
        assert_eq!(
            parse_model_target_for_provider("meta-llama/Llama-2-7b", ModelProvider::Hf),
            Ok(ModelTarget::Model("meta-llama/Llama-2-7b".to_string()))
        );
        assert_eq!(
            parse_model_target_for_provider("top", ModelProvider::Hf),
            Ok(ModelTarget::Top)
        );
        assert_eq!(
            parse_model_target_for_provider(
                "https://huggingface.co/openai-community/gpt2",
                ModelProvider::Hf,
            ),
            Ok(ModelTarget::Model("openai-community/gpt2".to_string()))
        );
    }

    #[test]
    fn parses_civitai_targets() {
        assert_eq!(
            parse_model_target_for_provider("827184", ModelProvider::Civitai),
            Ok(ModelTarget::Model("827184".to_string()))
        );
        assert_eq!(
            parse_model_target_for_provider("newest", ModelProvider::Civitai),
            Ok(ModelTarget::Latest)
        );
        assert_eq!(
            parse_model_target_for_provider(
                "https://civitai.com/models/827184",
                ModelProvider::Civitai,
            ),
            Ok(ModelTarget::Model("827184".to_string()))
        );
        assert_eq!(
            parse_model_target_for_provider("realistic-vision", ModelProvider::Civitai),
            Ok(ModelTarget::Org("realistic-vision".to_string()))
        );
    }

    #[test]
    fn parses_ollama_targets() {
        assert_eq!(
            parse_model_target_for_provider("llama3", ModelProvider::Ollama),
            Ok(ModelTarget::Org("llama3".to_string()))
        );
        assert_eq!(
            parse_model_target_for_provider("trending", ModelProvider::Ollama),
            Ok(ModelTarget::Top)
        );
    }

    #[test]
    fn provider_slugs_round_trip() {
        assert_eq!(ModelProvider::from_slug("hf"), Some(ModelProvider::Hf));
        assert_eq!(
            ModelProvider::from_slug("hugging-face"),
            Some(ModelProvider::Hf)
        );
        assert_eq!(ModelProvider::from_slug("olm"), Some(ModelProvider::Ollama));
        assert_eq!(
            ModelProvider::from_slug("civ"),
            Some(ModelProvider::Civitai)
        );
        assert_eq!(ModelProvider::from_slug("unknown"), None);
        assert_eq!(ModelProvider::Civitai.to_string(), "civitai");
    }

    #[test]
    fn target_helpers_describe_targets() {
        let model = ModelTarget::Model("owner/model".to_string());
        assert!(!model.is_preset());
        assert_eq!(model.kind(), "model");
        assert_eq!(model.value(), Some("owner/model"));

        let latest = ModelTarget::Latest;
        assert!(latest.is_preset());
        assert_eq!(latest.kind(), "latest");
        assert_eq!(latest.value(), None);
    }
}