from __future__ import annotations

import html
import json
import tempfile
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path
from typing import Any

import cv2
import fsspec
import gradio as gr
import numpy as np
import pyarrow.parquet as pq
from huggingface_hub import HfApi, hf_hub_download


GOLD_REPO = "macrodata/whats_going_on_bench"
GOLD_PARQUET = "whats_going_on_bench.parquet"
GOLD_INDEX = "whats_going_on_bench_index.parquet"
GOLD_PARQUET_URL = f"https://huggingface.co/datasets/{GOLD_REPO}/resolve/main/{GOLD_PARQUET}"
DEFAULT_RESULTS = (
    "hf://datasets/macrodata/whats_going_on_runs/runs/"
    "smoke_gemini_lite_3_retry/subtask_segmentation_eval"
)
IOU_THRESHOLD = 0.5
CACHE_DIR = Path(tempfile.gettempdir()) / "wasup_results_viewer"
CACHE_DIR.mkdir(parents=True, exist_ok=True)


@dataclass(frozen=True)
class Segment:
    start_sec: float
    end_sec: float
    label: str

    def valid(self) -> bool:
        return bool(self.label.strip()) and self.end_sec > self.start_sec


@dataclass(frozen=True)
class EpisodeScore:
    precision: float
    recall: float
    f1: float
    matches: int
    predicted: int
    gold: int


@lru_cache(maxsize=1)
def load_gold_index() -> dict[str, dict[str, Any]]:
    path = hf_hub_download(repo_id=GOLD_REPO, repo_type="dataset", filename=GOLD_INDEX)
    rows = pq.read_table(path).to_pylist()
    return {row["bench_id"]: row for row in rows}


@lru_cache(maxsize=1)
def gold_parquet() -> pq.ParquetFile:
    fs = fsspec.filesystem("https", block_size=2**20)
    return pq.ParquetFile(fs.open(GOLD_PARQUET_URL, "rb"))


@lru_cache(maxsize=512)
def load_gold_row(bench_id: str) -> dict[str, Any]:
    index = load_gold_index()
    if bench_id not in index:
        raise KeyError(f"{bench_id} is not in {GOLD_REPO}")
    row_group = int(index[bench_id]["row_group"])
    return gold_parquet().read_row_group(row_group).to_pylist()[0]


def load_result_rows(source: str) -> tuple[list[dict[str, Any]], str]:
    source = (source or "").strip() or DEFAULT_RESULTS
    files = resolve_parquet_files(source)
    table = pq.read_table(files if len(files) > 1 else files[0])
    rows = table.to_pylist()
    rows.sort(key=lambda row: (str(row.get("source_dataset") or ""), str(row.get("bench_id") or "")))
    return rows, f"Loaded {len(rows)} prediction rows from {source}"


def resolve_parquet_files(source: str) -> list[str]:
    if source.startswith("hf://datasets/"):
        return download_hf_parquets(source)

    path = Path(source).expanduser()
    if path.is_dir():
        files = sorted(str(item) for item in path.rglob("*.parquet"))
    else:
        files = [str(path)]
    if not files:
        raise FileNotFoundError(f"No parquet files found at {source}")
    return files


def download_hf_parquets(source: str) -> list[str]:
    suffix = source.removeprefix("hf://datasets/").strip("/")
    parts = suffix.split("/")
    if len(parts) < 2:
        raise ValueError("HF source must look like hf://datasets/owner/repo[/path]")
    repo_id = "/".join(parts[:2])
    prefix = "/".join(parts[2:]).strip("/")
    api = HfApi()
    repo_files = api.list_repo_files(repo_id=repo_id, repo_type="dataset")
    if prefix:
        matches = [item for item in repo_files if item == prefix or item.startswith(prefix.rstrip("/") + "/")]
    else:
        matches = repo_files
    parquet_files = sorted(item for item in matches if item.endswith(".parquet"))
    if not parquet_files:
        raise FileNotFoundError(f"No parquet files found under {source}")
    return [
        hf_hub_download(repo_id=repo_id, repo_type="dataset", filename=item)
        for item in parquet_files
    ]


def choices(rows: list[dict[str, Any]]) -> list[str]:
    items = []
    for idx, row in enumerate(rows):
        status = row.get("status") or "unknown"
        bench_id = row.get("bench_id") or "missing"
        source = row.get("source_dataset") or "unknown"
        instruction = row.get("instruction") or ""
        items.append(f"{idx:04d} | {status} | {source} | {bench_id} | {instruction}")
    return items


def parse_choice(choice: str | None) -> int:
    if not choice:
        return 0
    try:
        return int(choice.split("|", 1)[0].strip())
    except Exception:
        return 0


def load_results(source: str) -> tuple[list[dict[str, Any]], gr.Dropdown, str]:
    rows, message = load_result_rows(source)
    items = choices(rows)
    value = items[0] if items else None
    return rows, gr.Dropdown(choices=items, value=value), message


def write_video(row: dict[str, Any]) -> str:
    digest = str(row.get("primary_video_sha256") or row["bench_id"])
    path = CACHE_DIR / f'{row["bench_id"]}_{digest[:12]}.mp4'
    if not path.exists():
        path.write_bytes(row["primary_video"])
    return str(path)


def gold_segments(row: dict[str, Any]) -> list[Segment]:
    fps = float(row.get("fps") or 10.0)
    explicit = []
    for item in row.get("subtasks") or []:
        segment = segment_from_mapping(item, fps=fps)
        if segment and segment.valid():
            explicit.append(segment)
    if explicit:
        return explicit
    return collapse_frame_annotations(row.get("frame_annotations") or [], fps=fps)


def segment_from_mapping(item: dict[str, Any], fps: float | None = None) -> Segment | None:
    label = str(item.get("subtask") or item.get("label") or item.get("text") or "").strip()
    try:
        if "start_sec" in item and "end_sec" in item:
            return Segment(float(item["start_sec"]), float(item["end_sec"]), label)
        if fps and "start_frame" in item and "end_frame" in item:
            return Segment(float(item["start_frame"]) / fps, float(item["end_frame"]) / fps, label)
    except (TypeError, ValueError):
        return None
    return None


def collapse_frame_annotations(frames: list[dict[str, Any]], fps: float) -> list[Segment]:
    if not frames:
        return []
    ordered = sorted(frames, key=lambda item: int(item.get("frame_index") or 0))
    segments: list[Segment] = []
    start = int(ordered[0].get("frame_index") or 0)
    previous = start
    label = str(ordered[0].get("subtask") or "")
    for item in ordered[1:]:
        frame = int(item.get("frame_index") or 0)
        next_label = str(item.get("subtask") or "")
        if next_label != label:
            segment = Segment(start / fps, (previous + 1) / fps, label)
            if segment.valid():
                segments.append(segment)
            start = frame
            label = next_label
        previous = frame
    final = Segment(start / fps, (previous + 1) / fps, label)
    if final.valid():
        segments.append(final)
    return segments


def predicted_segments(row: dict[str, Any]) -> list[Segment]:
    raw = decode_json(row.get("predicted_subtasks_json"), [])
    segments = []
    for item in raw:
        if isinstance(item, dict):
            segment = segment_from_mapping(item)
            if segment and segment.valid():
                segments.append(segment)
    return sorted(segments, key=lambda item: (item.start_sec, item.end_sec))


def candidate_judgments(row: dict[str, Any]) -> list[dict[str, Any]]:
    raw = decode_json(row.get("candidate_judgments_json"), [])
    return [item for item in raw if isinstance(item, dict)]


def decode_json(value: Any, default: Any) -> Any:
    if value is None:
        return default
    if isinstance(value, (list, dict)):
        return value
    if isinstance(value, bytes):
        value = value.decode("utf-8")
    if not isinstance(value, str) or not value.strip():
        return default
    try:
        return json.loads(value)
    except json.JSONDecodeError:
        return default


def temporal_iou(left: Segment, right: Segment) -> float:
    intersection = max(0.0, min(left.end_sec, right.end_sec) - max(left.start_sec, right.start_sec))
    union = max(left.end_sec, right.end_sec) - min(left.start_sec, right.start_sec)
    return intersection / union if union > 0 else 0.0


def episode_score(gold: list[Segment], pred: list[Segment], judgments: list[dict[str, Any]]) -> tuple[EpisodeScore, set[tuple[int, int]]]:
    candidates: list[tuple[float, int, int]] = []
    for item in judgments:
        if not item.get("semantic_match"):
            continue
        try:
            pred_index = int(item["pred_segment_index"])
            gold_index = int(item["gt_segment_index"])
            iou = float(item["temporal_iou"])
        except (KeyError, TypeError, ValueError):
            continue
        if iou < IOU_THRESHOLD:
            continue
        if 0 <= pred_index < len(pred) and 0 <= gold_index < len(gold):
            candidates.append((iou, pred_index, gold_index))

    candidates.sort(reverse=True)
    used_pred: set[int] = set()
    used_gold: set[int] = set()
    matched_pairs: set[tuple[int, int]] = set()
    for _, pred_index, gold_index in candidates:
        if pred_index in used_pred or gold_index in used_gold:
            continue
        used_pred.add(pred_index)
        used_gold.add(gold_index)
        matched_pairs.add((pred_index, gold_index))

    matches = len(matched_pairs)
    precision = matches / len(pred) if pred else 0.0
    recall = matches / len(gold) if gold else 0.0
    f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0
    return EpisodeScore(precision, recall, f1, matches, len(pred), len(gold)), matched_pairs


def segment_at(segments: list[Segment], time_sec: float) -> str:
    for segment in segments:
        if segment.start_sec <= time_sec < segment.end_sec:
            return segment.label
    return ""


def render_frame(video_path: str, frame_index: int) -> np.ndarray | None:
    capture = cv2.VideoCapture(video_path)
    if not capture.isOpened():
        return None
    capture.set(cv2.CAP_PROP_POS_FRAMES, int(frame_index or 0))
    ok, frame = capture.read()
    capture.release()
    if not ok or frame is None:
        return None
    return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)


def timeline_html(gold: list[Segment], pred: list[Segment], matched: set[tuple[int, int]], duration: float) -> str:
    duration = max(duration, max([s.end_sec for s in gold + pred], default=1.0), 1.0)
    rows = [
        ("Gold", gold, "gold", set(range(len(gold)))),
        ("Pred", pred, "pred", {pred_index for pred_index, _ in matched}),
    ]
    parts = [
        "<style>",
        ".timeline{font-family:Arial,sans-serif}.lane{margin:10px 0}.lane-title{font-weight:700;margin-bottom:4px}",
        ".track{position:relative;height:46px;background:#f1f5f9;border:1px solid #cbd5e1;border-radius:6px;overflow:hidden}",
        ".seg{position:absolute;top:5px;height:34px;border-radius:5px;padding:3px 6px;box-sizing:border-box;font-size:12px;line-height:14px;overflow:hidden;color:#0f172a;border:1px solid rgba(15,23,42,.2)}",
        ".gold{background:#bfdbfe}.pred{background:#fecaca}.hit{outline:3px solid #22c55e}.miss{opacity:.8}",
        "</style><div class='timeline'>",
    ]
    for title, segments, klass, hit_indices in rows:
        parts.append(f"<div class='lane'><div class='lane-title'>{html.escape(title)}</div><div class='track'>")
        for idx, segment in enumerate(segments):
            left = max(0.0, min(100.0, 100 * segment.start_sec / duration))
            width = max(0.4, min(100.0 - left, 100 * (segment.end_sec - segment.start_sec) / duration))
            status = "hit" if idx in hit_indices and klass == "pred" else "miss"
            if klass == "gold":
                status = ""
            label = html.escape(f"{idx}: {segment.label} ({segment.start_sec:.1f}-{segment.end_sec:.1f}s)")
            parts.append(
                f"<div class='seg {klass} {status}' style='left:{left:.3f}%;width:{width:.3f}%' title='{label}'>{label}</div>"
            )
        parts.append("</div></div>")
    parts.append("</div>")
    return "".join(parts)


def segment_table(segments: list[Segment], kind: str, matched_pred: set[int] | None = None) -> list[list[Any]]:
    matched_pred = matched_pred or set()
    rows = []
    for idx, segment in enumerate(segments):
        matched = idx in matched_pred if kind == "pred" else ""
        rows.append([idx, segment.start_sec, segment.end_sec, segment.end_sec - segment.start_sec, segment.label, matched])
    return rows


def judgment_table(judgments: list[dict[str, Any]]) -> list[list[Any]]:
    rows = []
    for item in judgments:
        rows.append(
            [
                item.get("pred_segment_index"),
                item.get("gt_segment_index"),
                item.get("temporal_iou"),
                item.get("semantic_match"),
                item.get("pred_subtask"),
                item.get("gt_subtask"),
                item.get("judge_raw_output"),
            ]
        )
    rows.sort(key=lambda row: (row[0] if row[0] is not None else 9999, row[1] if row[1] is not None else 9999))
    return rows


def load_episode(rows: list[dict[str, Any]], choice: str | None, frame_index: int = 0) -> tuple[Any, ...]:
    if not rows:
        return empty_episode("No results loaded.")
    idx = max(0, min(parse_choice(choice), len(rows) - 1))
    pred_row = rows[idx]
    bench_id = pred_row.get("bench_id")
    if not bench_id:
        return empty_episode("Selected prediction row has no bench_id.")

    gold_row = load_gold_row(str(bench_id))
    video_path = write_video(gold_row)
    gold = gold_segments(gold_row)
    pred = predicted_segments(pred_row)
    judgments = candidate_judgments(pred_row)
    score, matched = episode_score(gold, pred, judgments)
    matched_pred = {pred_index for pred_index, _ in matched}
    duration = float(gold_row.get("duration_sec") or gold_row.get("num_frames") or 1)
    fps = float(gold_row.get("fps") or 10.0)
    max_frame = max(0, int(gold_row.get("num_frames") or 1) - 1)
    frame_index = max(0, min(int(frame_index or 0), max_frame))
    time_sec = frame_index / fps if fps else 0.0
    active = {
        "frame_index": frame_index,
        "time_sec": time_sec,
        "gold_subtask": segment_at(gold, time_sec),
        "predicted_subtask": segment_at(pred, time_sec),
    }
    summary = {
        "bench_id": bench_id,
        "source_dataset": gold_row.get("source_dataset"),
        "status": pred_row.get("status"),
        "error": pred_row.get("error"),
        "gold_instruction": gold_row.get("instruction"),
        "predicted_instruction": pred_row.get("predicted_instruction"),
        "annotation_model": pred_row.get("annotation_model"),
        "judge_model": pred_row.get("judge_model"),
        "score": score.__dict__,
    }
    status = (
        f"{idx + 1}/{len(rows)} | {bench_id} | "
        f"F1={score.f1:.3f}, P={score.precision:.3f}, R={score.recall:.3f}, "
        f"matches={score.matches}/{score.gold}, pred={score.predicted}"
    )
    slider = gr.Slider(value=frame_index, minimum=0, maximum=max_frame, step=1)
    return (
        video_path,
        status,
        summary,
        timeline_html(gold, pred, matched, duration),
        segment_table(gold, "gold"),
        segment_table(pred, "pred", matched_pred),
        judgment_table(judgments),
        slider,
        render_frame(video_path, frame_index),
        active,
        pred_row.get("raw_annotation_output") or "",
    )


def empty_episode(message: str) -> tuple[Any, ...]:
    return (
        None,
        message,
        {},
        "",
        [],
        [],
        [],
        gr.Slider(value=0, minimum=0, maximum=1, step=1),
        None,
        {},
        "",
    )


def step_episode(rows: list[dict[str, Any]], choice: str | None, delta: int) -> tuple[Any, ...]:
    items = choices(rows)
    if not items:
        return (gr.Dropdown(choices=[], value=None),) + empty_episode("No results loaded.")
    idx = max(0, min(parse_choice(choice) + delta, len(items) - 1))
    return (gr.Dropdown(choices=items, value=items[idx]),) + load_episode(rows, items[idx], 0)


def render_selected_frame(rows: list[dict[str, Any]], choice: str | None, frame_index: int) -> tuple[np.ndarray | None, dict[str, Any]]:
    if not rows:
        return None, {}
    idx = max(0, min(parse_choice(choice), len(rows) - 1))
    pred_row = rows[idx]
    gold_row = load_gold_row(str(pred_row["bench_id"]))
    video_path = write_video(gold_row)
    fps = float(gold_row.get("fps") or 10.0)
    gold = gold_segments(gold_row)
    pred = predicted_segments(pred_row)
    max_frame = max(0, int(gold_row.get("num_frames") or 1) - 1)
    frame_index = max(0, min(int(frame_index or 0), max_frame))
    time_sec = frame_index / fps if fps else 0.0
    return render_frame(video_path, frame_index), {
        "frame_index": frame_index,
        "time_sec": time_sec,
        "gold_subtask": segment_at(gold, time_sec),
        "predicted_subtask": segment_at(pred, time_sec),
    }


def step_frame(rows: list[dict[str, Any]], choice: str | None, frame_index: int, delta: int) -> tuple[gr.Slider, np.ndarray | None, dict[str, Any]]:
    if not rows:
        return gr.Slider(value=0, minimum=0, maximum=1, step=1), None, {}
    idx = max(0, min(parse_choice(choice), len(rows) - 1))
    gold_row = load_gold_row(str(rows[idx]["bench_id"]))
    max_frame = max(0, int(gold_row.get("num_frames") or 1) - 1)
    next_frame = max(0, min(int(frame_index or 0) + delta, max_frame))
    image, active = render_selected_frame(rows, choice, next_frame)
    return gr.Slider(value=next_frame, minimum=0, maximum=max_frame, step=1), image, active


with gr.Blocks(title="Wasup Results Viewer") as demo:
    gr.Markdown("# Wasup Results Viewer")
    gr.Markdown("Compare Refiner prediction parquet against gold `macrodata/whats_going_on_bench` segments.")

    rows_state = gr.State([])
    with gr.Row():
        results_path = gr.Textbox(label="Results parquet/folder/HF prefix", value=DEFAULT_RESULTS, scale=5)
        load_results_btn = gr.Button("Load results", scale=1)
    load_status = gr.Textbox(label="Load status", interactive=False)
    episode = gr.Dropdown(label="Episode", choices=[], value=None)

    with gr.Row():
        prev_btn = gr.Button("Previous")
        load_episode_btn = gr.Button("Load episode")
        next_btn = gr.Button("Next")

    video = gr.Video(label="Gold video")
    score_status = gr.Textbox(label="Episode score", interactive=False)
    metadata = gr.JSON(label="Episode metadata")
    timeline = gr.HTML(label="Gold/predicted timeline")

    with gr.Row():
        gold_table = gr.Dataframe(
            label="Gold segments",
            headers=["idx", "start_sec", "end_sec", "duration", "subtask", "matched"],
            interactive=False,
        )
        pred_table = gr.Dataframe(
            label="Predicted segments",
            headers=["idx", "start_sec", "end_sec", "duration", "subtask", "matched"],
            interactive=False,
        )

    judgments = gr.Dataframe(
        label="Candidate semantic judgments",
        headers=["pred_idx", "gold_idx", "temporal_iou", "semantic_match", "pred_subtask", "gold_subtask", "judge_raw_output"],
        interactive=False,
    )

    with gr.Row():
        frame_prev = gr.Button("Previous frame")
        frame_next = gr.Button("Next frame")
    frame_slider = gr.Slider(label="Frame", minimum=0, maximum=1, value=0, step=1)
    frame_image = gr.Image(label="Selected frame", interactive=False)
    active_frame = gr.JSON(label="Subtask at selected frame")
    raw_output = gr.Textbox(label="Raw annotation output", lines=8, interactive=False)

    episode_outputs = [
        video,
        score_status,
        metadata,
        timeline,
        gold_table,
        pred_table,
        judgments,
        frame_slider,
        frame_image,
        active_frame,
        raw_output,
    ]

    load_results_event = load_results_btn.click(load_results, inputs=results_path, outputs=[rows_state, episode, load_status])
    load_results_event.then(load_episode, inputs=[rows_state, episode], outputs=episode_outputs)
    load_episode_btn.click(load_episode, inputs=[rows_state, episode], outputs=episode_outputs)
    episode.change(load_episode, inputs=[rows_state, episode], outputs=episode_outputs)
    prev_btn.click(
        lambda rows, choice: step_episode(rows, choice, -1),
        inputs=[rows_state, episode],
        outputs=[episode] + episode_outputs,
    )
    next_btn.click(
        lambda rows, choice: step_episode(rows, choice, 1),
        inputs=[rows_state, episode],
        outputs=[episode] + episode_outputs,
    )
    frame_slider.change(render_selected_frame, inputs=[rows_state, episode, frame_slider], outputs=[frame_image, active_frame])
    frame_prev.click(
        lambda rows, choice, frame: step_frame(rows, choice, frame, -1),
        inputs=[rows_state, episode, frame_slider],
        outputs=[frame_slider, frame_image, active_frame],
    )
    frame_next.click(
        lambda rows, choice, frame: step_frame(rows, choice, frame, 1),
        inputs=[rows_state, episode, frame_slider],
        outputs=[frame_slider, frame_image, active_frame],
    )
    load_event = demo.load(load_results, inputs=results_path, outputs=[rows_state, episode, load_status])
    load_event.then(load_episode, inputs=[rows_state, episode], outputs=episode_outputs)


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