#!/usr/bin/env python3
"""GUI to label bifurcation yes/no for sampled frames across videos.

Designed for files like:
- evals/frame_bank_merged/new_bifurcation_samples_300.jsonl
"""

from __future__ import annotations

import argparse
import json
import os
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.widgets import Button, Slider

from deepivus.io.dicom import read_dicom


@dataclass
class Sample:
    sample_id: str
    group: str
    dicom_stem: str
    dicom_path: Path
    frame: int
    bifurcation: bool | None


def _read_jsonl(path: Path) -> list[dict[str, Any]]:
    rows = []
    with path.open("r", encoding="utf-8") as fp:
        for raw in fp:
            line = raw.strip()
            if not line:
                continue
            rows.append(json.loads(line))
    return rows


def _write_jsonl_atomic(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=path.parent, delete=False) as tmp:
        tmp_path = Path(tmp.name)
        for rec in rows:
            tmp.write(json.dumps(rec) + "\n")
        tmp.flush()
        os.fsync(tmp.fileno())
    os.replace(tmp_path, path)


def _load_samples(path: Path) -> tuple[dict[str, Any], list[Sample]]:
    rows = _read_jsonl(path)
    if not rows:
        raise ValueError(f"Empty file: {path}")
    meta = rows[0]
    samples: list[Sample] = []
    for rec in rows[1:]:
        if rec.get("record_type") != "frame":
            continue
        samples.append(
            Sample(
                sample_id=str(rec["sample_id"]),
                group=str(rec["group"]),
                dicom_stem=str(rec["dicom_stem"]),
                dicom_path=Path(rec["dicom_path"]),
                frame=int(rec["frame"]),
                bifurcation=rec.get("bifurcation", None),
            )
        )
    if not samples:
        raise ValueError(f"No frame records in {path}")
    return meta, samples


def _save_samples(path: Path, meta: dict[str, Any], samples: list[Sample]) -> None:
    rows: list[dict[str, Any]] = [meta]
    for s in samples:
        rows.append(
            {
                "record_type": "frame",
                "sample_id": s.sample_id,
                "group": s.group,
                "dicom_stem": s.dicom_stem,
                "dicom_path": str(s.dicom_path),
                "frame": int(s.frame),
                "bifurcation": s.bifurcation,
            }
        )
    _write_jsonl_atomic(path, rows)


class BifurcationSampleEditor:
    def __init__(self, samples_path: Path) -> None:
        self.samples_path = samples_path
        self.meta, self.samples = _load_samples(samples_path)
        self.current = 0
        self.is_updating_slider = False
        self.image_cache: dict[Path, np.ndarray] = {}

        self.fig = None
        self.ax_image = None
        self.image_artist = None
        self.slider = None
        self.status_text = None
        self.toggle_button = None
        self.save_button = None

    def _sample(self) -> Sample:
        return self.samples[self.current]

    def _load_frame(self, sample: Sample) -> np.ndarray:
        if sample.dicom_path not in self.image_cache:
            _, stack = read_dicom(str(sample.dicom_path))
            self.image_cache[sample.dicom_path] = stack
        stack = self.image_cache[sample.dicom_path]
        if sample.frame < 0 or sample.frame >= stack.shape[0]:
            raise IndexError(
                f"Frame {sample.frame} out of range for {sample.dicom_path} (num_frames={stack.shape[0]})"
            )
        return stack[sample.frame]

    def _autosave(self) -> None:
        _save_samples(self.samples_path, self.meta, self.samples)

    def _set_current(self, idx: int) -> None:
        self.current = int(np.clip(idx, 0, len(self.samples) - 1))
        self.is_updating_slider = True
        self.slider.set_val(self.current)
        self.is_updating_slider = False
        self._render()

    def _render(self) -> None:
        s = self._sample()
        frame = self._load_frame(s)
        if self.image_artist is None:
            self.image_artist = self.ax_image.imshow(frame, cmap="gray")
        else:
            self.image_artist.set_data(frame)

        label = s.bifurcation
        if label is True:
            txt = "Bifurcation: Yes"
            color = "#D6F5D6"
        elif label is False:
            txt = "Bifurcation: No"
            color = "#F5D6D6"
        else:
            txt = "Bifurcation: Unset"
            color = "#F5F1D6"

        self.toggle_button.label.set_text(txt)
        self.toggle_button.ax.set_facecolor(color)

        done = sum(1 for x in self.samples if x.bifurcation is not None)
        self.status_text.set_text(
            f"Sample {self.current + 1}/{len(self.samples)} ({done} labeled)  "
            f"{s.sample_id}  frame={s.frame}  "
            "Keys: left/right navigate, y=yes, n=no, u=unset, s=save"
        )
        self.fig.canvas.draw_idle()

    def _set_label(self, value: bool | None) -> None:
        self._sample().bifurcation = value
        self._autosave()
        self._render()

    def _on_slider(self, val: float) -> None:
        if self.is_updating_slider:
            return
        self._set_current(int(val))

    def _on_prev(self, _event: Any) -> None:
        self._set_current(self.current - 1)

    def _on_next(self, _event: Any) -> None:
        self._set_current(self.current + 1)

    def _on_toggle(self, _event: Any) -> None:
        cur = self._sample().bifurcation
        if cur is None:
            self._set_label(True)
        else:
            self._set_label(not cur)

    def _on_save(self, _event: Any) -> None:
        self._autosave()
        self._render()

    def _on_key(self, event: Any) -> None:
        if event.key in {"left", "a"}:
            self._on_prev(event)
        elif event.key in {"right", "d"}:
            self._on_next(event)
        elif event.key == "y":
            self._set_label(True)
        elif event.key == "n":
            self._set_label(False)
        elif event.key == "u":
            self._set_label(None)
        elif event.key == "s":
            self._on_save(event)

    def show(self) -> None:
        self.fig = plt.figure(figsize=(12, 9))
        self.fig.canvas.manager.set_window_title("Bifurcation Sample Labeler")
        self.ax_image = self.fig.add_axes([0.05, 0.22, 0.9, 0.74])
        self.ax_image.set_axis_off()
        self.ax_image.set_title("Bifurcation Frame Labeling", fontsize=14, weight="bold")

        slider_ax = self.fig.add_axes([0.12, 0.13, 0.76, 0.035])
        self.slider = Slider(
            ax=slider_ax,
            label="Sample",
            valmin=0,
            valmax=len(self.samples) - 1,
            valinit=0,
            valstep=1,
            color="#1f77b4",
        )

        prev_ax = self.fig.add_axes([0.12, 0.05, 0.1, 0.055])
        next_ax = self.fig.add_axes([0.24, 0.05, 0.1, 0.055])
        toggle_ax = self.fig.add_axes([0.38, 0.05, 0.2, 0.055])
        save_ax = self.fig.add_axes([0.62, 0.05, 0.12, 0.055])

        prev_button = Button(prev_ax, "Prev", color="#E0E0E0", hovercolor="#D0D0D0")
        next_button = Button(next_ax, "Next", color="#E0E0E0", hovercolor="#D0D0D0")
        self.toggle_button = Button(toggle_ax, "Bifurcation: Unset", color="#F5F1D6", hovercolor="#EDE6BF")
        self.save_button = Button(save_ax, "Save", color="#D6F5D6", hovercolor="#BFF0BF")

        self.status_text = self.fig.text(0.05, 0.18, "", fontsize=10)

        self.slider.on_changed(self._on_slider)
        prev_button.on_clicked(self._on_prev)
        next_button.on_clicked(self._on_next)
        self.toggle_button.on_clicked(self._on_toggle)
        self.save_button.on_clicked(self._on_save)
        self.fig.canvas.mpl_connect("key_press_event", self._on_key)

        self._render()
        plt.show()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--samples-jsonl",
        type=Path,
        default=Path("evals/frame_bank_merged/new_bifurcation_samples_300.jsonl"),
        help="Path to sampled frames JSONL.",
    )
    args = parser.parse_args()

    editor = BifurcationSampleEditor(args.samples_jsonl)
    editor.show()


if __name__ == "__main__":
    main()