'''
ostris/ai-toolkit on https://modal.com  —  rewritten for Modal >= 1.0 API.

Run training (from inside the ai-toolkit directory) with:
    modal run run_modal.py --config-file-list-str=/root/ai-toolkit/config/numzoo_klein.yaml
'''

import os
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
os.environ["DISABLE_TELEMETRY"] = "YES"
import modal

# --- paths -----------------------------------------------------------------
LOCAL_AI_TOOLKIT = os.path.dirname(os.path.abspath(__file__))  # this repo, shipped to the container
REMOTE_AI_TOOLKIT = "/root/ai-toolkit"
MOUNT_DIR = "/root/ai-toolkit/modal_output"  # volume mount (must be an empty path)

# --- persistent volume for outputs (LoRA weights, samples) -----------------
model_volume = modal.Volume.from_name("flux-lora-models", create_if_missing=True)

# --- container image: ai-toolkit deps + the repo + dataset shipped in ------
# Use ai-toolkit's OWN pinned requirements (requirements.txt -> requirements_base.txt).
# They pin a specific diffusers commit that supports FLUX.2 klein and avoid the pip
# "resolution-too-deep" blowup that unpinned deps (esp. old controlnet_aux) caused.
image = (
    modal.Image.debian_slim(python_version="3.11")
    .apt_install("libgl1", "libglib2.0-0", "git")
    # CUDA-enabled torch first, pinned to what torchao/torchcodec expect.
    # torchaudio is imported by ai-toolkit's config_modules — must match torch ver.
    .pip_install("torch==2.7.1", "torchvision==0.22.1", "torchaudio==2.7.1")
    # copy both requirements files into the build context, then install (pinned)
    .add_local_file(os.path.join(LOCAL_AI_TOOLKIT, "requirements.txt"),
                    "/tmp/reqs/requirements.txt", copy=True)
    .add_local_file(os.path.join(LOCAL_AI_TOOLKIT, "requirements_base.txt"),
                    "/tmp/reqs/requirements_base.txt", copy=True)
    .run_commands("cd /tmp/reqs && pip install -r requirements.txt")
    # ship the ai-toolkit repo (code + numzoo-dataset/) into the container LAST
    # (copy=False runtime mount must be the final image layer).
    .add_local_dir(
        LOCAL_AI_TOOLKIT,
        REMOTE_AI_TOOLKIT,
        ignore=[
            ".git", "**/.git", "**/node_modules", "**/__pycache__", "**/*.pyc",
            "modal_output", "**/.venv", "venv", "ui/node_modules", "output",
        ],
    )
)

app = modal.App(name="numzoo-lora-training", image=image, volumes={MOUNT_DIR: model_volume})


@app.function(
    gpu="A100",                 # 40GB — plenty for a quantized 4B LoRA (<24GB)
    timeout=7200,               # 2h ceiling; a 1500-step run is ~45 min
    secrets=[modal.Secret.from_name("huggingface")],  # provides HF_TOKEN for gated model
)
def main(config_file_list_str: str, recover: bool = False, name: str = None):
    import sys
    sys.path.insert(0, REMOTE_AI_TOOLKIT)
    from toolkit.job import get_job  # imported in-container (pulls torch etc.)

    config_file_list = config_file_list_str.split(",")
    jobs_completed = 0
    jobs_failed = 0
    print(f"Running {len(config_file_list)} job{'' if len(config_file_list) == 1 else 's'}")

    for config_file in config_file_list:
        try:
            job = get_job(config_file, name)
            job.config['process'][0]['training_folder'] = MOUNT_DIR
            os.makedirs(MOUNT_DIR, exist_ok=True)
            print(f"Training outputs will be saved to: {MOUNT_DIR}")
            job.run()
            model_volume.commit()   # persist weights/samples after each job
            job.cleanup()
            jobs_completed += 1
        except Exception as e:
            print(f"Error running job: {e}")
            jobs_failed += 1
            if not recover:
                raise

    print("========================================")
    print(f"Result: {jobs_completed} completed, {jobs_failed} failed")
    print("========================================")


@app.local_entrypoint()
def run(config_file_list_str: str, recover: bool = False, name: str = None):
    main.remote(config_file_list_str=config_file_list_str, recover=recover, name=name)