"""Modal deployment of the simulation worker (the default SIM_BACKEND).

Deploy from the repository root (the image reuses api/ and worker/)::

    modal secret create rosdiff-r2 R2_ACCOUNT_ID=... R2_ACCESS_KEY_ID=... R2_SECRET_ACCESS_KEY=... R2_BUCKET=...
    modal deploy -m worker.modal_app

These environment variables are read at deploy time:

    MODAL_APP              app name (default rosdiff-sim; the API's MODAL_APP must match)
    MODAL_GPU              GPU type (default T4, the cheapest); "none" renders on CPU instead
    MODAL_MAX_CONTAINERS   runs executing at once (default 2); more runs wait in Modal's queue
    RUN_EXECUTION_TIMEOUT_S  Modal stops a run after this long (default 600)
    MODAL_R2_SECRET        name of the Modal secret holding the R2 credentials (default rosdiff-r2)
"""

from __future__ import annotations

import os
from pathlib import Path

import modal

APP_NAME = os.environ.get("MODAL_APP", "rosdiff-sim")
_gpu = os.environ.get("MODAL_GPU", "T4").strip()
GPU = None if _gpu.lower() in ("", "none", "cpu") else _gpu
GL = "egl" if GPU else "osmesa"
MENAGERIE_DIR = "/opt/mujoco_menagerie"
ROBOT_DIRS = ("unitree_g1", "unitree_h1", "unitree_go2", "franka_emika_panda")  # api.robots.ROBOTS

image = (
    modal.Image.debian_slim(python_version="3.11")
    .apt_install("git", "libglvnd0", "libgl1", "libegl1", "libgles2", "libosmesa6")
    .run_commands(
        # Tell libglvnd where the NVIDIA EGL implementation is (the driver is provided by Modal's GPU hosts).
        "mkdir -p /usr/share/glvnd/egl_vendor.d && printf "
        '\'{"file_format_version":"1.0.0","ICD":{"library_path":"libEGL_nvidia.so.0"}}\\n\' '
        "> /usr/share/glvnd/egl_vendor.d/10_nvidia.json",
        # Robot models (only the ones generated scenes may include).
        "git clone --depth 1 --filter=blob:none --sparse https://github.com/google-deepmind/mujoco_menagerie.git "
        f"{MENAGERIE_DIR} && cd {MENAGERIE_DIR} && git sparse-checkout set {' '.join(ROBOT_DIRS)} && rm -rf .git",
    )
    .pip_install_from_requirements(str(Path(__file__).with_name("requirements.txt")))
    .env({"MENAGERIE_DIR": MENAGERIE_DIR, "MUJOCO_GL": GL, "PYOPENGL_PLATFORM": GL, "PYTHONUNBUFFERED": "1"})
    .add_local_python_source("api", "worker")
)

app = modal.App(APP_NAME)


@app.function(
    image=image,
    gpu=GPU,
    timeout=int(os.environ.get("RUN_EXECUTION_TIMEOUT_S", "600")),
    max_containers=int(os.environ.get("MODAL_MAX_CONTAINERS", "2")),
    scaledown_window=20,  # stop paying for an idle container quickly (runs are sporadic)
    secrets=[modal.Secret.from_name(os.environ.get("MODAL_R2_SECRET", "rosdiff-r2"))],
)
def run_simulation(job_input: dict) -> dict:
    from worker.handler import handler

    return handler({"input": {**job_input, "function": "run_simulation"}})


@app.function(
    image=image,
    gpu=GPU,
    timeout=900,
    max_containers=int(os.environ.get("MODAL_MAX_CONTAINERS", "2")),
    scaledown_window=20,
)
def sweep(job_input: dict) -> dict:
    """Robustness sweep with MuJoCo Warp; results return inline (no storage secret needed)."""
    from worker.handler import handler

    return handler({"input": {**job_input, "function": "sweep"}})


# --------------------------------------------------------------------------- ROS 2 sandbox (/code/check)


def _ros_image(distro: str) -> modal.Image:
    # Same Dockerfile CI builds and tests; Modal's own Python is added alongside the image's (the checks call
    # ROS through bash with a clean environment, so the two never mix).
    return modal.Image.from_dockerfile(
        str(Path(__file__).parent / "ros" / "Dockerfile"),
        context_dir=str(Path(__file__).parent.parent),
        build_args={"ROS_DISTRO": distro},
        add_python="3.11",
    ).add_local_python_source("worker")


# Generated code is untrusted: no network, no secrets. Results come back inline and the API stores them.
_ROS_SANDBOX = dict(
    cpu=2.0,
    memory=4096,
    timeout=900,
    max_containers=int(os.environ.get("MODAL_MAX_CONTAINERS", "2")),
    scaledown_window=20,
    block_network=True,
)


@app.function(image=_ros_image("humble"), **_ROS_SANDBOX)
def ros_check_humble(job_input: dict) -> dict:
    return _ros_check(job_input, "humble")


@app.function(image=_ros_image("jazzy"), **_ROS_SANDBOX)
def ros_check_jazzy(job_input: dict) -> dict:
    return _ros_check(job_input, "jazzy")


def _ros_check(job_input: dict, distro: str) -> dict:
    import traceback

    from worker.ros_runner import run_check_job

    try:
        return run_check_job(job_input, distro)
    except Exception as e:
        return {"error": f"{type(e).__name__}: {e}", "traceback": traceback.format_exc()[-2000:]}


# --------------------------------------------------------------------------- Cosmos Reason review (L4)
# Opt-in (ENABLE_COSMOS=1 at deploy): it needs the rosdiff-hf secret and keeps the ~16 GB model in a Modal
# Volume, and volume storage is billed even when the spend limit is $0. Without it nothing is stored on Modal.
ENABLE_COSMOS = os.environ.get("ENABLE_COSMOS", "0") == "1"

if ENABLE_COSMOS:
    COSMOS_MODEL = os.environ.get("COSMOS_MODEL", "nvidia/Cosmos-Reason1-7B")
    models = modal.Volume.from_name("rosdiff-models", create_if_missing=True)
    cosmos_image = (
        modal.Image.debian_slim(python_version="3.11")
        .pip_install(
            "torch==2.7.1",
            "torchvision==0.22.1",
            "transformers>=4.51,<4.60",
            "accelerate>=1.6",
            "qwen-vl-utils>=0.0.8",
            "pillow",
            "imageio-ffmpeg==0.6.0",
            "huggingface_hub>=0.30",
        )
        .env({"HF_HOME": "/models/hf", "COSMOS_MODEL": COSMOS_MODEL})
        .add_local_python_source("worker")
    )

    @app.function(
        image=cosmos_image,
        volumes={"/models": models},
        timeout=3600,
        secrets=[modal.Secret.from_name(os.environ.get("MODAL_HF_SECRET", "rosdiff-hf"))],
    )
    def download_cosmos() -> str:
        """One-off, CPU only: `modal run -m worker.modal_app::download_cosmos` fills the models volume."""
        from huggingface_hub import snapshot_download

        path = snapshot_download(COSMOS_MODEL, local_dir=f"/models/{COSMOS_MODEL}")
        models.commit()
        return path

    @app.cls(
        image=cosmos_image,
        gpu="L4",
        volumes={"/models": models},
        timeout=600,
        max_containers=1,
        scaledown_window=20,
    )
    class Reviewer:
        @modal.enter()
        def load(self):
            from worker.cosmos_review import CosmosModel

            self.model = CosmosModel(f"/models/{COSMOS_MODEL}")

        @modal.method()
        def review(self, job_input: dict) -> dict:
            import traceback

            from worker.cosmos_review import run_review_job

            try:
                return run_review_job({**job_input, "model": COSMOS_MODEL}, self.model)
            except Exception as e:
                return {"error": f"{type(e).__name__}: {e}", "traceback": traceback.format_exc()[-2000:]}

    @app.function(image=cosmos_image, timeout=600)
    def review_video(job_input: dict) -> dict:
        """Entry point the API spawns; hands the job to the (GPU, model-loaded) Reviewer."""
        return Reviewer().review.remote(job_input)


# --------------------------------------------------------------------------- LeRobot dataset export (CPU)

lerobot_image = (
    modal.Image.debian_slim(python_version="3.11")
    .apt_install("ffmpeg")
    .pip_install("lerobot==0.4.4", "imageio-ffmpeg==0.6.0")
    .add_local_python_source("worker")
)


@app.function(
    image=lerobot_image,
    cpu=2.0,
    memory=4096,
    timeout=900,
    max_containers=1,
    scaledown_window=20,
    # Only with MODAL_HF_SECRET=rosdiff-hf at deploy (a secret holding HF_TOKEN): datasets go to the Hub.
    # Without it, each dataset comes back as a zip stored in R2 next to the job.
    secrets=[modal.Secret.from_name(os.environ["MODAL_HF_SECRET"])] if os.environ.get("MODAL_HF_SECRET") else [],
)
def export_lerobot(job_input: dict) -> dict:
    """Runs -> LeRobot dataset, pushed to the Hub with the secret's HF_TOKEN (else returned as a zip)."""
    import traceback

    from worker.lerobot_export import run_export_job

    try:
        return run_export_job(job_input)
    except Exception as e:
        return {"error": f"{type(e).__name__}: {e}", "traceback": traceback.format_exc()[-2000:]}