from __future__ import annotations

import os
import tempfile
from contextlib import contextmanager
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, TypeAlias
from urllib.parse import urlparse

import av

from daft.daft import FileInfos, ImageMode
from daft.datatype import DataType
from daft.filesystem import _infer_filesystem, glob_path_with_stats
from daft.io import DataSource, DataSourceTask
from daft.recordbatch import RecordBatch
from daft.schema import Schema

if TYPE_CHECKING:
    from collections.abc import AsyncIterator, Generator
    from fractions import Fraction

    from av.video import VideoFrame

    from daft.daft import IOConfig
    from daft.dependencies import np
    from daft.io.pushdowns import Pushdowns

    _VideoFrameData: TypeAlias = np.typing.NDArray[Any]
else:
    _VideoFrameData: TypeAlias = Any


@dataclass
class _VideoFrame:
    """Represents a single video frame.

    Note:
        The field name 'data' is required due to a casting bug.
        See: https://github.com/Eventual-Inc/Daft/issues/4872
    """

    path: str
    frame_index: int
    frame_time: float
    frame_time_base: Fraction
    frame_pts: int
    frame_dts: int | None
    frame_duration: int | None
    is_key_frame: bool
    data: _VideoFrameData


@dataclass
class _VideoFramesSource(DataSource):
    """DataSource for streaming video files into rows of images.

    Note:
        This API requires [PyAV](https://github.com/PyAV-Org/PyAV), please do `pip install av`.

    Attributes:
        paths (list[str]): Path(s) to the video file(s).
        image_height (int): Height to which each frame will be resized.
        image_width (int): Width to which each frame will be resized.
        is_key_frame (bool|None): If True, only include key frames; if False, only non-key frames; if None, include all frames.
        io_config (IOConfig|None): Optional IOConfig.
    """

    paths: list[str]
    image_height: int
    image_width: int
    is_key_frame: bool | None = None
    io_config: IOConfig | None = None
    sample_interval_seconds: float | None = None

    @property
    def name(self) -> str:
        return "VideoFramesSource"

    @property
    def schema(self) -> Schema:
        return _schema(
            image_height=self.image_height,
            image_width=self.image_width,
        )

    def _list_file_infos(self) -> Generator[FileInfos]:
        for file_paths in self.paths:
            yield glob_path_with_stats(file_paths, file_format=None, io_config=self.io_config)

    def _list_file_paths(self) -> Generator[str]:
        all_youtube = all(_is_youtube_url(p) for p in self.paths)
        any_youtube = any(_is_youtube_url(p) for p in self.paths)
        if all_youtube:
            _assert_youtube_available()
            yield from self.paths
        elif any_youtube:
            raise ValueError("Either all or none of the paths must be YouTube URLs.")
        else:
            for file_infos in self._list_file_infos():
                yield from file_infos.file_paths

    async def get_tasks(self, pushdowns: Pushdowns) -> AsyncIterator[DataSourceTask]:
        for path in self._list_file_paths():
            yield _VideoFramesSourceTask(
                path=path,
                image_height=self.image_height,
                image_width=self.image_width,
                is_key_frame=self.is_key_frame,
                io_config=self.io_config,
                sample_interval_seconds=self.sample_interval_seconds,
            )


@dataclass
class _VideoFramesSourceTask(DataSourceTask):
    """DataSourceTask which yields record batches of images from a video file."""

    path: str
    image_height: int
    image_width: int
    is_key_frame: bool | None
    io_config: IOConfig | None
    sample_interval_seconds: float | None = None

    _max_partition_size = 10 * 1024 * 1024  # 10 MB

    @property
    def schema(self) -> Schema:
        return _schema(
            image_height=self.image_height,
            image_width=self.image_width,
        )

    def _list_frames(self, path: str, file: Any) -> Generator[_VideoFrame]:
        container = None
        try:
            container = av.open(file)

            # TODO support reading frames for multiple video streams
            stream = next((s for s in container.streams if s.type == "video"), None)
            if stream is None:
                container.close()
                raise RuntimeError(f"No video stream found in file: {path}")

            if self.is_key_frame:
                # https://pyav.org/docs/develop/cookbook/basics.html#saving-keyframes
                stream.codec_context.skip_frame = "NONKEY"

            sample_interval = self.sample_interval_seconds
            if sample_interval is not None and sample_interval <= 0:
                raise ValueError("sample_interval_seconds must be positive if provided")

            next_sample_time: float | None = 0.0 if sample_interval is not None else None
            # Small tolerance for floating point comparisons
            epsilon: float = 1e-9 if sample_interval is None else max(1e-9, sample_interval * 1e-6)

            frame_index: int = 0
            frame: VideoFrame
            while True:
                try:
                    frame = next(container.decode(stream))
                except av.EOFError:
                    break
                except StopIteration:
                    break

                frame_time = frame.time
                should_emit = True

                if sample_interval is not None:
                    if frame_time is None:
                        # Skip frames without timestamps when sampling is enabled.
                        should_emit = False
                    else:
                        assert next_sample_time is not None
                        if frame_time + epsilon < next_sample_time:
                            should_emit = False
                        else:
                            # Emit this frame and advance to the next target timestamp(s).
                            should_emit = True
                            while next_sample_time is not None and frame_time + epsilon >= next_sample_time:
                                next_sample_time += sample_interval

                if should_emit:
                    reformatted = frame.reformat(
                        width=self.image_width,
                        height=self.image_height,
                    )

                    yield _VideoFrame(
                        path=path,
                        frame_index=frame_index,
                        frame_time=frame_time,
                        frame_time_base=frame.time_base,
                        frame_pts=frame.pts,
                        frame_dts=frame.dts,
                        frame_duration=frame.duration,
                        is_key_frame=frame.key_frame,
                        data=reformatted.to_ndarray(format="rgb24"),
                    )

                frame_index += 1
        finally:
            if container:
                container.close()

    def _open(self) -> Any:
        if _is_youtube_url(self.path):
            return self._open_youtube_file()
        else:
            fp, fs, _ = _infer_filesystem(self.path, io_config=self.io_config)
            return fs.open_input_file(fp)

    @contextmanager
    def _open_youtube_file(self) -> Any:
        import yt_dlp

        def selector(ctx):  # type: ignore
            best_fmt = None
            best_fit = float("inf")  # lower is better

            for fmt in ctx["formats"]:
                if fmt.get("ext") != "mp4":
                    continue

                h, w = fmt.get("height"), fmt.get("width")
                if h is None or w is None:
                    continue

                fit = abs(h - self.image_height) + abs(w - self.image_width)
                if fit == 0:
                    yield fmt
                    return
                if fit < best_fit:
                    best_fmt = fmt
                    best_fit = fit

            yield best_fmt

        # Note:
        #   Streaming youtube downloads requires deeper work and was error-prone.
        #   The parsed youtube urls use m3u8 which requires decoding which ydl
        #   will handle for us. We cannot reliable pass a file-like HTTP response
        #   to PyAV; hence why this will download to a tempfile then open the file.

        temp_file = next(tempfile._get_candidate_names())  # type: ignore

        params = {
            "format": selector,
            "quiet": True,
            "outtmpl": temp_file,
            "no_warnings": True,
            "extract_flat": False,
            "no_check_certificate": True,
            "ignoreerrors": False,
            "consoletitle": False,
            "noprogress": True,
        }

        try:
            with yt_dlp.YoutubeDL(params) as ydl:
                ydl.download([self.path])
            yield open(temp_file, mode="rb")
        finally:
            if os.path.exists(temp_file):
                os.remove(temp_file)

    async def read(self) -> AsyncIterator[RecordBatch]:
        with self._open() as file:
            buffer = _VideoFramesBuffer(
                image_height=self.image_height,
                image_width=self.image_width,
            )
            for frame in self._list_frames(self.path, file):
                buffer.append(frame)
                if buffer.size() >= self._max_partition_size:
                    yield buffer.to_recordbatch()
                    buffer.clear()
            if buffer and buffer.size() > 0:
                yield buffer.to_recordbatch()


class _VideoFramesBuffer:
    """A record batch buffer/builder for video frames.

    This enables decoupling the video source from a particular library,
    making it possible for the VideoSource to leverage an Iterable[_VideoFrame[T]]
    at some later time. How the iterator is implemented (e.g. OpenCV vs PyAV)
    is not important, just that we have an iterable of frames which this builder
    and the source can stream as appropriately sized record batches.
    """

    image_height: int
    image_width: int

    _arr_path: list[str]
    _arr_frame_index: list[int]
    _arr_frame_time: list[float]
    _arr_frame_time_base: list[str]
    _arr_frame_pts: list[int]
    _arr_frame_dts: list[int | None]
    _arr_frame_duration: list[int | None]
    _arr_is_key_frame: list[bool]
    _arr_data: list[_VideoFrameData]
    _size_in_bytes: int
    _size_of_metadata = 64

    def __init__(self, image_height: int, image_width: int):
        self.image_height = image_height
        self.image_width = image_width
        self.clear()

    def clear(self) -> None:
        self._arr_path = []
        self._arr_frame_index = []
        self._arr_frame_time = []
        self._arr_frame_time_base = []
        self._arr_frame_pts = []
        self._arr_frame_dts = []
        self._arr_frame_duration = []
        self._arr_is_key_frame = []
        self._arr_data = []
        self._size_in_bytes = 0

    def size(self) -> int:
        """Returns the current size of the partition in bytes."""
        return self._size_in_bytes

    def append(self, frame: _VideoFrame) -> None:
        """Appends the frame to this partition builder.

        Note:
            We encode the time_base as a fraction string.
            See: https://github.com/Eventual-Inc/Daft/issues/4971
        """
        self._arr_path.append(frame.path)
        self._arr_frame_index.append(frame.frame_index)
        self._arr_frame_time.append(frame.frame_time)
        self._arr_frame_time_base.append(str(frame.frame_time_base))
        self._arr_frame_pts.append(frame.frame_pts)
        self._arr_frame_dts.append(frame.frame_dts)
        self._arr_frame_duration.append(frame.frame_duration)
        self._arr_is_key_frame.append(frame.is_key_frame)
        self._arr_data.append(frame.data)
        self._size_in_bytes += frame.data.nbytes + self._size_of_metadata

    def to_recordbatch(self) -> RecordBatch:
        """Returns a RecordBatch for this builder."""
        return RecordBatch.from_pydict(
            {
                "path": self._arr_path,
                "frame_index": self._arr_frame_index,
                "frame_time": self._arr_frame_time,
                "frame_time_base": self._arr_frame_time_base,
                "frame_pts": self._arr_frame_pts,
                "frame_dts": self._arr_frame_dts,
                "frame_duration": self._arr_frame_duration,
                "is_key_frame": self._arr_is_key_frame,
                "data": self._arr_data,
            }
        )


def _schema(image_height: int, image_width: int) -> Schema:
    """Returns the common schema which is needed in several places."""
    return Schema.from_pydict(
        {
            "path": DataType.string(),
            "frame_index": DataType.int64(),
            "frame_time": DataType.float64(),
            "frame_time_base": DataType.string(),
            "frame_pts": DataType.int64(),
            "frame_dts": DataType.int64(),
            "frame_duration": DataType.int64(),
            "is_key_frame": DataType.bool(),
            "data": DataType.image(
                height=image_height,
                width=image_width,
                mode=ImageMode.RGB,  # make configurable at later time
            ),
        }
    )


def _is_youtube_url(url: str) -> bool:
    hostname = urlparse(url).hostname
    if hostname is None:
        return False
    return hostname in {"www.youtube.com", "youtube.com", "youtu.be"}


def _assert_youtube_available() -> None:
    try:
        import yt_dlp  # noqa: F401
    except ImportError as e:
        raise ImportError(
            "YouTube video support requires the 'yt-dlp' package. Please install it with `pip install yt-dlp`."
        ) from e
