"""Modal app for testing the SCAIL-2 LTX-2 training/inference path.

Free-tier friendly: `verify` runs on CPU and `smoke` on a cheap T4, both using
tiny random-init models + synthetic data (no checkpoint, no dataset) so a full run
costs pennies. `train` is the scale-up entrypoint that runs the real trainer once
you upload a checkpoint + preprocessed data to the `scail-data` Volume -- that one
uses a real GPU and is NOT free.

Setup (once):
    pip install modal && modal setup

Run:
    modal run modal/app.py::verify     # CPU plumbing check (~free)
    modal run modal/app.py::smoke      # same checks on a T4 GPU (cents)
    modal run modal/app.py::train --config-rel configs/scail_animation_lora.yaml
"""

import subprocess
from pathlib import Path

import modal

REPO = Path(__file__).parent.parent

# Build the workspace once into the image. ltx-kernels (CUDA-compiled) is excluded
# from the workspace and not needed -- attention falls back to SDPA.
image = (
    modal.Image.debian_slim(python_version="3.11")
    .apt_install("git")
    .pip_install("uv")
    .env({"UV_LINK_MODE": "copy", "UV_PROJECT_ENVIRONMENT": "/root/LTX-2/.venv"})
    .add_local_dir(
        str(REPO),
        "/root/LTX-2",
        copy=True,
        ignore=[".git", ".venv", "**/__pycache__", "**/*.pyc", "outputs", "wandb", "**/.pytest_cache"],
    )
    .run_commands("cd /root/LTX-2 && uv sync --package ltx-trainer")
)

app = modal.App("scail-ltx2", image=image)

# Persistent volume for the (large) checkpoint + preprocessed dataset used by `train`.
data_volume = modal.Volume.from_name("scail-data", create_if_missing=True)

_UV_PY = ["uv", "run", "--package", "ltx-trainer", "python"]


def _run(args: list[str]) -> None:
    subprocess.run([*_UV_PY, *args], cwd="/root/LTX-2", check=True)


@app.function(timeout=1800)
def verify() -> None:
    """Run the SCAIL-2 plumbing checks on CPU (near-free)."""
    _run(["modal/checks.py"])


@app.function(gpu="T4", timeout=1800)
def smoke() -> None:
    """Run the same checks on a T4 GPU to validate the CUDA path (cents)."""
    _run(["modal/checks.py"])


@app.function(gpu="A10G", timeout=60 * 60 * 6, volumes={"/data": data_volume})
def train(config_rel: str = "configs/scail_animation_lora.yaml") -> None:
    """Run the real SCAIL trainer. NOT free -- needs a real checkpoint + data.

    Upload your base checkpoint, Gemma text encoder, and preprocessed data to the
    ``scail-data`` Volume (mounted at /data), and point the config's paths at /data.
    For the full 19B model use a bigger GPU (e.g. gpu="A100") and expect real cost.
    """
    _run(["packages/ltx-trainer/scripts/train.py", config_rel])
    data_volume.commit()


@app.local_entrypoint()
def main(target: str = "verify", config_rel: str = "configs/scail_animation_lora.yaml") -> None:
    if target == "smoke":
        smoke.remote()
    elif target == "train":
        train.remote(config_rel)
    else:
        verify.remote()