mod transfer;

pub(crate) use transfer::{
    delete_managed_artifact, download_managed_artifact, validate_managed_artifact,
};

use serde::{Deserialize, Serialize};
use std::fmt;
use std::fs;
use std::path::{Path, PathBuf};

pub(crate) const MODEL_CHOICES: [ModelChoice; 4] = [
    ModelChoice::DeepSeekV4Flash0731,
    ModelChoice::DeepSeekV4Flash,
    ModelChoice::DeepSeekV4Pro,
    ModelChoice::Glm52,
];
pub(crate) const MANAGED_ARTIFACTS: [ManagedArtifactId; 7] = [
    ManagedArtifactId::DeepSeekV4Flash0731,
    ManagedArtifactId::DeepSeekV4Flash0731Dspark,
    ManagedArtifactId::DeepSeekV4Flash,
    ManagedArtifactId::DeepSeekV4FlashMtp,
    ManagedArtifactId::DeepSeekV4FlashDspark,
    ManagedArtifactId::DeepSeekV4Pro,
    ManagedArtifactId::Glm52,
];

const DEEPSEEK_REPOSITORY: &str = "antirez/deepseek-v4-gguf";
const GLM_REPOSITORY: &str = "antirez/glm-5.2-gguf";

const FLASH: Artifact = Artifact {
    label: "DeepSeek V4 Flash model",
    file_name: "DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf",
    repository: DEEPSEEK_REPOSITORY,
    size: 86_720_111_488,
    sha256: "efc7ed607ff27076e3e501fc3fefefa33c0ed8cf1eff483a2b7fdc0c2e616668",
    support: Some(false),
};
const FLASH_DSPARK: Artifact = Artifact {
    label: "DSpark support",
    file_name: "DeepSeek-V4-Flash-DSpark-support.gguf",
    repository: DEEPSEEK_REPOSITORY,
    size: 5_989_114_272,
    sha256: "8b3adf5942bec22ae2ea867cd7079cf13530ba83ffcffaf00f5de48664a1a34e",
    support: Some(true),
};
const FLASH_MTP: Artifact = Artifact {
    label: "Legacy MTP support",
    file_name: "DeepSeek-V4-Flash-MTP-Q4K-Q8_0-F32.gguf",
    repository: DEEPSEEK_REPOSITORY,
    size: 3_807_602_400,
    sha256: "afd481ee689dce9037f70f39085fcdae5a5b096d521cdad43b19fa52bf8f4083",
    support: Some(true),
};
const FLASH_0731: Artifact = Artifact {
    label: "DeepSeek V4 Flash 0731 model",
    file_name: "DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-0731.gguf",
    repository: DEEPSEEK_REPOSITORY,
    size: 86_720_111_488,
    sha256: "ca22ae2f838e14077c22bc1c1417b71b45b5e5a3687bd96c2ac6e17fdb6261c0",
    support: Some(false),
};
const FLASH_0731_DSPARK: Artifact = Artifact {
    label: "DeepSeek V4 Flash 0731 DSpark support",
    file_name: "DeepSeek-V4-Flash-DSpark-support-0731.gguf",
    repository: DEEPSEEK_REPOSITORY,
    size: 5_989_114_272,
    sha256: "7e319924541db3f7a163ed7e11d7532a70d48228ab59d36cb81e1d4511885360",
    support: Some(true),
};
const PRO: Artifact = Artifact {
    label: "DeepSeek V4 Pro 0813 model",
    file_name: "DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix-0813.gguf",
    repository: DEEPSEEK_REPOSITORY,
    size: 464_627_334_560,
    sha256: "c4d997ab9894b6c78b759f7869fe1726b6314b6515f6ff82607df3797c5eb193",
    support: Some(false),
};
const GLM: Artifact = Artifact {
    label: "GLM 5.2 model",
    file_name: "GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf",
    repository: GLM_REPOSITORY,
    size: 211_075_856_448,
    sha256: "a49de64c5020432bdae23de36a423a9660a5621bc0db8d12b66bd8814b07fea0",
    support: Some(false),
};

#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
pub(crate) enum ModelChoice {
    #[serde(rename = "deepseek-v4-flash")]
    DeepSeekV4Flash,
    #[default]
    #[serde(rename = "deepseek-v4-flash-0731")]
    DeepSeekV4Flash0731,
    #[serde(rename = "deepseek-v4-pro")]
    DeepSeekV4Pro,
    #[serde(rename = "glm-5.2")]
    Glm52,
}

impl ModelChoice {
    pub(crate) fn id(self) -> &'static str {
        match self {
            Self::DeepSeekV4Flash => "deepseek-v4-flash",
            Self::DeepSeekV4Flash0731 => "deepseek-v4-flash-0731",
            Self::DeepSeekV4Pro => "deepseek-v4-pro",
            Self::Glm52 => "glm-5.2",
        }
    }

    pub(crate) fn from_id(id: &str) -> Option<Self> {
        MODEL_CHOICES.into_iter().find(|model| model.id() == id)
    }

    pub(crate) fn supports_legacy_mtp(self) -> bool {
        self == Self::DeepSeekV4Flash
    }

    pub(crate) fn supports_dspark(self) -> bool {
        matches!(self, Self::DeepSeekV4Flash | Self::DeepSeekV4Flash0731)
    }

    fn main_artifact(self) -> &'static Artifact {
        match self {
            Self::DeepSeekV4Flash => &FLASH,
            Self::DeepSeekV4Flash0731 => &FLASH_0731,
            Self::DeepSeekV4Pro => &PRO,
            Self::Glm52 => &GLM,
        }
    }

    fn dspark_artifact(self) -> Option<&'static Artifact> {
        match self {
            Self::DeepSeekV4Flash => Some(&FLASH_DSPARK),
            Self::DeepSeekV4Flash0731 => Some(&FLASH_0731_DSPARK),
            Self::DeepSeekV4Pro | Self::Glm52 => None,
        }
    }

    #[cfg(test)]
    fn artifacts(
        self,
        legacy_mtp_enabled: bool,
        dspark_enabled: bool,
    ) -> impl Iterator<Item = &'static Artifact> {
        [
            Some(self.main_artifact()),
            (self.supports_legacy_mtp() && legacy_mtp_enabled).then_some(&FLASH_MTP),
            dspark_enabled.then(|| self.dspark_artifact()).flatten(),
        ]
        .into_iter()
        .flatten()
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) struct EngineArtifacts {
    pub(crate) model: PathBuf,
    pub(crate) mtp: Option<PathBuf>,
}

pub(crate) fn engine_artifacts(
    model: ModelChoice,
    legacy_mtp_enabled: bool,
    dspark_enabled: bool,
    models_path: &Path,
) -> EngineArtifacts {
    EngineArtifacts {
        model: model.main_artifact().path(model, models_path),
        mtp: if model.supports_legacy_mtp() && legacy_mtp_enabled {
            Some(FLASH_MTP.path(model, models_path))
        } else if dspark_enabled {
            model
                .dspark_artifact()
                .map(|artifact| artifact.path(model, models_path))
        } else {
            None
        },
    }
}

pub(crate) fn validate_engine_artifacts(
    model: ModelChoice,
    legacy_mtp_enabled: bool,
    dspark_enabled: bool,
    artifacts: &EngineArtifacts,
) -> Result<(), String> {
    if legacy_mtp_enabled && dspark_enabled {
        return Err("Legacy MTP and DSpark cannot be enabled together".into());
    }
    if legacy_mtp_enabled && !model.supports_legacy_mtp() {
        return Err(format!("Legacy MTP is not compatible with {model}"));
    }
    if dspark_enabled && !model.supports_dspark() {
        return Err(format!("DSpark is not compatible with {model}"));
    }
    model
        .main_artifact()
        .validate_installed_path(&artifacts.model)?;
    let expected_support = if legacy_mtp_enabled {
        model.supports_legacy_mtp().then_some(&FLASH_MTP)
    } else if dspark_enabled {
        model.dspark_artifact()
    } else {
        None
    };
    match (expected_support, artifacts.mtp.as_deref()) {
        (Some(expected), Some(path)) => expected.validate_installed_path(path),
        (None, None) => Ok(()),
        (Some(_), None) => Err(format!("{model} is missing its required support GGUF")),
        (None, Some(path)) => Err(format!(
            "{} is not compatible with the selected {model} checkpoint",
            path.display()
        )),
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum ManagedArtifactId {
    DeepSeekV4Flash,
    DeepSeekV4FlashMtp,
    DeepSeekV4FlashDspark,
    DeepSeekV4Flash0731,
    DeepSeekV4Flash0731Dspark,
    DeepSeekV4Pro,
    Glm52,
}

impl ManagedArtifactId {
    pub(crate) fn model(self) -> ModelChoice {
        match self {
            Self::DeepSeekV4Flash | Self::DeepSeekV4FlashMtp | Self::DeepSeekV4FlashDspark => {
                ModelChoice::DeepSeekV4Flash
            }
            Self::DeepSeekV4Flash0731 | Self::DeepSeekV4Flash0731Dspark => {
                ModelChoice::DeepSeekV4Flash0731
            }
            Self::DeepSeekV4Pro => ModelChoice::DeepSeekV4Pro,
            Self::Glm52 => ModelChoice::Glm52,
        }
    }

    fn artifact(self) -> &'static Artifact {
        match self {
            Self::DeepSeekV4Flash => &FLASH,
            Self::DeepSeekV4FlashMtp => &FLASH_MTP,
            Self::DeepSeekV4FlashDspark => &FLASH_DSPARK,
            Self::DeepSeekV4Flash0731 => &FLASH_0731,
            Self::DeepSeekV4Flash0731Dspark => &FLASH_0731_DSPARK,
            Self::DeepSeekV4Pro => &PRO,
            Self::Glm52 => &GLM,
        }
    }
}

impl fmt::Display for ManagedArtifactId {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(self.artifact().label)
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum ManagedArtifactState {
    Missing,
    Partial,
    NeedsVerification,
    Ready,
}

#[derive(Clone, Debug)]
pub(crate) struct ManagedArtifact {
    pub(crate) id: ManagedArtifactId,
    pub(crate) stored: u64,
    pub(crate) expected: u64,
    pub(crate) state: ManagedArtifactState,
}

impl ManagedArtifact {
    pub(crate) fn can_validate(&self) -> bool {
        matches!(
            self.state,
            ManagedArtifactState::NeedsVerification | ManagedArtifactState::Ready
        )
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum DownloadPhase {
    Pending(&'static str),
    Downloading(&'static str),
    Verifying(&'static str),
    Complete,
}

#[derive(Clone, Debug)]
pub(crate) struct DownloadProgress {
    pub(crate) downloaded: u64,
    pub(crate) total: u64,
    pub(crate) phase: DownloadPhase,
    pub(crate) verification: Option<VerificationProgress>,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) struct VerificationProgress {
    pub(crate) verified: u64,
    pub(crate) total: u64,
}

impl DownloadProgress {
    pub(crate) fn remaining(&self) -> u64 {
        self.active_total().saturating_sub(self.completed())
    }

    pub(crate) fn fraction(&self) -> f32 {
        let total = self.active_total();
        if total == 0 {
            0.0
        } else {
            self.completed() as f32 / total as f32
        }
    }

    pub(crate) fn completed(&self) -> u64 {
        self.verification
            .map_or(self.downloaded, |progress| progress.verified)
    }

    pub(crate) fn active_total(&self) -> u64 {
        self.verification
            .map_or(self.total, |progress| progress.total)
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum DownloadOutcome {
    Complete,
    Stopped,
}

impl fmt::Display for ModelChoice {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            Self::DeepSeekV4Flash => "DeepSeek V4 Flash (deprecated preview)",
            Self::DeepSeekV4Flash0731 => "DeepSeek V4 Flash 0731",
            Self::DeepSeekV4Pro => "DeepSeek V4 Pro 0813",
            Self::Glm52 => "GLM 5.2",
        })
    }
}

struct Artifact {
    label: &'static str,
    file_name: &'static str,
    repository: &'static str,
    size: u64,
    sha256: &'static str,
    support: Option<bool>,
}

impl Artifact {
    fn validate_installed_path(&self, path: &Path) -> Result<(), String> {
        if path.file_name().and_then(|name| name.to_str()) != Some(self.file_name) {
            return Err(format!(
                "{} is not the expected {} artifact",
                path.display(),
                self.label
            ));
        }
        let size = path
            .metadata()
            .map_err(|error| format!("Could not inspect {}: {error}", path.display()))?
            .len();
        if size != self.size {
            return Err(format!(
                "{} has {size} bytes, expected {}",
                path.display(),
                self.size
            ));
        }
        let verification = fs::read_to_string(path.with_extension("gguf.sha256"))
            .map_err(|_| format!("{} has not passed checksum verification", path.display()))?;
        if verification.trim() != self.sha256 {
            return Err(format!(
                "{} has the wrong checkpoint identity",
                path.display()
            ));
        }
        Ok(())
    }

    fn path(&self, model: ModelChoice, models_path: &Path) -> PathBuf {
        models_path.join(model.id()).join(self.file_name)
    }

    fn url(&self) -> String {
        format!(
            "https://huggingface.co/{}/resolve/main/{}",
            self.repository, self.file_name
        )
    }

    fn verification_path(&self, model: ModelChoice, models_path: &Path) -> PathBuf {
        self.path(model, models_path).with_extension("gguf.sha256")
    }

    fn partial_path(&self, model: ModelChoice, models_path: &Path) -> PathBuf {
        self.path(model, models_path).with_extension("gguf.part")
    }

    fn is_installed(&self, model: ModelChoice, models_path: &Path) -> bool {
        self.path(model, models_path)
            .metadata()
            .is_ok_and(|metadata| metadata.len() == self.size)
            && fs::read_to_string(self.verification_path(model, models_path))
                .is_ok_and(|sha256| sha256.trim() == self.sha256)
    }

    fn downloaded_bytes(&self, model: ModelChoice, models_path: &Path) -> u64 {
        if self.is_installed(model, models_path) {
            return self.size;
        }
        self.path(model, models_path)
            .metadata()
            .or_else(|_| self.partial_path(model, models_path).metadata())
            .map_or(0, |metadata| metadata.len().min(self.size))
    }

    fn stored_bytes(&self, model: ModelChoice, models_path: &Path) -> u64 {
        [
            &self.path(model, models_path),
            &self.partial_path(model, models_path),
        ]
        .into_iter()
        .filter_map(|path| path.metadata().ok())
        .map(|metadata| metadata.len())
        .sum()
    }
}

pub(crate) fn managed_artifacts(models_path: &Path) -> Vec<ManagedArtifact> {
    MANAGED_ARTIFACTS
        .into_iter()
        .map(|id| {
            let model = id.model();
            let artifact = id.artifact();
            let stored = artifact.stored_bytes(model, models_path);
            let complete_file = artifact
                .path(model, models_path)
                .metadata()
                .or_else(|_| artifact.partial_path(model, models_path).metadata())
                .is_ok_and(|metadata| metadata.len() == artifact.size);
            let state = if artifact.is_installed(model, models_path) {
                ManagedArtifactState::Ready
            } else if complete_file {
                ManagedArtifactState::NeedsVerification
            } else if stored > 0 {
                ManagedArtifactState::Partial
            } else {
                ManagedArtifactState::Missing
            };
            ManagedArtifact {
                id,
                stored,
                expected: artifact.size,
                state,
            }
        })
        .collect()
}

pub(crate) fn installed_models(models_path: &Path) -> Vec<ModelChoice> {
    MODEL_CHOICES
        .into_iter()
        .filter(|model| model.main_artifact().is_installed(*model, models_path))
        .collect()
}

pub(crate) fn artifact_download_progress(
    id: ManagedArtifactId,
    models_path: &Path,
) -> DownloadProgress {
    let model = id.model();
    let artifact = id.artifact();
    let downloaded = artifact.downloaded_bytes(model, models_path);
    let installed = artifact.is_installed(model, models_path);
    let verification =
        (!installed && downloaded >= artifact.size).then_some(VerificationProgress {
            verified: 0,
            total: artifact.size,
        });
    let phase = if installed {
        DownloadPhase::Complete
    } else if verification.is_some() {
        DownloadPhase::Verifying(artifact.label)
    } else if downloaded > 0 {
        DownloadPhase::Downloading(artifact.label)
    } else {
        DownloadPhase::Pending(artifact.label)
    };
    DownloadProgress {
        downloaded,
        total: artifact.size,
        phase,
        verification,
    }
}

pub(crate) fn artifact_verification_progress(
    id: ManagedArtifactId,
    verified: u64,
) -> DownloadProgress {
    let artifact = id.artifact();
    DownloadProgress {
        downloaded: artifact.size,
        total: artifact.size,
        phase: DownloadPhase::Verifying(artifact.label),
        verification: Some(VerificationProgress {
            verified: verified.min(artifact.size),
            total: artifact.size,
        }),
    }
}
