#!/usr/bin/env python3
"""Evaluate a Modal Volume Clover inpainting candidate against the release."""

from __future__ import annotations

import os
import subprocess
import sys
from pathlib import Path

import modal

APP_NAME = "clover-image-tiny-inpaint-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",
        "datasets==4.8.5",
        "diffusers==0.39.0",
        "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=2 * 60 * 60, cpu=4, memory=24576)
def evaluate(candidate_name: str, evaluation_name: str, sample_count: int = 6) -> 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"),
            "HF_DATASETS_CACHE": str(CACHE_ROOT / "huggingface" / "datasets"),
            "TOKENIZERS_PARALLELISM": "false",
        }
    )
    command = [
        sys.executable,
        "-u",
        "/root/inpainting/evaluate.py",
        "--baseline_model",
        "neonforestmist/Clover-Image-Tiny-Inpaint",
        "--baseline_revision",
        "1b6f8ae3db51900520369d5522c7dc7c2a97e21e",
        "--candidate_model",
        str(candidate),
        "--dataset_name",
        "prithivMLmods/Caption3o-Opt",
        "--dataset_revision",
        "17e893f785fcd3f5d6fc4a5d65a914b9f7b1ff5b",
        "--sample_count",
        str(sample_count),
        "--output_dir",
        str(destination),
    ]
    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, sample_count: int = 6) -> None:
    result = evaluate.remote(candidate_name, evaluation_name, sample_count)
    print(f"Evaluation is available in Modal Volume {OUTPUT_VOLUME_NAME}: {result}")