"""Per-frame mask directories (`<tMs>.png`), as written by segmentPrompted."""

from __future__ import annotations

import bisect
import os
import re

import cv2
import numpy as np

from ..errors import BadRequest, MediaError

MASK_NAME = re.compile(r"^(\d+)\.png$")


class MaskIndex:
    def __init__(self, directory: str, frame_ms: float) -> None:
        entries = []
        for name in os.listdir(directory):
            match = MASK_NAME.match(name)
            if match:
                entries.append((int(match.group(1)), os.path.join(directory, name)))
        if not entries:
            raise BadRequest("mask directory contains no <tMs>.png masks")
        entries.sort()
        self.times = [t for t, _ in entries]
        self.files = [p for _, p in entries]
        gaps = np.diff(self.times) if len(self.times) > 1 else np.array([frame_ms])
        # A mask covers frames up to 1.5 sampling intervals away (at least one frame).
        self.tolerance = max(float(np.median(gaps)) * 1.5, frame_ms)
        self._cache: tuple[int, np.ndarray] | None = None

    def lookup(self, t_ms: int, width: int, height: int) -> np.ndarray | None:
        """Binary mask (uint8 0/255) for the frame at t_ms, resized to the frame, or None."""
        i = bisect.bisect_left(self.times, t_ms)
        best = None
        for j in (i - 1, i):
            if not 0 <= j < len(self.times) or abs(self.times[j] - t_ms) > self.tolerance:
                continue
            if best is None or abs(self.times[j] - t_ms) < abs(self.times[best] - t_ms):
                best = j
        if best is None:
            return None
        if self._cache is not None and self._cache[0] == best:
            return self._cache[1]
        mask = cv2.imread(self.files[best], cv2.IMREAD_GRAYSCALE)
        if mask is None:
            raise MediaError("internal", f"unreadable mask {os.path.basename(self.files[best])}")
        if mask.shape != (height, width):
            mask = cv2.resize(mask, (width, height), interpolation=cv2.INTER_NEAREST)
        mask = np.where(mask > 127, 255, 0).astype(np.uint8)
        self._cache = (best, mask)
        return mask