"""Record tactile arrow overlays to video without stalling the pipeline.

Drawing the arrow grid costs a few milliseconds per frame (hundreds of
``cv2.arrowedLine`` calls) and encoding costs a few more, so doing it inline
would cut a stream's throughput. Instead:

* :meth:`StreamRecorder.on_frame` is designed to be the pipeline's ``on_frame``
  callback: it only performs a bounded, drop-oldest queue push, so the slowest
  pipeline stage is unaffected and the run never blocks on the encoder;
* a dedicated writer thread does the drawing and the ``cv2.VideoWriter`` calls;
* when the writer falls behind, the newest frames win and the rest are counted
  in :meth:`StreamRecorder.stats` so a recording gap is never silent.

Typical use::

    recorder = StreamRecorder("run/stream_0.mp4", label="#0",
                              valid_mask=lambda: runtime.valid_mask)
    pipe = StreamPipeline(sensor, produce, include=("image", "flow"),
                          on_frame=recorder.on_frame)
    pipe.start()
    ...
    print(recorder.close())        # written / dropped / fps

Per-stream files are the natural unit because each stream is its own process;
:func:`mosaic_videos` tiles them into one grid afterwards, which costs nothing
during the run.

Arrow parameters
----------------
``draw_arrows`` divides the sampled displacement by ``arrow_scale``, so the
defaults assume the finger's UV domain, where the active-region displacement is
only ~1-3 px: ``grid_spacing=14, arrow_scale=0.15, threshold=0.1`` magnifies it
~6.7x, which is what makes the direction field legible instead of a scatter of
dots (the ``draw_arrows`` defaults of ``arrow_scale=1.0, threshold=5`` render the
raw pixel displacement and are tuned for large-motion scenes). Length saturates
at ``grid_spacing * margin`` px, so colour - not length - carries the magnitude.
A box/mini stream at 365x365 needs its own calibration of ``arrow_scale``.
"""

from __future__ import annotations

import os
import queue
import threading
import time
from typing import Any, Callable, Iterable, Sequence

import cv2
import numpy as np

from .util import draw_arrows, draw_contact

__all__ = ["StreamRecorder", "mosaic_videos"]


class StreamRecorder:
    """Draw arrow overlays from snapshots and encode them to a video file."""

    def __init__(
        self,
        path: str,
        *,
        fps: float = 30.0,
        every: int = 1,
        fourcc: str = "mp4v",
        valid_mask: np.ndarray | Callable[[], np.ndarray | None] | None = None,
        grid_spacing: int = 14,
        threshold: float = 0.1,
        arrow_scale: float = 0.15,
        margin: float = 0.8,
        draw_contact: bool = False,
        label: str | None = None,
        overlay_values: bool = True,
        queue_depth: int = 8,
        verbose: bool = False,
    ) -> None:
        self.path = str(path)
        self.fps = max(1.0, float(fps))
        self.every = max(1, int(every))
        self.fourcc = fourcc
        # A callable is resolved once per frame: the SDK rebuilds
        # ``Sensor.valid_mask`` during ``get_img``, so holding an array taken
        # before the first frame would capture ``None``.
        self.valid_mask = valid_mask
        self.grid_spacing = int(grid_spacing)
        self.threshold = float(threshold)
        self.arrow_scale = float(arrow_scale)
        self.margin = float(margin)
        self.draw_contact_flag = bool(draw_contact)
        self.label = label
        self.overlay_values = bool(overlay_values)
        self.verbose = bool(verbose)

        self._queue: queue.Queue = queue.Queue(maxsize=max(1, int(queue_depth)))
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None
        self._seen = 0
        self._written = 0
        self._dropped = 0
        self._skipped = 0
        self._error: BaseException | None = None
        self._writer: cv2.VideoWriter | None = None
        self._size: tuple[int, int] | None = None

    # ------------------------------------------------------------------ ingest
    def on_frame(self, snap: Any) -> None:
        """Pipeline callback: enqueue one snapshot. Never blocks."""
        self._seen += 1
        if self.every > 1 and (self._seen - 1) % self.every:
            self._skipped += 1
            return
        while True:
            try:
                self._queue.put_nowait(snap)
                break
            except queue.Full:
                # Drop the oldest so the newest frame is always kept.
                try:
                    self._queue.get_nowait()
                    self._dropped += 1
                except queue.Empty:
                    pass
        self._ensure_thread()

    def _ensure_thread(self) -> None:
        if self._thread is None:
            self._thread = threading.Thread(target=self._loop, daemon=True,
                                            name="orisys-recorder")
            self._thread.start()

    # ------------------------------------------------------------------ writer
    def _open(self, frame: np.ndarray) -> None:
        h, w = frame.shape[:2]
        os.makedirs(os.path.dirname(os.path.abspath(self.path)) or ".", exist_ok=True)
        fourcc = cv2.VideoWriter_fourcc(*self.fourcc) if len(self.fourcc) == 4 else 0
        writer = cv2.VideoWriter(self.path, fourcc, self.fps, (w, h), True)
        if not writer.isOpened():
            raise RuntimeError(
                f"cv2.VideoWriter could not open {self.path!r} with fourcc={self.fourcc!r}"
            )
        self._writer = writer
        self._size = (w, h)
        if self.verbose:
            print(f"[recorder] {self.path} {w}x{h} @ {self.fps:g} fps fourcc={self.fourcc}")

    def _render(self, snap: Any) -> np.ndarray:
        if snap.image is None or snap.flow is None:
            raise RuntimeError(
                "recording needs image and flow in the snapshot; build the pipeline "
                'with include=("image", "flow")'
            )
        mask = self.valid_mask() if callable(self.valid_mask) else self.valid_mask
        image = draw_arrows(
            snap.image,
            snap.flow,
            threshold=self.threshold,
            grid_spacing=self.grid_spacing,
            arrow_scale=self.arrow_scale,
            margin=self.margin,
            valid_mask=mask,
        )
        if self.draw_contact_flag and snap.contour is not None and snap.centroid is not None:
            image = draw_contact(image, snap.contour, snap.centroid)
        if self.overlay_values or self.label:
            self._draw_hud(image, snap)
        return image

    def _draw_hud(self, image: np.ndarray, snap: Any) -> None:
        lines = []
        if self.label:
            lines.append(str(self.label))
        if self.overlay_values:
            lines.append(f"{snap.unit} fn={snap.normal_force:.3g} "
                         f"fx={snap.shear_x:.3g} fy={snap.shear_y:.3g}")
            lines.append(f"{snap.fps:.0f} FPS  {snap.age_ms:.1f} ms")
        y = 14
        for text in lines:
            cv2.putText(image, text, (4, y), cv2.FONT_HERSHEY_SIMPLEX, 0.38,
                        (0, 0, 0), 3, cv2.LINE_AA)
            cv2.putText(image, text, (4, y), cv2.FONT_HERSHEY_SIMPLEX, 0.38,
                        (255, 255, 255), 1, cv2.LINE_AA)
            y += 14

    def _loop(self) -> None:
        try:
            while True:
                try:
                    snap = self._queue.get(timeout=0.05)
                except queue.Empty:
                    if self._stop.is_set():
                        break
                    continue
                if snap is None:
                    break
                try:
                    frame = self._render(snap)
                    if self._writer is None:
                        self._open(frame)
                    self._writer.write(frame)
                    self._written += 1
                except BaseException as exc:
                    self._error = exc
                    break
        finally:
            if self._writer is not None:
                self._writer.release()
                self._writer = None

    # ------------------------------------------------------------------- close
    def close(self, timeout: float = 30.0) -> dict:
        """Drain, stop the writer thread and release the file."""
        if self._thread is not None:
            self._stop.set()
            # The stop flag is the real mechanism - the writer drains what is
            # queued and exits on its next empty poll. The sentinel only makes
            # that immediate, so it must never block: if the writer died on an
            # error the queue may be full and nobody would ever drain it.
            try:
                self._queue.put_nowait(None)
            except queue.Full:
                pass
            self._thread.join(timeout=timeout)
            alive = self._thread.is_alive()
            self._thread = None
        else:
            alive = False
        stats = self.stats()
        stats["thread_alive"] = alive
        if self._error is not None:
            stats["error"] = f"{type(self._error).__name__}: {self._error}"
            if self.verbose:
                print(f"[recorder] {self.path}: {stats['error']}")
        return stats

    def stats(self) -> dict:
        return {
            "path": self.path,
            "seen": self._seen,
            "written": self._written,
            "dropped": self._dropped,
            "skipped": self._skipped,
            "fps": self.fps,
            "fourcc": self.fourcc,
        }

    def __enter__(self) -> "StreamRecorder":
        return self

    def __exit__(self, exc_type, exc, tb) -> bool:
        self.close()
        return False


def mosaic_videos(
    paths: Sequence[str],
    out_path: str,
    *,
    cols: int | None = None,
    labels: Iterable[str] | None = None,
    scale: float = 1.0,
    gap: int = 4,
    fps: float | None = None,
    fourcc: str = "mp4v",
) -> dict:
    """Tile per-stream recordings into one grid video.

    Post-processing only: run it after the streams have finished, so it costs
    nothing while they are running. Shorter inputs hold their last frame.
    """
    paths = [str(p) for p in paths]
    if not paths:
        raise ValueError("mosaic_videos needs at least one input")
    for path in paths:
        if not os.path.isfile(path):
            raise FileNotFoundError(path)

    caps = [cv2.VideoCapture(p) for p in paths]
    try:
        for path, cap in zip(paths, caps):
            if not cap.isOpened():
                raise RuntimeError(f"cannot open {path!r}")
        tiles = []
        for cap in caps:
            ok, frame = cap.read()
            if not ok:
                raise RuntimeError("an input video has no frames")
            tiles.append(frame)

        if scale != 1.0:
            tiles = [cv2.resize(t, None, fx=scale, fy=scale,
                                interpolation=cv2.INTER_AREA) for t in tiles]
        th, tw = tiles[0].shape[:2]
        n = len(tiles)
        cols = int(cols) if cols else int(np.ceil(np.sqrt(n)))
        cols = max(1, min(cols, n))
        rows = int(np.ceil(n / cols))
        label_list = [str(x) for x in labels] if labels is not None else []

        out_h = rows * th + (rows - 1) * gap
        out_w = cols * tw + (cols - 1) * gap
        canvas = np.zeros((out_h, out_w, 3), dtype=np.uint8)
        positions = []
        for i in range(n):
            r, c = divmod(i, cols)
            positions.append((r * (th + gap), c * (tw + gap)))
        # Frame 0 is already in hand; seed the canvas with it so the output is
        # never short by one frame.
        for i, tile in enumerate(tiles):
            y, x = positions[i]
            canvas[y:y + th, x:x + tw] = tile

        out_fps = float(fps) if fps else float(caps[0].get(cv2.CAP_PROP_FPS) or 30.0)
        os.makedirs(os.path.dirname(os.path.abspath(out_path)) or ".", exist_ok=True)
        writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*fourcc), out_fps,
                                 (out_w, out_h), True)
        if not writer.isOpened():
            raise RuntimeError(f"cannot open {out_path!r} for writing")

        def emit() -> None:
            # Labels go on the canvas, not on the tiles: the read loop below
            # overwrites each tile region every frame.
            for i, text in enumerate(label_list):
                if i >= n:
                    break
                y, x = positions[i]
                org = (x + 4, y + th - 6)
                cv2.putText(canvas, text, org, cv2.FONT_HERSHEY_SIMPLEX, 0.45,
                            (0, 0, 0), 3, cv2.LINE_AA)
                cv2.putText(canvas, text, org, cv2.FONT_HERSHEY_SIMPLEX, 0.45,
                            (255, 255, 255), 1, cv2.LINE_AA)
            writer.write(canvas)

        written = 0
        try:
            # One frame per input per iteration; a spent input keeps its last
            # frame so the grid stays full until every input has run out.
            while True:
                emit()
                written += 1
                progressed = False
                for i, cap in enumerate(caps):
                    ok, frame = cap.read()
                    if not ok:
                        continue
                    progressed = True
                    if scale != 1.0:
                        frame = cv2.resize(frame, (tw, th), interpolation=cv2.INTER_AREA)
                    y, x = positions[i]
                    canvas[y:y + th, x:x + tw] = frame
                if not progressed:
                    break
        finally:
            writer.release()
    finally:
        for cap in caps:
            cap.release()

    return {
        "path": out_path,
        "frames": written,
        "fps": out_fps,
        "size": (out_w, out_h),
        "grid": (rows, cols),
        "scale": scale,
    }