"""
RoboMind VLA — Task 1: rollout video + ground-truth label generation (MINIMAL).

Runs ENTIRELY on Modal (CPU container, no GPU). Nothing runs locally.

What it does
------------
1. Downloads one Minari offline-RL dataset (`mujoco/humanoid/expert-v0`).
2. For each of N episodes, reconstructs the simulator state from the stored
   observations, sets it on a recovered MuJoCo env, and renders RGB frames
   headlessly via OSMesa (software rendering -> no GPU needed).
3. Writes an `.mp4` per episode plus a `metadata.jsonl` line carrying the
   GROUND-TRUTH labels (episode return, fell flag, quality tier, #steps).

Why reconstruct state instead of replaying actions open-loop:
    The humanoid is chaotic; open-loop action replay from a fresh reset diverges
    instantly and destroys the quality spectrum. Gymnasium's Humanoid observation
    contains qpos[2:] and qvel, so we can rebuild (qpos, qvel) exactly and
    set_state() each step -> the rendered video matches the dataset trajectory,
    and the dataset's reward/termination labels stay valid for that video.

Run it:
    modal run data_gen_modal.py
    # then inspect the volume:
    modal volume ls robomind-data rollouts

Done when:
    `modal volume ls robomind-data rollouts` shows a few .mp4 files and a
    metadata.jsonl with one line per rendered episode.
"""

from __future__ import annotations

import modal

# --- Dataset to render in this minimal task -------------------------------
DATASET_ID = "mujoco/humanoid/expert-v0"   # tier="expert" -> high-quality rollouts
N_EPISODES = 5                              # keep tiny for the first run
RENDER_FPS = 30

# --- Modal image: MuJoCo + OSMesa (CPU software rendering) + Minari --------
image = (
    modal.Image.debian_slim(python_version="3.11")
    .apt_install(
        "libosmesa6",
        "libosmesa6-dev",
        "libgl1-mesa-glx",
        "libglfw3",
        "libglew2.2",
        "patchelf",
        "ffmpeg",
    )
    .pip_install(
        "minari[all]",
        "gymnasium[mujoco]",
        "mujoco>=3.1.0",
        "imageio",
        "imageio-ffmpeg",
        "numpy",
    )
    # Force headless software rendering -> no GPU required.
    .env({"MUJOCO_GL": "osmesa", "PYOPENGL_PLATFORM": "osmesa"})
)

app = modal.App("robomind-vla-data")

# Persisted output so later tasks (dataset build, fine-tune) can read it.
volume = modal.Volume.from_name("robomind-data", create_if_missing=True)
OUT_DIR = "/data/rollouts"


@app.function(image=image, volumes={"/data": volume}, timeout=3600)
def generate(dataset_id: str = DATASET_ID, n_episodes: int = N_EPISODES) -> dict:
    import json
    import os

    import imageio
    import minari
    import numpy as np

    os.makedirs(OUT_DIR, exist_ok=True)

    # tier is the last path segment minus the version suffix, e.g. "expert-v0" -> "expert"
    tier = dataset_id.split("/")[-1].split("-v")[0]
    env_name = dataset_id.split("/")[1]  # "humanoid"

    print(f"[data] loading Minari dataset: {dataset_id}")
    dataset = minari.load_dataset(dataset_id, download=True)
    env = dataset.recover_environment(render_mode="rgb_array")

    nq = int(env.unwrapped.model.nq)   # full position dim (humanoid: 24)
    nv = int(env.unwrapped.model.nv)   # velocity dim (humanoid: 23)
    # Gymnasium MuJoCo excludes the root x,y from the observation by default,
    # so the obs starts at qpos[2:]. Position part length = nq - 2.
    pos_len = nq - 2
    print(f"[data] env={env_name} nq={nq} nv={nv} (obs pos_len={pos_len})")

    manifest = []
    written = 0
    for ep in dataset.iterate_episodes():
        if written >= n_episodes:
            break

        obs = np.asarray(ep.observations)        # shape (T+1, obs_dim)
        rewards = np.asarray(ep.rewards, dtype=float)
        terminations = np.asarray(ep.terminations, dtype=bool)
        truncations = np.asarray(ep.truncations, dtype=bool)

        frames = []
        env.reset()
        for t in range(obs.shape[0]):
            o = obs[t]
            qpos = np.concatenate([[0.0, 0.0], o[:pos_len]])     # restore root x,y as 0
            qvel = o[pos_len:pos_len + nv]
            env.unwrapped.set_state(qpos, qvel)
            frame = env.render()
            if frame is not None:
                frames.append(np.asarray(frame))

        if not frames:
            print(f"[data] WARNING: no frames rendered for episode {ep.id}, skipping")
            continue

        fell = bool(terminations.any() and not truncations.all())
        ep_return = float(rewards.sum())
        n_steps = int(rewards.shape[0])

        vid_name = f"{env_name}_{tier}_ep{ep.id}.mp4"
        vid_path = os.path.join(OUT_DIR, vid_name)
        imageio.mimwrite(vid_path, frames, fps=RENDER_FPS, macro_block_size=None)

        record = {
            "video": vid_name,
            "env": env_name,
            "tier": tier,
            "episode_id": int(ep.id),
            "num_steps": n_steps,
            "return": ep_return,
            "fell": fell,
        }
        manifest.append(record)
        written += 1
        print(f"[data] wrote {vid_name}  steps={n_steps} return={ep_return:.1f} fell={fell}")

    # Append to a JSONL manifest (one line per episode).
    meta_path = os.path.join(OUT_DIR, "metadata.jsonl")
    with open(meta_path, "a") as f:
        for r in manifest:
            f.write(json.dumps(r) + "\n")

    volume.commit()  # persist writes to the Modal volume
    print(f"[data] done: {len(manifest)} episodes -> {OUT_DIR}")
    return {"written": len(manifest), "dataset_id": dataset_id, "out_dir": OUT_DIR}


@app.local_entrypoint()
def main():
    result = generate.remote()
    print("RESULT:", result)