modelforge 0.1.0

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

pub type Result<T> = std::result::Result<T, ModelForgeError>;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ModelProvider {
    Hf,
    Ollama,
    Civitai,
}

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

    pub fn short_label(self) -> &'static str {
        match self {
            Self::Hf => "HF",
            Self::Ollama => "OLM",
            Self::Civitai => "CIV",
        }
    }
}

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

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ModelTarget {
    Org(String),
    Model(String),
    Top,
    Latest,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ModelForgeError {
    InvalidTarget(String),
}

impl ModelForgeError {
    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 {}

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]))),
    }
}

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