"""segmentPrompted: SAM2 (Apache-2.0) from a click or box, propagated over the video.

Frames are sampled at up to 15 fps and fed to the SAM2 video predictor in
chunks so memory stays bounded; each chunk is re-prompted with the bounding
box of the previous chunk's last mask. Masks are written as `<tMs>.png`
(0 or 255) at the source resolution.
"""

from __future__ import annotations

import os
import tempfile
import threading
from collections.abc import Iterator
from typing import Any

import cv2
import numpy as np

from .. import paths
from ..errors import BadRequest, Cancelled, MediaError
from ..media_io import FFMPEG, run_cancellable, scaled_size
from ..schemas import SegmentPromptedRequest
from .common import Ctx, clamp01, norm_box, open_input
from .devices import torch_device
from .masks import MASK_NAME

SEGMENT_MAX_FPS = 15.0
SEGMENT_MAX_SIDE = 1024
CHUNK_FRAMES = 64

MaskStream = Iterator[tuple[int, np.ndarray, float]]


def _mask_box(mask: np.ndarray) -> tuple[int, int, int, int] | None:
    ys, xs = np.nonzero(mask)
    if xs.size == 0:
        return None
    return int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1


class Sam2Segmenter:
    def __init__(self, checkpoint: str, settings: Any) -> None:
        import torch
        from sam2.build_sam import build_sam2_video_predictor

        self.torch = torch
        self.device = torch_device(settings)
        self.predictor = build_sam2_video_predictor(
            "configs/sam2.1/sam2.1_hiera_t.yaml", checkpoint, device=self.device
        )
        self._lock = threading.Lock()

    def _chunk(
        self,
        frames: list[str],
        prompt_at: int,
        box: tuple[int, int, int, int] | None,
        point: tuple[int, int] | None,
        reverse: bool,
        cancel: threading.Event,
    ) -> MaskStream:
        with tempfile.TemporaryDirectory(prefix="sam2-") as chunk_dir:
            for i, source in enumerate(frames):
                os.symlink(source, os.path.join(chunk_dir, f"{i}.jpg"))
            torch = self.torch
            with torch.inference_mode():
                state = self.predictor.init_state(video_path=chunk_dir, offload_video_to_cpu=True)
                kwargs: dict[str, Any] = {}
                if box is not None:
                    kwargs["box"] = np.array(box, dtype=np.float32)
                if point is not None:
                    kwargs["points"] = np.array([point], dtype=np.float32)
                    kwargs["labels"] = np.array([1], dtype=np.int32)
                self.predictor.add_new_points_or_box(state, frame_idx=prompt_at, obj_id=1, **kwargs)
                for index, _, logits in self.predictor.propagate_in_video(
                    state, start_frame_idx=prompt_at, reverse=reverse
                ):
                    if cancel.is_set():
                        raise Cancelled()
                    values = logits[0, 0].float().cpu().numpy()
                    mask = values > 0
                    confidence = float(1.0 / (1.0 + np.exp(-values[mask])).mean()) if mask.any() else 0.0
                    yield index, mask, confidence
                self.predictor.reset_state(state)

    def segment(
        self,
        frames: list[str],
        prompt_index: int,
        box: tuple[int, int, int, int] | None,
        point: tuple[int, int] | None,
        propagate: bool,
        cancel: threading.Event,
    ) -> MaskStream:
        with self._lock:
            if not propagate:
                for _, mask, confidence in self._chunk([frames[prompt_index]], 0, box, point, False, cancel):
                    yield prompt_index, mask, confidence
                return
            yield from self._forward(frames, prompt_index, box, point, cancel)
            yield from self._backward(frames, prompt_index, box, point, cancel)

    def _forward(
        self,
        frames: list[str],
        anchor: int,
        box: tuple[int, int, int, int] | None,
        point: tuple[int, int] | None,
        cancel: threading.Event,
    ) -> MaskStream:
        emit_anchor = True
        while True:
            chunk = frames[anchor : anchor + CHUNK_FRAMES]
            last: np.ndarray | None = None
            for index, mask, confidence in self._chunk(chunk, 0, box, point, False, cancel):
                if index == 0 and not emit_anchor:
                    continue
                last = mask
                yield anchor + index, mask, confidence
            end = anchor + len(chunk) - 1
            box = _mask_box(last) if last is not None else None
            if end >= len(frames) - 1 or box is None:
                return
            anchor, point, emit_anchor = end, None, False

    def _backward(
        self,
        frames: list[str],
        anchor: int,
        box: tuple[int, int, int, int] | None,
        point: tuple[int, int] | None,
        cancel: threading.Event,
    ) -> MaskStream:
        while anchor > 0:
            start = max(0, anchor - CHUNK_FRAMES + 1)
            last: np.ndarray | None = None
            for index, mask, confidence in self._chunk(
                frames[start : anchor + 1], anchor - start, box, point, True, cancel
            ):
                if start + index == anchor:
                    continue  # emitted by the forward pass
                last = mask
                yield start + index, mask, confidence
            box = _mask_box(last) if last is not None else None
            if start == 0 or box is None:
                return
            anchor, point = start, None


def load_segmenter(files: dict[str, str], settings: Any) -> Sam2Segmenter:
    return Sam2Segmenter(files["sam2.1_hiera_tiny.pt"], settings)


def _extract_frames(
    path: str,
    fps: float,
    size: tuple[int, int],
    directory: str,
    cancel: threading.Event,
    at_ms: int | None = None,
) -> list[str]:
    argv = [FFMPEG, "-v", "error", "-nostdin"]
    if at_ms is not None:
        argv += ["-ss", f"{at_ms / 1000:.3f}"]
    argv += ["-i", path, "-map", "0:v:0", "-an", "-sn"]
    filters = [] if at_ms is not None else [f"fps={fps:.6f}"]
    filters.append(f"scale={size[0]}:{size[1]}:flags=area")
    argv += ["-vf", ",".join(filters)]
    if at_ms is not None:
        argv += ["-frames:v", "1"]
    argv += ["-q:v", "2", "-start_number", "0", os.path.join(directory, "%06d.jpg")]
    result = run_cancellable(argv, cancel)
    if result.returncode != 0:
        raise MediaError("internal", "frame extraction failed")
    names = sorted(n for n in os.listdir(directory) if n.endswith(".jpg"))
    if not names:
        raise MediaError("internal", "no frames extracted")
    return [os.path.join(directory, n) for n in names]


def run(ctx: Ctx, body: SegmentPromptedRequest, cancel: threading.Event) -> dict[str, Any]:
    real, info = open_input(ctx, body.path, cancel)
    out_dir = paths.output_dir(ctx.settings.media_root, body.outDir)
    segmenter = ctx.registry.runtime("segmentPrompted")
    fps = min(info.fps, SEGMENT_MAX_FPS)
    width, height = scaled_size(info.width, info.height, SEGMENT_MAX_SIDE)
    prompt = body.prompt
    if prompt.tMs > int(info.duration_s * 1000) + 1:
        raise BadRequest("prompt.tMs is after the end of the video")
    with tempfile.TemporaryDirectory(prefix="segment-") as work:
        if body.propagate:
            frames = _extract_frames(real, fps, (width, height), work, cancel)
            times = [int(round(i * 1000.0 / fps)) for i in range(len(frames))]
            prompt_index = int(np.argmin([abs(t - prompt.tMs) for t in times]))
        else:
            frames = _extract_frames(real, fps, (width, height), work, cancel, at_ms=prompt.tMs)
            times = [prompt.tMs]
            prompt_index = 0
        box = None
        if prompt.box is not None:
            b = prompt.box
            box = (
                int(round(b.x * width)),
                int(round(b.y * height)),
                int(round(clamp01(b.x + b.w) * width)),
                int(round(clamp01(b.y + b.h) * height)),
            )
        point = (
            (int(round(prompt.point.x * width)), int(round(prompt.point.y * height)))
            if prompt.point
            else None
        )

        # Replace masks from an earlier run in the same directory.
        for name in os.listdir(out_dir):
            if MASK_NAME.match(name):
                os.unlink(os.path.join(out_dir, name))
        keyframes = []
        for index, mask, confidence in segmenter.segment(
            frames, prompt_index, box, point, body.propagate, cancel
        ):
            if cancel.is_set():
                raise Cancelled()
            t_ms = times[index]
            full = cv2.resize(
                mask.astype(np.uint8) * 255, (info.width, info.height), interpolation=cv2.INTER_NEAREST
            )
            found = _mask_box(full > 0)
            if found is None:
                continue
            if not cv2.imwrite(os.path.join(out_dir, f"{t_ms}.png"), full):
                raise MediaError("internal", "could not write mask")
            x0, y0, x1, y1 = found
            keyframes.append(
                {
                    "tMs": t_ms,
                    "box": norm_box(x0, y0, x1 - x0, y1 - y0, info.width, info.height),
                    "confidence": clamp01(confidence),
                }
            )
    keyframes.sort(key=lambda k: k["tMs"])
    score = float(np.mean([k["confidence"] for k in keyframes])) if keyframes else 0.0
    return {"maskDir": body.outDir, "keyframes": keyframes, "score": score}