#!/usr/bin/env python3
"""Materialize a resumable training checkpoint as a Diffusers pipeline."""

from __future__ import annotations

import json
import os
from pathlib import Path

import modal

APP_NAME = "clover-image-tiny-inpaint-snapshot"
OUTPUT_VOLUME_NAME = "clover-image-tiny-inpaint-output"
CACHE_VOLUME_NAME = "clover-image-tiny-inpaint-cache"
OUTPUT_ROOT = Path("/outputs")
CACHE_ROOT = Path("/cache")
INITIAL_MODEL = "neonforestmist/Clover-Image-Tiny-Inpaint"
INITIAL_REVISION = "1b6f8ae3db51900520369d5522c7dc7c2a97e21e"

image = modal.Image.debian_slim(python_version="3.11").pip_install(
    "accelerate==1.14.0",
    "diffusers==0.39.0",
    "huggingface_hub==0.36.0",
    "safetensors==0.8.0",
    "torch==2.7.0",
    "transformers==4.57.6",
)

output_volume = modal.Volume.from_name(OUTPUT_VOLUME_NAME, create_if_missing=True)
cache_volume = modal.Volume.from_name(CACHE_VOLUME_NAME, create_if_missing=True)
app = modal.App(
    APP_NAME,
    image=image,
    volumes={
        str(OUTPUT_ROOT): output_volume,
        str(CACHE_ROOT): cache_volume,
    },
)


@app.function(timeout=60 * 60, cpu=4, memory=16384)
def materialize(source_name: str, checkpoint_step: int, output_name: str) -> str:
    import torch
    from diffusers import StableDiffusionInpaintPipeline
    from safetensors.torch import load_file

    source = OUTPUT_ROOT / source_name
    checkpoint = source / "checkpoints" / f"checkpoint-{checkpoint_step}"
    weights = checkpoint / "model.safetensors"
    if not weights.is_file():
        raise RuntimeError(f"Missing checkpoint weights: {weights}")
    destination = OUTPUT_ROOT / output_name
    if destination.exists():
        raise RuntimeError(f"Snapshot output already exists: {destination}")

    os.environ.update(
        {
            "HF_HOME": str(CACHE_ROOT / "huggingface"),
            "HF_HUB_CACHE": str(CACHE_ROOT / "huggingface" / "hub"),
            "TOKENIZERS_PARALLELISM": "false",
        }
    )
    pipeline = StableDiffusionInpaintPipeline.from_pretrained(
        INITIAL_MODEL,
        revision=INITIAL_REVISION,
        torch_dtype=torch.float32,
    )
    state = load_file(str(weights), device="cpu")
    incompatible = pipeline.unet.load_state_dict(state, strict=True)
    if incompatible.missing_keys or incompatible.unexpected_keys:
        raise RuntimeError(f"Checkpoint state mismatch: {incompatible}")
    pipeline.save_pretrained(destination, safe_serialization=True)

    metadata = {
        "snapshot_type": "training_checkpoint",
        "source_name": source_name,
        "checkpoint_step": checkpoint_step,
        "initial_model": INITIAL_MODEL,
        "initial_revision": INITIAL_REVISION,
    }
    (destination / "snapshot.json").write_text(json.dumps(metadata, indent=2) + "\n")
    (destination / "training-complete.json").write_text(
        json.dumps(metadata, indent=2) + "\n"
    )
    output_volume.commit()
    cache_volume.commit()
    return str(destination)


@app.local_entrypoint()
def main(source_name: str, checkpoint_step: int, output_name: str) -> None:
    result = materialize.remote(source_name, checkpoint_step, output_name)
    print(f"Snapshot is available in Modal Volume {OUTPUT_VOLUME_NAME}: {result}")