modelforge 0.1.7

//! 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 {
    /// Stable provider list in CLI display order.
    pub const ALL: [Self; 3] = [Self::Hf, Self::Ollama, Self::Civitai];

    /// Return all supported providers in stable CLI display order.
    pub fn all() -> &'static [Self] {
        &Self::ALL
    }

    /// 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 namespace, organization, owner, or search keyword targets.
    pub fn is_org(&self) -> bool {
        matches!(self, Self::Org(_))
    }

    /// Return true for concrete model id targets.
    pub fn is_model(&self) -> bool {
        matches!(self, Self::Model(_))
    }

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

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

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

    /// 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()));
        }
    }

    if let ModelProvider::Ollama = provider {
        if let Some(path) = normalized
            .strip_prefix("https://ollama.com/library/")
            .or_else(|| normalized.strip_prefix("http://ollama.com/library/"))
        {
            let id = first_clean_path_segment(path);
            if !id.is_empty() {
                return Ok(ModelTarget::Model(id.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 cleaned = path.split(['?', '#']).next().unwrap_or(path);
    let parts = cleaned
        .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]))),
    }
}

fn first_clean_path_segment(path: &str) -> &str {
    path.split(['/', '?', '#']).next().unwrap_or("").trim()
}

#[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(
                "https://ollama.com/library/llama3.2/tags",
                ModelProvider::Ollama,
            ),
            Ok(ModelTarget::Model("llama3.2".to_string()))
        );
        assert_eq!(
            parse_model_target_for_provider("trending", ModelProvider::Ollama),
            Ok(ModelTarget::Top)
        );
    }

    #[test]
    fn strips_url_query_and_fragment_noise() {
        assert_eq!(
            parse_model_target_for_provider(
                "https://huggingface.co/openai-community/gpt2?library=transformers#files",
                ModelProvider::Hf,
            ),
            Ok(ModelTarget::Model("openai-community/gpt2".to_string()))
        );
        assert_eq!(
            parse_model_target_for_provider(
                "https://ollama.com/library/deepseek-r1?sort=popular",
                ModelProvider::Ollama,
            ),
            Ok(ModelTarget::Model("deepseek-r1".to_string()))
        );
    }

    #[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_model());
        assert!(!model.is_org());
        assert!(!model.is_preset());
        assert_eq!(model.kind(), "model");
        assert_eq!(model.model(), Some("owner/model"));
        assert_eq!(model.org(), None);
        assert_eq!(model.value(), Some("owner/model"));

        let org = ModelTarget::Org("owner".to_string());
        assert!(org.is_org());
        assert!(!org.is_model());
        assert_eq!(org.org(), Some("owner"));
        assert_eq!(org.model(), None);

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

    #[test]
    fn providers_iterate_in_stable_display_order() {
        assert_eq!(ModelProvider::all(), &ModelProvider::ALL);
        assert_eq!(ModelProvider::all()[0], ModelProvider::Hf);
        assert_eq!(ModelProvider::all()[1], ModelProvider::Ollama);
        assert_eq!(ModelProvider::all()[2], ModelProvider::Civitai);
    }
}