#!/usr/bin/env python3
"""Run context-aware Clover Image Tiny inpainting distillation on Modal."""

from __future__ import annotations

import os
import subprocess
import sys
import time
from pathlib import Path

import modal

APP_NAME = "clover-image-tiny-inpaint-v2"
OUTPUT_VOLUME_NAME = "clover-image-tiny-inpaint-output"
CACHE_VOLUME_NAME = "clover-image-tiny-inpaint-cache"
OUTPUT_ROOT = Path("/outputs")
CACHE_ROOT = Path("/cache")

image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install(
        "accelerate==1.14.0",
        "datasets==4.8.5",
        "diffusers==0.39.0",
        "ftfy==6.3.1",
        "huggingface_hub==0.36.0",
        "numpy==2.2.6",
        "pillow==12.3.0",
        "safetensors==0.8.0",
        "torch==2.7.0",
        "torchvision==0.22.0",
        "transformers==4.57.6",
    )
    .add_local_dir("inpainting", remote_path="/root/inpainting")
)

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(gpu="A10", timeout=12 * 60 * 60, cpu=8, memory=32768)
def train(
    *,
    base_model: str = "neonforestmist/Clover-Image-Tiny",
    base_revision: str = "63b0e9f6be9c00888ff464f342a9ef052bf76681",
    initial_inpaint_model: str = "neonforestmist/Clover-Image-Tiny-Inpaint",
    initial_inpaint_revision: str = "1b6f8ae3db51900520369d5522c7dc7c2a97e21e",
    teacher_model: str = "stable-diffusion-v1-5/stable-diffusion-inpainting",
    teacher_revision: str = "8a4288a76071f7280aedbdb3253bdb9e9d5d84bb",
    dataset: str = "prithivMLmods/Caption3o-Opt",
    dataset_revision: str = "17e893f785fcd3f5d6fc4a5d65a914b9f7b1ff5b",
    dataset_split: str = "train",
    image_column: str = "image",
    caption_column: str = "caption",
    max_train_steps: int = 12000,
    output_name: str = "clover-image-tiny-inpaint-v2",
    max_train_samples: int | None = None,
    learning_rate: float = 5e-6,
    lr_warmup_steps: int | None = None,
    teacher_loss_weight: float = 0.75,
    ground_truth_loss_weight: float = 0.25,
    context_loss_weight: float = 0.25,
    masked_loss_weight: float = 2.5,
    boundary_loss_weight: float = 2.0,
    seed: int = 20260811,
    checkpointing_steps: int = 500,
    resume: bool = False,
) -> str:
    """Train v2 and return its persistent Modal Volume path."""

    output_dir = OUTPUT_ROOT / output_name
    if (output_dir / "training-complete.json").exists():
        raise RuntimeError(f"Completed output already exists: {output_dir}")
    if output_dir.exists() and not resume:
        raise RuntimeError(
            f"Partial output already exists; choose another name or enable resume: {output_dir}"
        )

    command = [
        sys.executable,
        "-u",
        "/root/inpainting/train.py",
        "--pretrained_model_name_or_path",
        base_model,
        "--revision",
        base_revision,
        "--initial_inpaint_model",
        initial_inpaint_model,
        "--teacher_model_name_or_path",
        teacher_model,
        "--teacher_revision",
        teacher_revision,
        "--dataset_name",
        dataset,
        "--dataset_revision",
        dataset_revision,
        "--dataset_split",
        dataset_split,
        "--image_column",
        image_column,
        "--caption_column",
        caption_column,
        "--max_train_steps",
        str(max_train_steps),
        "--learning_rate",
        str(learning_rate),
        "--lr_warmup_steps",
        str(
            lr_warmup_steps
            if lr_warmup_steps is not None
            else min(500, max(1, max_train_steps // 20))
        ),
        "--gradient_accumulation_steps",
        "4",
        "--gradient_checkpointing",
        "--mixed_precision",
        "bf16",
        "--random_flip",
        "--caption_dropout_probability",
        "0.10",
        "--validation_samples",
        "128",
        "--teacher_loss_weight",
        str(teacher_loss_weight),
        "--ground_truth_loss_weight",
        str(ground_truth_loss_weight),
        "--context_loss_weight",
        str(context_loss_weight),
        "--masked_loss_weight",
        str(masked_loss_weight),
        "--boundary_loss_weight",
        str(boundary_loss_weight),
        "--snr_gamma",
        "5.0",
        "--checkpointing_steps",
        str(checkpointing_steps),
        "--checkpoints_total_limit",
        "3",
        "--seed",
        str(seed),
        "--output_dir",
        str(output_dir),
    ]
    if initial_inpaint_revision:
        command.extend(["--initial_inpaint_revision", initial_inpaint_revision])
    if max_train_samples is not None:
        command.extend(["--max_train_samples", str(max_train_samples)])
    if resume:
        command.extend(["--resume_from_checkpoint", "latest"])

    env = os.environ.copy()
    env.update(
        {
            "HF_HOME": str(CACHE_ROOT / "huggingface"),
            "HF_HUB_CACHE": str(CACHE_ROOT / "huggingface" / "hub"),
            "HF_DATASETS_CACHE": str(CACHE_ROOT / "huggingface" / "datasets"),
            "TOKENIZERS_PARALLELISM": "false",
            "PYTHONUNBUFFERED": "1",
        }
    )
    process = subprocess.Popen(command, env=env)
    last_commit = time.monotonic()
    while process.poll() is None:
        time.sleep(30)
        if time.monotonic() - last_commit >= 300:
            output_volume.commit()
            cache_volume.commit()
            last_commit = time.monotonic()
    if process.returncode:
        output_volume.commit()
        cache_volume.commit()
        raise subprocess.CalledProcessError(process.returncode, command)

    output_volume.commit()
    cache_volume.commit()
    return str(output_dir)


@app.local_entrypoint()
def main(
    smoke: bool = False,
    steps: int = 12000,
    output_name: str = "clover-image-tiny-inpaint-v2",
    initial_inpaint_model: str = "neonforestmist/Clover-Image-Tiny-Inpaint",
    initial_inpaint_revision: str | None = "1b6f8ae3db51900520369d5522c7dc7c2a97e21e",
    learning_rate: float = 5e-6,
    lr_warmup_steps: int | None = None,
    teacher_loss_weight: float = 0.75,
    ground_truth_loss_weight: float = 0.25,
    context_loss_weight: float = 0.25,
    masked_loss_weight: float = 2.5,
    boundary_loss_weight: float = 2.0,
    resume: bool = False,
) -> None:
    """Launch a small pipeline test or the full bounded A100 run."""

    if smoke:
        steps = min(steps, 4)
        samples = 8
        output_name = f"{output_name}-smoke-{int(time.time())}"
        checkpointing_steps = 2
    else:
        samples = None
        checkpointing_steps = 500
    result = train.remote(
        max_train_steps=steps,
        max_train_samples=samples,
        output_name=output_name,
        initial_inpaint_model=initial_inpaint_model,
        initial_inpaint_revision=initial_inpaint_revision,
        learning_rate=learning_rate,
        lr_warmup_steps=lr_warmup_steps,
        teacher_loss_weight=teacher_loss_weight,
        ground_truth_loss_weight=ground_truth_loss_weight,
        context_loss_weight=context_loss_weight,
        masked_loss_weight=masked_loss_weight,
        boundary_loss_weight=boundary_loss_weight,
        checkpointing_steps=checkpointing_steps,
        resume=resume,
    )
    print(f"Training output is available in Modal Volume {OUTPUT_VOLUME_NAME}: {result}")