#!/usr/bin/env python3
"""Run controlled semantic inpainting evaluation from a Modal Volume candidate."""

from __future__ import annotations

import os
import subprocess
import sys
from pathlib import Path

import modal

APP_NAME = "clover-image-tiny-inpaint-semantic-eval"
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",
        "diffusers==0.39.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=3 * 60 * 60, cpu=4, memory=24576)
def evaluate(
    candidate_name: str,
    evaluation_name: str,
    guidance_scale: float = 7.5,
    steps: int = 30,
    mask_crop_padding: int = 0,
) -> str:
    candidate = OUTPUT_ROOT / candidate_name
    if not (candidate / "training-complete.json").exists():
        raise RuntimeError(f"Candidate is not complete: {candidate}")
    destination = OUTPUT_ROOT / "evaluations" / evaluation_name
    if destination.exists():
        raise RuntimeError(f"Evaluation output already exists: {destination}")
    env = os.environ.copy()
    env.update(
        {
            "HF_HOME": str(CACHE_ROOT / "huggingface"),
            "HF_HUB_CACHE": str(CACHE_ROOT / "huggingface" / "hub"),
            "TOKENIZERS_PARALLELISM": "false",
        }
    )
    command = [
        sys.executable,
        "-u",
        "/root/inpainting/evaluate_semantic.py",
        "--base_model",
        "neonforestmist/Clover-Image-Tiny",
        "--base_revision",
        "63b0e9f6be9c00888ff464f342a9ef052bf76681",
        "--baseline_model",
        "neonforestmist/Clover-Image-Tiny-Inpaint",
        "--baseline_revision",
        "1b6f8ae3db51900520369d5522c7dc7c2a97e21e",
        "--teacher_model",
        "stable-diffusion-v1-5/stable-diffusion-inpainting",
        "--teacher_revision",
        "8a4288a76071f7280aedbdb3253bdb9e9d5d84bb",
        "--teacher_variant",
        "fp16",
        "--candidate_model",
        str(candidate),
        "--clip_model",
        "openai/clip-vit-base-patch32",
        "--clip_revision",
        "3d74acf9a28c67741b2f4f2ea7635f0aaf6f0268",
        "--output_dir",
        str(destination),
        "--guidance_scale",
        str(guidance_scale),
        "--steps",
        str(steps),
        "--mask_crop_padding",
        str(mask_crop_padding),
    ]
    subprocess.run(command, check=True, env=env)
    output_volume.commit()
    cache_volume.commit()
    return str(destination)


@app.local_entrypoint()
def main(
    candidate_name: str,
    evaluation_name: str,
    guidance_scale: float = 7.5,
    steps: int = 30,
    mask_crop_padding: int = 0,
) -> None:
    result = evaluate.remote(
        candidate_name,
        evaluation_name,
        guidance_scale,
        steps,
        mask_crop_padding,
    )
    print(f"Evaluation is available in Modal Volume {OUTPUT_VOLUME_NAME}: {result}")