use anyhow::{Context, Result};
use hf_hub::{BucketTreeEntry, HFClient};
use sha2::{Digest, Sha256};

/// The embedded canonical job script.
pub const EMBEDDED_SCRIPT: &str = include_str!("scripts/split-model-job.sh");

// Generated by build.rs — provides EMBEDDED_SCRIPT_SHA256 and EMBEDDED_SCRIPT_SIZE.
include!(concat!(env!("OUT_DIR"), "/script_hash.rs"));

/// Result of comparing the bucket script against the embedded version.
#[derive(Debug)]
pub struct ScriptFreshness {
    /// Size of the script in the bucket (bytes).
    pub bucket_size: u64,
    /// Size of the embedded script (bytes).
    pub expected_size: u64,
    /// Whether the bucket script matches the embedded version.
    pub is_current: bool,
    /// Human-readable reason when the script is not current.
    pub mismatch_reason: Option<String>,
}

/// Compare the bucket script's size against the embedded version.
///
/// If sizes match, downloads the ~10 KB file and verifies the SHA-256 hash.
pub async fn check_bucket_script(client: &HFClient) -> Result<ScriptFreshness> {
    let bucket = client.bucket("meshllm", "layer-split-output");
    let paths = vec!["split-model-job.sh".to_string()];

    let entries = bucket
        .get_paths_info(&paths)
        .await
        .context("check bucket script metadata")?;

    let bucket_size = entries
        .into_iter()
        .find_map(|entry| match entry {
            BucketTreeEntry::File { size, .. } => Some(size),
            _ => None,
        })
        .context("split-model-job.sh not found in meshllm/layer-split-output bucket")?;

    if bucket_size != EMBEDDED_SCRIPT_SIZE {
        return Ok(ScriptFreshness {
            bucket_size,
            expected_size: EMBEDDED_SCRIPT_SIZE,
            is_current: false,
            mismatch_reason: Some(format!(
                "size differs: bucket has {bucket_size} bytes, embedded script has {expected_size} bytes",
                expected_size = EMBEDDED_SCRIPT_SIZE
            )),
        });
    }

    // Sizes match — download to a unique temp file and verify SHA-256.
    let tmp_dir = tempfile::tempdir().context("create bucket script check temp dir")?;
    let tmp_path = tmp_dir.path().join("split-model-job.sh");

    let download_params = hf_hub::BucketDownloadFilesParams {
        files: vec![("split-model-job.sh".to_string(), tmp_path.clone())],
    };

    bucket
        .download_files(&download_params, &Default::default())
        .await
        .context("download bucket script for hash verification")?;

    let bytes = std::fs::read(&tmp_path).context("read downloaded bucket script")?;
    let hash = Sha256::digest(&bytes);
    let hex: String = hash.iter().map(|b| format!("{b:02x}")).collect();

    let is_current = hex == EMBEDDED_SCRIPT_SHA256;

    Ok(ScriptFreshness {
        bucket_size,
        expected_size: EMBEDDED_SCRIPT_SIZE,
        is_current,
        mismatch_reason: (!is_current).then(|| {
            format!(
                "sha256 differs: bucket has {hex}, embedded script has {expected}",
                expected = EMBEDDED_SCRIPT_SHA256
            )
        }),
    })
}

/// Upload the embedded script to the meshllm bucket.
///
/// Requires meshllm org write access.
pub async fn update_bucket_script(client: &HFClient) -> Result<()> {
    let bucket = client.bucket("meshllm", "layer-split-output");

    // Write embedded script to a unique temp file for upload.
    let tmp_dir = tempfile::tempdir().context("create bucket script upload temp dir")?;
    let tmp_path = tmp_dir.path().join("split-model-job.sh");
    std::fs::write(&tmp_path, EMBEDDED_SCRIPT)?;

    let files = vec![(tmp_path.clone(), "split-model-job.sh".to_string())];
    bucket
        .upload_files(&files, &Default::default())
        .await
        .context("upload script to meshllm/layer-split-output bucket")?;

    Ok(())
}

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

    #[test]
    fn embedded_script_prints_mesh_llm_repo_shorthand() {
        assert!(EMBEDDED_SCRIPT.contains(r#"mesh-llm serve --model "{target_repo}" --split"#));
        assert!(EMBEDDED_SCRIPT.contains("mesh-llm serve --model ${TARGET_REPO} --split"));
        assert!(!EMBEDDED_SCRIPT.contains(r#"mesh-llm serve --model "hf://${TARGET_REPO}""#));
        assert!(!EMBEDDED_SCRIPT.contains("mesh-llm serve --model hf://${TARGET_REPO}"));
    }

    #[test]
    fn embedded_script_writes_rich_model_card() {
        assert!(EMBEDDED_SCRIPT.contains("pipeline_tag: text-generation"));
        assert!(EMBEDDED_SCRIPT.contains("- openai-compatible"));
        assert!(EMBEDDED_SCRIPT.contains("## Model Overview"));
        assert!(EMBEDDED_SCRIPT.contains("## Highlights"));
        assert!(EMBEDDED_SCRIPT.contains("## Recommended Use"));
        assert!(EMBEDDED_SCRIPT.contains("## Quickstart"));
        assert!(EMBEDDED_SCRIPT.contains("## Package Variant"));
        assert!(EMBEDDED_SCRIPT.contains("## What Is Included"));
        assert!(EMBEDDED_SCRIPT.contains("## Validation"));
        assert!(EMBEDDED_SCRIPT.contains("docs/mesh-llm-logo.svg"));
        assert!(EMBEDDED_SCRIPT.contains("style=for-the-badge"));
        assert!(EMBEDDED_SCRIPT.contains("https://www.meshllm.cloud"));
        assert!(EMBEDDED_SCRIPT.contains("https://discord.gg/rs6fmc63eN"));
        assert!(EMBEDDED_SCRIPT.contains("Distributed GGUF inference package for Mesh LLM"));
        assert!(EMBEDDED_SCRIPT.contains(
            "For upstream architecture details, chat template guidance, sampling recommendations"
        ));
        assert!(EMBEDDED_SCRIPT.contains("Package manifest SHA-256"));
        assert!(EMBEDDED_SCRIPT.contains("uploaded to this repository"));
    }

    #[test]
    fn embedded_script_streams_package_artifacts_to_hub() {
        assert!(EMBEDDED_SCRIPT.contains("SOURCE_QUANT"));
        assert!(EMBEDDED_SCRIPT.contains("SOURCE_TOTAL_BYTES"));
        assert!(EMBEDDED_SCRIPT.contains("Estimated fallback /bucket cache needed"));
        assert!(EMBEDDED_SCRIPT.contains("estimate_bucket_workspace_bytes"));
        assert!(EMBEDDED_SCRIPT.contains(r#"HF_HUB_CACHE="${HF_HUB_CACHE:-${HF_HOME}/hub}""#));
        assert!(EMBEDDED_SCRIPT.contains(r#"HF_XET_CACHE="${HF_XET_CACHE:-${HF_HOME}/xet}""#));
        assert!(
            EMBEDDED_SCRIPT.contains(r#"PACKAGE_DIR="${PACKAGE_DIR:-${LOCAL_WORK_DIR}/package}""#)
        );
        assert!(EMBEDDED_SCRIPT.contains(r#"JOB_TMP_DIR="${JOB_TMP_DIR:-${LOCAL_WORK_DIR}/tmp}""#));
        assert!(EMBEDDED_SCRIPT.contains(r#"LOCAL_WORK_DIR="${LOCAL_WORK_DIR:-/tmp/"#));
        assert!(
            EMBEDDED_SCRIPT.contains(r#"BUILD_TMP_DIR="${BUILD_TMP_DIR:-${LOCAL_WORK_DIR}/tmp}""#)
        );
        assert!(EMBEDDED_SCRIPT.contains(r#"TMPDIR="$BUILD_TMP_DIR""#));
        assert!(EMBEDDED_SCRIPT.contains(r#"TMPDIR="$JOB_TMP_DIR""#));
        assert!(EMBEDDED_SCRIPT.contains("export TMPDIR TEMP TMP"));
        assert!(EMBEDDED_SCRIPT
            .contains(r#"CARGO_TARGET_DIR="${CARGO_TARGET_DIR:-${LOCAL_WORK_DIR}/cargo-target}""#));
        assert!(EMBEDDED_SCRIPT
            .contains(r#"rm -rf "$BUILD_DIR" "$CARGO_TARGET_DIR" "$CARGO_HOME" "$RUSTUP_HOME""#));
        assert!(EMBEDDED_SCRIPT.contains(r#"SOURCE_REF="${SOURCE_REPO}:${SOURCE_QUANT}""#));
        assert!(EMBEDDED_SCRIPT
            .contains(r#"SOURCE_REF="${SOURCE_REPO}@${SOURCE_REVISION}:${SOURCE_QUANT}""#));
        assert!(EMBEDDED_SCRIPT.contains("log_storage_snapshot"));
        assert!(EMBEDDED_SCRIPT.contains("start_heartbeat"));
        assert!(EMBEDDED_SCRIPT.contains("Starting write-package"));
        assert!(EMBEDDED_SCRIPT.contains("upload-package-artifact.py"));
        assert!(EMBEDDED_SCRIPT.contains("SKIPPY_PACKAGE_ARTIFACT_PATH"));
        assert!(EMBEDDED_SCRIPT.contains("SKIPPY_PACKAGE_ARTIFACT_RELATIVE_PATH"));
        assert!(EMBEDDED_SCRIPT.contains(r#"--after-artifact-command "$ARTIFACT_UPLOAD_HOOK""#));
        assert!(EMBEDDED_SCRIPT.contains("Uploaded and removed"));
        assert!(!EMBEDDED_SCRIPT.contains("api.upload_folder"));
        assert!(EMBEDDED_SCRIPT.contains(r#"MOUNTED_SOURCE_PATH="/source/${SOURCE_FILE}""#));
        assert!(EMBEDDED_SCRIPT.contains(r#"WRITE_PACKAGE_INPUT="$MOUNTED_SOURCE_PATH""#));
        assert!(EMBEDDED_SCRIPT.contains(r#"--source-file "$SOURCE_FILE""#));
        assert!(EMBEDDED_SCRIPT.contains(r#"time "$SLICER" write-package "$WRITE_PACKAGE_INPUT""#));
        assert!(!EMBEDDED_SCRIPT.contains(r#"time $SLICER write-package "$SOURCE_PATH""#));
    }
}
