import os
import cv2
import numpy as np
import zarr
import subprocess
import tempfile
import shutil
from pathlib import Path
from huggingface_hub import hf_hub_download, HfFileSystem, list_repo_files
import gradio as gr

# ─────────────────────────────────────────────
# Constants
# ─────────────────────────────────────────────

EGOVERSE_REPO = "angkul07/egoverse-pick-tasks"
ABC_REPO      = "angkul07/abc-ego"
ABC_TASKS     = ["place_the_bread", "put_the_screwdriver_in_the_bin"]
# ABC_CAMERAS   = {
#     "/top-camera":         "Top",
#     "/left-wrist-camera":  "Left Wrist",
#     "/right-wrist-camera": "Right Wrist",
# }
FPS = 30

# Remove the static ABC_CAMERAS dict entirely

CAMERA_TOPIC_VARIANTS = [
    "/top-camera",
    "/top-left-camera", 
    "/top-right-camera",
    "/left-wrist-camera",
    "/right-wrist-camera",
]

TOPIC_LABELS = {
    "/top-camera":         "Top",
    "/top-left-camera":    "Top Left",
    "/top-right-camera":   "Top Right",
    "/left-wrist-camera":  "Left Wrist",
    "/right-wrist-camera": "Right Wrist",
}

# ─────────────────────────────────────────────
# UUID discovery helpers
# ─────────────────────────────────────────────

def list_egoverse_uuids(n: int = 50) -> list[str]:
    """Return up to n top-level UUID dirs from the egoverse repo."""
    fs = HfFileSystem()
    try:
        entries = fs.ls(f"datasets/{EGOVERSE_REPO}", detail=False)
        uuids = []
        for e in entries:
            name = Path(e).name
            # skip .gitattributes and similar
            if len(name) == 24 and name.isalnum():
                uuids.append(name)
            if len(uuids) >= n:
                break
        return sorted(uuids)
    except Exception as ex:
        print(f"[warn] Could not list egoverse uuids: {ex}")
        return []


def list_abc_uuids(task: str, n: int = 50) -> list[str]:
    fs = HfFileSystem()
    try:
        base = f"datasets/{ABC_REPO}/data/train/{task}"
        entries = fs.ls(base, detail=False)
        uuids = []
        for e in entries:
            name = Path(e).name
            if name.startswith("episode_"):
                uuids.append(name[len("episode_"):])  # strip prefix
            if len(uuids) >= n:
                break
        return sorted(uuids)
    except Exception as ex:
        print(f"[warn] Could not list abc uuids for {task}: {ex}")
        return []

def get_mcap_camera_topics(mcap_path: str) -> list[str]:
    """Return only the camera topics that actually exist in this mcap file."""
    with open(mcap_path, "rb") as f:
        reader = make_reader(f, decoder_factories=[DecoderFactory()])
        topics = {ch.topic for _, ch, _, _ in reader.iter_decoded_messages()}
    return [t for t in CAMERA_TOPIC_VARIANTS if t in topics]

# ─────────────────────────────────────────────
# Zarr helpers (egoverse)
# ─────────────────────────────────────────────

def unwrap_frame(frame):
    while isinstance(frame, np.ndarray) and frame.dtype == object:
        frame = frame.item()
    return frame


def decode_frame(frame):
    frame = unwrap_frame(frame)
    if isinstance(frame, np.ndarray):
        frame = frame.tobytes()
    buf = np.frombuffer(frame, dtype=np.uint8)
    img = cv2.imdecode(buf, cv2.IMREAD_COLOR)
    if img is None:
        raise RuntimeError("Failed to decode JPEG frame")
    return img

def convert_zarr_from_hf(uuid: str, tmp_dir: str) -> str:
    from huggingface_hub import snapshot_download

    yield_status = f"Downloading zarr for {uuid} ..."  # caller will show this

    # Download only the images.front_1 array for this episode
    local_dir = snapshot_download(
        repo_id=EGOVERSE_REPO,
        repo_type="dataset",
        allow_patterns=[f"{uuid}/images.front_1/**"],
        local_dir=os.path.join(tmp_dir, uuid),
    )

    array_path = os.path.join(local_dir, uuid, "images.front_1")
    root = zarr.open(array_path, mode="r", zarr_format=3)
    frames = root

    n_frames = frames.shape[0]
    first    = decode_frame(frames[0:1][0])
    h, w     = first.shape[:2]

    all_frames = frames[0:n_frames]  # single read, now local so fast

    temp_avi = os.path.join(tmp_dir, f"{uuid}.avi")
    output   = os.path.join(tmp_dir, f"{uuid}.mp4")

    writer = cv2.VideoWriter(
        temp_avi, cv2.VideoWriter_fourcc(*"MJPG"), FPS, (w, h)
    )
    for i in range(n_frames):
        writer.write(decode_frame(all_frames[i]))
    writer.release()

    subprocess.run(
        ["ffmpeg", "-y",
         "-i", temp_avi,
         "-c:v", "libx264", "-preset", "fast",
         "-crf", "18", "-pix_fmt", "yuv420p",
         "-movflags", "+faststart",
         output],
        check=True, capture_output=True,
    )
    os.remove(temp_avi)
    return output

# def convert_zarr_from_hf(uuid: str, tmp_dir: str) -> str:
#     """
#     Open images.front_1 directly from HF via fsspec, convert to mp4.
#     Returns path to the output mp4.
#     """
#     fs = HfFileSystem()
#     store_path = f"hf://datasets/{EGOVERSE_REPO}/{uuid}/images.front_1"

#     store = zarr.storage.FsspecStore.from_url(store_path)
#     root  = zarr.open(store, mode="r", zarr_format=3)

#     if "c" not in root and not hasattr(root, "shape"):
#         # root IS the array when opened at images.front_1
#         frames = root
#     else:
#         frames = root  # images.front_1 is the array itself

#     n_frames = frames.shape[0]
#     first    = decode_frame(frames[0])
#     h, w     = first.shape[:2]

#     temp_avi = os.path.join(tmp_dir, f"{uuid}.avi")
#     output   = os.path.join(tmp_dir, f"{uuid}.mp4")

#     # writer = cv2.VideoWriter(
#     #     temp_avi,
#     #     cv2.VideoWriter_fourcc(*"MJPG"),
#     #     FPS,
#     #     (w, h),
#     # )
#     # for i in range(n_frames):
#     #     writer.write(decode_frame(frames[i]))
#     # writer.release()

#     # Instead of fetching frame by frame
#     all_frames = frames[0:n_frames]  # single batched network read
    
#     writer = cv2.VideoWriter(temp_avi, cv2.VideoWriter_fourcc(*"MJPG"), FPS, (w, h))
#     for i in range(n_frames):
#         writer.write(decode_frame(all_frames[i]))
#     writer.release()

#     subprocess.run(
#         ["ffmpeg", "-y",
#          "-i", temp_avi,
#          "-c:v", "libx264", "-preset", "fast",
#          "-crf", "18", "-pix_fmt", "yuv420p",
#          "-movflags", "+faststart",
#          output],
#         check=True, capture_output=True,
#     )
#     os.remove(temp_avi)
#     return output

# ─────────────────────────────────────────────
# MCAP helpers (abc-teleop)
# ─────────────────────────────────────────────

try:
    from mcap.reader import make_reader
    from mcap_protobuf.decoder import DecoderFactory
    MCAP_AVAILABLE = True
except ImportError:
    MCAP_AVAILABLE = False
    print("[warn] mcap / mcap_protobuf not installed; abc-teleop conversion disabled.")


# def convert_mcap_camera(mcap_path: str, topic: str, tmp_dir: str, label: str) -> str:
#     """Extract one camera topic from mcap and encode to mp4. Returns output path."""
#     temp_h264 = os.path.join(tmp_dir, f"{label}.h264")
#     output    = os.path.join(tmp_dir, f"{label}.mp4")

#     with open(temp_h264, "wb") as out:
#         with open(mcap_path, "rb") as f:
#             reader = make_reader(f, decoder_factories=[DecoderFactory()])
#             for schema, channel, msg, decoded in reader.iter_decoded_messages():
#                 if channel.topic != topic:
#                     continue
#                 out.write(decoded.data)

#     subprocess.run(
#         ["ffmpeg", "-y",
#          "-framerate", str(FPS),
#          "-i", temp_h264,
#          "-c:v", "libx264", "-preset", "fast",
#          "-crf", "18", "-pix_fmt", "yuv420p",
#          "-movflags", "+faststart",
#          output],
#         check=True, capture_output=True,
#     )
#     os.remove(temp_h264)
#     return output

def convert_mcap_camera(mcap_path: str, topic: str, tmp_dir: str, label: str) -> str:
    temp_h264 = os.path.join(tmp_dir, f"{label}.h264")
    output    = os.path.join(tmp_dir, f"{label}.mp4")

    timestamps = []
    with open(temp_h264, "wb") as out:
        with open(mcap_path, "rb") as f:
            reader = make_reader(f, decoder_factories=[DecoderFactory()])
            for schema, channel, msg, decoded in reader.iter_decoded_messages():
                if channel.topic != topic:
                    continue
                timestamps.append(msg.log_time)  # nanoseconds
                out.write(decoded.data)

    # derive fps from median inter-frame interval
    if len(timestamps) >= 2:
        intervals = [timestamps[i+1] - timestamps[i] for i in range(len(timestamps) - 1)]
        median_ns  = sorted(intervals)[len(intervals) // 2]
        fps = round(1e9 / median_ns) if median_ns > 0 else FPS
    else:
        fps = FPS

    print(f"[{label}] detected {fps} fps from {len(timestamps)} frames")

    result = subprocess.run(
        ["ffmpeg", "-y",
         "-err_detect", "ignore_err",
         "-framerate", str(fps),
         "-i", temp_h264,
         "-c:v", "libx264", "-preset", "fast",
         "-crf", "18", "-pix_fmt", "yuv420p",
         "-movflags", "+faststart",
         output],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        print(f"[ffmpeg stderr] {result.stderr[-2000:]}")
        raise RuntimeError(f"ffmpeg failed (exit {result.returncode})")

    os.remove(temp_h264)
    return output


def download_mcap(task: str, episode_folder: str, tmp_dir: str) -> str:
    """Download just the episode.mcap file from HF to tmp_dir. Returns local path."""
    hf_path = f"data/train/{task}/{episode_folder}/episode.mcap"
    local = hf_hub_download(
        repo_id=ABC_REPO,
        repo_type="dataset",
        filename=hf_path,
        local_dir=tmp_dir,
    )
    return local

# ─────────────────────────────────────────────
# Core streaming generator
# ─────────────────────────────────────────────

def stream_videos(dataset: str, uuid_input: str, task: str):
    """
    Gradio generator. Yields (status_text, vid1, vid2, vid3) tuples progressively.
    vid2 and vid3 are only used for abc-teleop (3 cameras).
    """
    uuid = uuid_input.strip()
    if not uuid:
        yield "Enter a UUID to begin.", None, None, None
        return

    tmp_dir = tempfile.mkdtemp(prefix="viewer_")
    try:
        # ── Egoverse (zarr → single front camera) ───────────────────────
        # if dataset == "EgoVerse":
        #     yield f"Fetching zarr for {uuid} ...", None, None, None
        #     try:
        #         mp4 = convert_zarr_from_hf(uuid, tmp_dir)
        #         yield f"Done: {uuid}", mp4, None, None
        #     except Exception as e:
        #         yield f"Error: {e}", None, None, None

        if dataset == "EgoVerse":
            yield f"Downloading zarr for {uuid} ...", None, None, None
            try:
                mp4 = convert_zarr_from_hf(uuid, tmp_dir)
                yield f"Done: {uuid}", mp4, None, None
            except Exception as e:
                yield f"Error: {e}", None, None, None

        # ── ABC Teleop (mcap → 3 cameras, streamed one at a time) ───────
        elif dataset == "ABC Teleop":
            episode_folder = uuid if uuid.startswith("episode_") else f"episode_{uuid}"
        
            yield f"Downloading {episode_folder} / {task} ...", None, None, None
            try:
                mcap_path = download_mcap(task, episode_folder, tmp_dir)
            except Exception as e:
                yield f"Download failed: {e}", None, None, None
                return
        
            # discover which cameras this episode actually has
            topics = get_mcap_camera_topics(mcap_path)
            print(f"[{episode_folder}] found cameras: {topics}")
        
            videos = [None, None, None]
            for i, topic in enumerate(topics[:3]):   # max 3 panels
                label = TOPIC_LABELS[topic]
                yield f"Converting {label} camera ...", *videos
                try:
                    mp4 = convert_mcap_camera(mcap_path, topic, tmp_dir, label.replace(" ", "_").lower())
                    videos[i] = mp4
                except Exception as e:
                    yield f"Error on {label}: {e}", *videos
                    return
                yield f"Done: {label}", *videos
        
            yield f"All cameras ready: {episode_folder}", *videos

    finally:
        # cleanup is deferred — Gradio needs the file alive while streaming
        # Space restarts or the next call will clean /tmp automatically
        pass


# ─────────────────────────────────────────────
# UUID dropdown update callbacks
# ─────────────────────────────────────────────

def update_uuid_choices(dataset: str, task: str):
    """Refresh the UUID dropdown when dataset or task changes."""
    if dataset == "EgoVerse":
        choices = list_egoverse_uuids()
    else:
        choices = list_abc_uuids(task)
    return gr.update(choices=choices, value=None)


def toggle_task_visibility(dataset: str):
    return gr.update(visible=(dataset == "ABC Teleop"))


def toggle_camera_panels(dataset: str):
    show_multi = (dataset == "ABC Teleop")
    return (
        gr.update(visible=True),                # vid1 always visible
        gr.update(visible=show_multi, label="Left Wrist"),
        gr.update(visible=show_multi, label="Right Wrist"),
    )

# ─────────────────────────────────────────────
# UI
# ─────────────────────────────────────────────

CSS = """
/* ── Global ─────────────────────────────── */
@import url('https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:wght@400;600&family=IBM+Plex+Sans:wght@300;400;600&display=swap');

body, .gradio-container {
    background: #0d0f12 !important;
    color: #c8cdd6 !important;
    font-family: 'IBM Plex Sans', sans-serif !important;
}

/* ── Navbar ──────────────────────────────── */
#navbar {
    display: flex;
    align-items: center;
    gap: 16px;
    padding: 18px 24px;
    border-bottom: 1px solid #1f2937;
    background: #0d0f12;
    flex-wrap: wrap;
}

#brand {
    font-family: 'IBM Plex Mono', monospace;
    font-size: 13px;
    font-weight: 600;
    color: #5b8dee;
    letter-spacing: 0.08em;
    text-transform: uppercase;
    white-space: nowrap;
    margin-right: 8px;
}

/* ── Controls ────────────────────────────── */
.gr-dropdown > label > span,
.gr-textbox > label > span {
    font-family: 'IBM Plex Mono', monospace !important;
    font-size: 11px !important;
    color: #6b7280 !important;
    text-transform: uppercase;
    letter-spacing: 0.06em;
}

.gr-dropdown select,
.gr-textbox input,
.gr-textbox textarea {
    background: #141720 !important;
    border: 1px solid #1f2937 !important;
    color: #c8cdd6 !important;
    border-radius: 6px !important;
    font-family: 'IBM Plex Mono', monospace !important;
    font-size: 13px !important;
}

.gr-dropdown select:focus,
.gr-textbox input:focus {
    border-color: #5b8dee !important;
    outline: none !important;
    box-shadow: 0 0 0 2px rgba(91,141,238,0.15) !important;
}

/* ── Run button ──────────────────────────── */
#run-btn {
    background: #5b8dee !important;
    color: #0d0f12 !important;
    font-family: 'IBM Plex Mono', monospace !important;
    font-weight: 600 !important;
    font-size: 12px !important;
    text-transform: uppercase;
    letter-spacing: 0.08em;
    border: none !important;
    border-radius: 6px !important;
    padding: 10px 22px !important;
    cursor: pointer;
    transition: background 0.15s;
    white-space: nowrap;
}
#run-btn:hover { background: #4a7add !important; }

/* ── Status bar ─────────────────────────── */
#status-bar {
    padding: 10px 24px;
    font-family: 'IBM Plex Mono', monospace;
    font-size: 12px;
    color: #5b8dee;
    border-bottom: 1px solid #1f2937;
    min-height: 36px;
}

/* ── Video panels ────────────────────────── */
#video-row {
    display: flex;
    gap: 12px;
    padding: 20px 24px;
}

.video-panel {
    flex: 1;
    background: #141720;
    border: 1px solid #1f2937;
    border-radius: 8px;
    overflow: hidden;
}

.video-panel .panel-label {
    font-family: 'IBM Plex Mono', monospace;
    font-size: 11px;
    color: #6b7280;
    text-transform: uppercase;
    letter-spacing: 0.06em;
    padding: 10px 14px 6px;
}

.gr-video {
    background: #0a0c0f !important;
    border: none !important;
    border-radius: 0 !important;
}

/* scrollbar */
::-webkit-scrollbar { width: 6px; }
::-webkit-scrollbar-track { background: #0d0f12; }
::-webkit-scrollbar-thumb { background: #1f2937; border-radius: 3px; }
"""

with gr.Blocks(css=CSS, title="Dataset Viewer") as demo:

    # ── Navbar ────────────────────────────────────────────────────────
    with gr.Row(elem_id="navbar"):
        gr.HTML('<span id="brand">&#9632; Dataset Viewer</span>')

        dataset_dd = gr.Dropdown(
            choices=["EgoVerse", "ABC Teleop"],
            value="EgoVerse",
            label="Dataset",
            scale=1,
            min_width=160,
        )

        task_dd = gr.Dropdown(
            choices=ABC_TASKS,
            value=ABC_TASKS[0],
            label="Task",
            visible=False,
            scale=1,
            min_width=220,
        )

        uuid_dd = gr.Dropdown(
            choices=[],
            value=None,
            label="UUID",
            allow_custom_value=True,
            scale=3,
            min_width=340,
        )

        run_btn = gr.Button("Convert →", elem_id="run-btn", scale=0)

    # ── Status bar ────────────────────────────────────────────────────
    status = gr.Textbox(
        value="Select a dataset and UUID to begin.",
        show_label=False,
        interactive=False,
        elem_id="status-bar",
        lines=1,
    )

    # ── Video panels ──────────────────────────────────────────────────
    with gr.Row(elem_id="video-row"):
        vid1 = gr.Video(label="Front Camera",     show_label=True, visible=True)
        vid2 = gr.Video(label="Left Wrist",       show_label=True, visible=False)
        vid3 = gr.Video(label="Right Wrist",      show_label=True, visible=False)

    # ── Wiring ────────────────────────────────────────────────────────

    # When dataset changes: toggle task dropdown, reload UUIDs, toggle camera panels
    dataset_dd.change(
        fn=toggle_task_visibility,
        inputs=[dataset_dd],
        outputs=[task_dd],
    )
    dataset_dd.change(
        fn=update_uuid_choices,
        inputs=[dataset_dd, task_dd],
        outputs=[uuid_dd],
    )
    dataset_dd.change(
        fn=toggle_camera_panels,
        inputs=[dataset_dd],
        outputs=[vid1, vid2, vid3],
    )

    # When task changes (abc only): reload UUIDs
    task_dd.change(
        fn=update_uuid_choices,
        inputs=[dataset_dd, task_dd],
        outputs=[uuid_dd],
    )

    # On load: populate UUID dropdown for default dataset
    demo.load(
        fn=update_uuid_choices,
        inputs=[dataset_dd, task_dd],
        outputs=[uuid_dd],
    )

    # # Run button → streaming generator
    # run_btn.click(
    #     fn=stream_videos,
    #     inputs=[dataset_dd, uuid_dd, task_dd],
    #     outputs=[status, vid1, vid2, vid3],
    # )

    # At the bottom of the Blocks, add a cancel button next to run
    cancel_btn = gr.Button("Cancel", elem_id="cancel-btn", scale=0)
    
    # Update the click wiring
    click_event = run_btn.click(
        fn=stream_videos,
        inputs=[dataset_dd, uuid_dd, task_dd],
        outputs=[status, vid1, vid2, vid3],
        concurrency_limit=1,
    )
    
    cancel_btn.click(fn=None, cancels=[click_event])


if __name__ == "__main__":
    demo.queue()
    demo.launch()