"""视频读取模块"""

import cv2
from pathlib import Path
from typing import Iterator, Optional
import numpy as np


class VideoReader:
    """读取视频文件，按帧提取"""

    def __init__(self, path: str | Path):
        self.path = Path(path)
        self._cap: Optional[cv2.VideoCapture] = None
        self._frame_count: Optional[int] = None
        self._fps: Optional[float] = None
        self._width: Optional[int] = None
        self._height: Optional[int] = None

    @property
    def frame_count(self) -> int:
        if self._frame_count is None:
            self._open()
        return self._frame_count

    @property
    def fps(self) -> float:
        if self._fps is None:
            self._open()
        return self._fps

    @property
    def width(self) -> int:
        if self._width is None:
            self._open()
        return self._width

    @property
    def height(self) -> int:
        if self._height is None:
            self._open()
        return self._height

    def _open(self) -> None:
        self._cap = cv2.VideoCapture(str(self.path))
        if not self._cap.isOpened():
            raise FileNotFoundError(f"Cannot open video: {self.path}")
        self._frame_count = int(self._cap.get(cv2.CAP_PROP_FRAME_COUNT))
        self._fps = self._cap.get(cv2.CAP_PROP_FPS)
        self._width = int(self._cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self._height = int(self._cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    def read_frame(self, index: int) -> Optional[np.ndarray]:
        """读取指定索引的一帧 (BGR) — 随机访问，可能慢"""
        if self._cap is None:
            self._open()
        self._cap.set(cv2.CAP_PROP_POS_FRAMES, index)
        ret, frame = self._cap.read()
        if not ret:
            return None
        return frame

    def reset(self) -> None:
        """重开视频，回到第 0 帧，用于顺序读取"""
        self.close()
        self._open()

    def read_next(self) -> Optional[np.ndarray]:
        """顺序读取下一帧（不 seek，比 read_frame 可靠）"""
        if self._cap is None:
            self._open()
        ret, frame = self._cap.read()
        if not ret:
            return None
        return frame

    def iter_frames(self) -> Iterator[tuple[int, np.ndarray]]:
        """逐帧迭代，返回 (index, frame)"""
        self.reset()
        idx = 0
        while True:
            frame = self.read_next()
            if frame is None:
                break
            yield idx, frame
            idx += 1

    def close(self) -> None:
        if self._cap is not None:
            self._cap.release()
            self._cap = None

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()