from __future__ import annotations

from pathlib import Path

from PIL import Image

try:
    import av
except Exception:  # pragma: no cover - environment dependent
    av = None  # type: ignore[assignment]


class VideoDecodeError(RuntimeError):
    pass


class FrameIndexError(VideoDecodeError):
    pass


class PyAVVideoDecoder:
    """Lazy decoder optimized for fast startup and on-demand frame access."""

    def __init__(self, video_path: Path, max_cache: int = 256) -> None:
        self.video_path = video_path
        self._max_cache = max_cache
        self._frames: list[Image.Image] = []
        self._cache_start = 0
        self._container = None
        self._stream = None
        self._decode_iter = None
        self._decoded_complete = False
        self._frame_count = 0
        self._fps: float | None = None
        self._open()

    @property
    def frame_count(self) -> int:
        frame_count = getattr(self, "_frame_count", 0)
        if frame_count:
            return frame_count
        frames = getattr(self, "_frames", [])
        cache_start = getattr(self, "_cache_start", 0)
        return cache_start + len(frames)

    @property
    def has_known_frame_count(self) -> bool:
        return bool(getattr(self, "_frame_count", 0))

    @property
    def fps(self) -> float | None:
        return self._fps

    def frame_index_for_seconds(self, seconds: float) -> int:
        if seconds < 0:
            raise ValueError("Seconds must be non-negative")
        if not self._fps:
            raise VideoDecodeError("FPS is unavailable for this video")
        return max(0, int(round(seconds * self._fps)))

    def _open(self) -> None:
        if av is None:
            raise VideoDecodeError("PyAV is not available in this environment")
        if not self.video_path.exists():
            raise VideoDecodeError(f"Video not found: {self.video_path}")

        try:
            self._container = av.open(str(self.video_path))
            self._stream = next(
                (s for s in self._container.streams if s.type == "video"), None
            )
            if self._stream is None:
                raise VideoDecodeError("No video stream found")
            if self._stream.average_rate:
                self._fps = float(self._stream.average_rate)
            self._frame_count = self._infer_frame_count()
            self._decode_iter = self._container.decode(self._stream)
        except VideoDecodeError:
            raise
        except Exception as exc:  # pragma: no cover - depends on media/backend
            raise VideoDecodeError(f"Unable to decode video: {exc}") from exc

    def get_frame(self, frame_index: int) -> Image.Image:
        frames = getattr(self, "_frames", [])
        decode_iter = getattr(self, "_decode_iter", None)
        decoded_complete = getattr(self, "_decoded_complete", False)
        frame_count = getattr(self, "_frame_count", 0)
        cache_start = getattr(self, "_cache_start", 0)
        max_cache = max(1, int(getattr(self, "_max_cache", 1)))
        if frame_index < 0:
            raise FrameIndexError(f"Frame index out of bounds: {frame_index}")
        if frame_count and frame_index >= frame_count:
            raise FrameIndexError(f"Frame index out of bounds: {frame_index}")
        cache_end = cache_start + len(frames)
        if cache_start <= frame_index < cache_end:
            return frames[frame_index - cache_start]
        if frame_index < cache_start:
            if not self._seek_to_frame(frame_index):
                self._reset_decode()
            frames = self._frames
            decode_iter = self._decode_iter
            decoded_complete = self._decoded_complete
            cache_start = self._cache_start
            cache_end = cache_start + len(frames)
            if cache_start <= frame_index < cache_end:
                return frames[frame_index - cache_start]
        if decode_iter is None:
            raise FrameIndexError(f"Frame index out of bounds: {frame_index}")

        if frame_index >= cache_end + max_cache:
            if self._seek_to_frame(frame_index):
                frames = self._frames
                decode_iter = self._decode_iter
                decoded_complete = self._decoded_complete
                cache_start = self._cache_start
                cache_end = cache_start + len(frames)
                if cache_start <= frame_index < cache_end:
                    return frames[frame_index - cache_start]

        while (cache_start + len(frames)) <= frame_index and not decoded_complete:
            try:
                frame = next(decode_iter)
            except StopIteration:
                decoded_complete = True
                self._decoded_complete = True
                break
            except Exception as exc:  # pragma: no cover - backend/media dependent
                raise VideoDecodeError(f"Unable to decode frame: {exc}") from exc
            image = frame.to_image().convert("RGB")
            pts = getattr(frame, "pts", None)
            if not frames and pts is not None and self._fps:
                time_base = getattr(frame, "time_base", None) or self._stream.time_base
                cache_start = self._frame_index_from_pts(pts, time_base)
                self._cache_start = cache_start
            frames.append(image)
            if len(frames) > max_cache:
                frames.pop(0)
                cache_start += 1
                self._cache_start = cache_start
            self._frames = frames

        cache_end = cache_start + len(frames)
        if cache_start <= frame_index < cache_end:
            return frames[frame_index - cache_start]

        if decoded_complete:
            self._frame_count = cache_start + len(frames)
            raise FrameIndexError(f"Frame index out of bounds: {frame_index}")

        raise VideoDecodeError(f"Unable to decode frame: {frame_index}")

    def _frame_index_from_pts(self, pts: int, time_base) -> int:
        if not self._fps or not time_base:
            return self._cache_start
        seconds = float(pts * time_base)
        return max(0, int(round(seconds * self._fps)))

    def _seek_to_frame(self, frame_index: int) -> bool:
        if self._container is None or self._stream is None or not self._fps:
            return False
        time_base = self._stream.time_base
        if not time_base:
            return False
        seconds = frame_index / self._fps
        offset = int(seconds / float(time_base))
        try:
            self._container.seek(
                offset,
                stream=self._stream,
                backward=True,
                any_frame=False,
            )
        except Exception:
            return False
        self._frames = []
        self._cache_start = 0
        self._decoded_complete = False
        self._decode_iter = self._container.decode(self._stream)
        return True

    def _infer_frame_count(self) -> int:
        if self._stream is None or self._container is None:
            return 0
        if self._stream.frames:
            return int(self._stream.frames)
        if self._stream.duration and self._stream.average_rate:
            seconds = float(self._stream.duration * self._stream.time_base)
            fps = float(self._stream.average_rate)
            if seconds > 0 and fps > 0:
                return max(1, int(round(seconds * fps)))
        if self._container.duration and self._stream.average_rate:
            seconds = float(self._container.duration) / 1_000_000.0
            fps = float(self._stream.average_rate)
            if seconds > 0 and fps > 0:
                return max(1, int(round(seconds * fps)))
        return 0

    def _reset_decode(self) -> None:
        if self._container is not None:
            self._container.close()
        self._frames = []
        self._cache_start = 0
        self._decoded_complete = False
        self._decode_iter = None
        self._container = None
        self._stream = None
        self._open()

    def close(self) -> None:
        if self._container is not None:
            self._container.close()
            self._container = None
            self._stream = None
            self._decode_iter = None
