# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.


import io
import json
import os
import shutil
import sys
import warnings
from contextlib import nullcontext
from pathlib import Path

import torch
from torch.library import get_ctx, register_fake
from torchcodec._internally_replaced_utils import (  # @manual=//pytorch/torchcodec/src:internally_replaced_utils
    load_torchcodec_shared_libraries,
)


expose_ffmpeg_dlls = nullcontext
if sys.platform == "win32" and hasattr(os, "add_dll_directory"):
    # On windows we try to locate the FFmpeg DLLs and temporarily add them to
    # the DLL search path. This seems to be needed on some users machine, but
    # not on our CI. We don't know why.
    if ffmpeg_path := shutil.which("ffmpeg"):

        def expose_ffmpeg_dlls():  # noqa: F811
            ffmpeg_dir = Path(ffmpeg_path).parent.absolute()
            return os.add_dll_directory(str(ffmpeg_dir))  # that's the actual CM


with expose_ffmpeg_dlls():
    ffmpeg_major_version, core_library_path, _pybind_ops = (
        load_torchcodec_shared_libraries()
    )


# Note: We use disallow_in_graph because PyTorch does constant propagation of
# factory functions.
create_from_file = torch._dynamo.disallow_in_graph(
    torch.ops.torchcodec_ns.create_from_file.default
)
create_from_tensor = torch._dynamo.disallow_in_graph(
    torch.ops.torchcodec_ns.create_from_tensor.default
)
_create_from_file_like = torch._dynamo.disallow_in_graph(
    torch.ops.torchcodec_ns._create_from_file_like.default
)
_add_video_stream_raw = torch.ops.torchcodec_ns.add_video_stream.default
_add_video_stream = torch.ops.torchcodec_ns._add_video_stream.default


def add_video_stream(
    decoder: torch.Tensor,
    *,
    num_threads: int | None = None,
    dimension_order: str | None = None,
    stream_index: int | None = None,
    device: str = "cpu",
    device_variant: str = "default",
    transform_specs: str = "",
    custom_frame_mappings: (
        tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None
    ) = None,
    output_dtype: str = "uint8",
) -> None:
    custom_frame_mappings_pts: torch.Tensor | None = None
    custom_frame_mappings_keyframe_indices: torch.Tensor | None = None
    custom_frame_mappings_duration: torch.Tensor | None = None
    if custom_frame_mappings is not None:
        (
            custom_frame_mappings_pts,
            custom_frame_mappings_keyframe_indices,
            custom_frame_mappings_duration,
        ) = custom_frame_mappings
    _add_video_stream_raw(
        decoder,
        num_threads=num_threads,
        dimension_order=dimension_order,
        stream_index=stream_index,
        device=device,
        device_variant=device_variant,
        transform_specs=transform_specs,
        custom_frame_mappings_pts=custom_frame_mappings_pts,
        custom_frame_mappings_duration=custom_frame_mappings_duration,
        custom_frame_mappings_keyframe_indices=custom_frame_mappings_keyframe_indices,
        output_dtype=output_dtype,
    )


add_audio_stream = torch.ops.torchcodec_ns.add_audio_stream.default
seek_to_pts = torch.ops.torchcodec_ns.seek_to_pts.default
get_next_frame = torch.ops.torchcodec_ns.get_next_frame.default
get_frame_at_pts = torch.ops.torchcodec_ns.get_frame_at_pts.default
get_frame_at_index = torch.ops.torchcodec_ns.get_frame_at_index.default
_get_frames_at_indices_tensor_input = (
    torch.ops.torchcodec_ns.get_frames_at_indices.default
)
_get_frames_by_pts_tensor_input = torch.ops.torchcodec_ns.get_frames_by_pts.default
get_frames_in_range = torch.ops.torchcodec_ns.get_frames_in_range.default
get_frames_by_pts_in_range = torch.ops.torchcodec_ns.get_frames_by_pts_in_range.default
get_frames_by_pts_in_range_audio = (
    torch.ops.torchcodec_ns.get_frames_by_pts_in_range_audio.default
)
get_json_metadata = torch.ops.torchcodec_ns.get_json_metadata.default

# Building-block ops for torchcodec.decoders._blocks (CPU). See
# API_breakdown_claude_plan.md.
_blocks_create_demuxer = torch.ops.torchcodec_ns._blocks_create_demuxer.default
_blocks_demuxer_next_packet = (
    torch.ops.torchcodec_ns._blocks_demuxer_next_packet.default
)
_blocks_create_packet_decoder = (
    torch.ops.torchcodec_ns._blocks_create_packet_decoder.default
)
_blocks_packet_decoder_send_packet = (
    torch.ops.torchcodec_ns._blocks_packet_decoder_send_packet.default
)
_blocks_packet_decoder_send_eof = (
    torch.ops.torchcodec_ns._blocks_packet_decoder_send_eof.default
)
_blocks_packet_decoder_receive_frame = (
    torch.ops.torchcodec_ns._blocks_packet_decoder_receive_frame.default
)
_blocks_create_color_converter = (
    torch.ops.torchcodec_ns._blocks_create_color_converter.default
)
_blocks_convert_frame = torch.ops.torchcodec_ns._blocks_convert_frame.default

_test_frame_pts_equality = torch.ops.torchcodec_ns._test_frame_pts_equality.default
_get_container_json_metadata = (
    torch.ops.torchcodec_ns.get_container_json_metadata.default
)
_get_key_frame_indices = torch.ops.torchcodec_ns._get_key_frame_indices.default
scan_all_streams_to_update_metadata = (
    torch.ops.torchcodec_ns.scan_all_streams_to_update_metadata.default
)
_get_stream_json_metadata = torch.ops.torchcodec_ns.get_stream_json_metadata.default
_get_json_ffmpeg_library_versions = (
    torch.ops.torchcodec_ns._get_json_ffmpeg_library_versions.default
)
_get_backend_details = torch.ops.torchcodec_ns._get_backend_details.default
create_streaming_encoder = torch._dynamo.disallow_in_graph(
    torch.ops.torchcodec_ns.create_streaming_encoder.default
)
streaming_encoder_close = torch.ops.torchcodec_ns.streaming_encoder_close.default
streaming_encoder_add_video_stream = (
    torch.ops.torchcodec_ns.streaming_encoder_add_video_stream.default
)
streaming_encoder_add_audio_stream = (
    torch.ops.torchcodec_ns.streaming_encoder_add_audio_stream.default
)
streaming_encoder_open_file = (
    torch.ops.torchcodec_ns.streaming_encoder_open_file.default
)
_streaming_encoder_open_file_like = (
    torch.ops.torchcodec_ns.streaming_encoder_open_file_like.default
)
streaming_encoder_add_frames = (
    torch.ops.torchcodec_ns.streaming_encoder_add_frames.default
)
streaming_encoder_add_samples = (
    torch.ops.torchcodec_ns.streaming_encoder_add_samples.default
)
set_nvdec_cache_capacity = torch.ops.torchcodec_ns.set_nvdec_cache_capacity.default
get_nvdec_cache_capacity = torch.ops.torchcodec_ns.get_nvdec_cache_capacity.default
_get_nvdec_cache_size = torch.ops.torchcodec_ns._get_nvdec_cache_size.default
_set_cpp_log_level = torch.ops.torchcodec_ns._set_cpp_log_level.default
_get_log_level = torch.ops.torchcodec_ns._get_log_level.default
create_wav_decoder_from_file = (
    torch.ops.torchcodec_ns.create_wav_decoder_from_file.default
)
create_wav_decoder_from_tensor = (
    torch.ops.torchcodec_ns.create_wav_decoder_from_tensor.default
)
_create_wav_decoder_from_file_like = (
    torch.ops.torchcodec_ns._create_wav_decoder_from_file_like.default
)
get_wav_samples_in_range = torch.ops.torchcodec_ns.get_wav_samples_in_range.default
get_wav_metadata_from_decoder = (
    torch.ops.torchcodec_ns.get_wav_metadata_from_decoder.default
)


# =============================
# Functions not related to custom ops, but similar implementation to c++ ops
# =============================
def create_from_bytes(video_bytes: bytes, seek_mode: str | None = None) -> torch.Tensor:
    with warnings.catch_warnings():
        # Ignore warning stating that the underlying video_bytes buffer is
        # non-writable.
        warnings.filterwarnings("ignore", category=UserWarning)
        buffer = torch.frombuffer(video_bytes, dtype=torch.uint8)
    return create_from_tensor(buffer, seek_mode)


def create_from_file_like(
    file_like: io.RawIOBase | io.BufferedReader, seek_mode: str | None = None
) -> torch.Tensor:
    assert _pybind_ops is not None
    return _create_from_file_like(
        _pybind_ops.create_file_like_context(
            file_like, False  # False means not for writing
        ),
        seek_mode,
    )


def create_wav_decoder_from_bytes(wav_bytes: bytes) -> torch.Tensor:
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", category=UserWarning)
        buffer = torch.frombuffer(wav_bytes, dtype=torch.uint8)
    return create_wav_decoder_from_tensor(buffer)


def create_wav_decoder_from_file_like(
    file_like: io.RawIOBase | io.BufferedReader,
) -> torch.Tensor:
    assert _pybind_ops is not None
    return _create_wav_decoder_from_file_like(
        _pybind_ops.create_file_like_context(
            file_like, False  # False means not for writing
        ),
    )


def streaming_encoder_open_file_like(
    encoder: torch.Tensor,
    format: str,
    file_like: io.RawIOBase | io.BufferedIOBase,
) -> None:
    assert _pybind_ops is not None
    _streaming_encoder_open_file_like(
        encoder,
        format,
        _pybind_ops.create_file_like_context(file_like, True),
    )


def get_frames_at_indices(
    decoder: torch.Tensor, *, frame_indices: torch.Tensor | list[int]
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    if isinstance(frame_indices, torch.Tensor):
        # Ensure indices is the correct dtype (int64)
        frame_indices = frame_indices.to(torch.int64)
    else:
        # Convert list to tensor for dispatch
        frame_indices = torch.tensor(frame_indices)
    return _get_frames_at_indices_tensor_input(decoder, frame_indices=frame_indices)


def get_frames_by_pts(
    decoder: torch.Tensor, *, timestamps: torch.Tensor | list[float]
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    if isinstance(timestamps, torch.Tensor):
        # Ensure indices is the correct dtype (float64)
        timestamps = timestamps.to(torch.float64)
    else:
        # Convert list to tensor for dispatch
        try:
            timestamps = torch.tensor(timestamps, dtype=torch.float64)
        except Exception as e:
            raise ValueError("Couldn't convert timestamps input to a tensor") from e
    return _get_frames_by_pts_tensor_input(decoder, timestamps=timestamps)


# ==============================
# Abstract impl for the operators. Needed by torch.compile.
# ==============================
@register_fake("torchcodec_ns::create_from_file")
def create_from_file_abstract(filename: str, seek_mode: str | None) -> torch.Tensor:
    return torch.empty([], dtype=torch.long)


@register_fake("torchcodec_ns::_create_from_file_like")
def _create_from_file_like_abstract(
    file_like: int, seek_mode: str | None
) -> torch.Tensor:
    return torch.empty([], dtype=torch.long)


@register_fake("torchcodec_ns::create_from_tensor")
def create_from_tensor_abstract(
    video_tensor: torch.Tensor, seek_mode: str | None
) -> torch.Tensor:
    return torch.empty([], dtype=torch.long)


@register_fake("torchcodec_ns::_add_video_stream")
def _add_video_stream_abstract(
    decoder: torch.Tensor,
    *,
    num_threads: int | None = None,
    dimension_order: str | None = None,
    stream_index: int | None = None,
    device: str = "cpu",
    device_variant: str = "default",
    transform_specs: str = "",
    custom_frame_mappings_pts: torch.Tensor | None = None,
    custom_frame_mappings_duration: torch.Tensor | None = None,
    custom_frame_mappings_keyframe_indices: torch.Tensor | None = None,
    color_conversion_library: str | None = None,
    output_dtype: str = "uint8",
) -> None:
    return


@register_fake("torchcodec_ns::add_video_stream")
def add_video_stream_abstract(
    decoder: torch.Tensor,
    *,
    num_threads: int | None = None,
    dimension_order: str | None = None,
    stream_index: int | None = None,
    device: str = "cpu",
    device_variant: str = "default",
    transform_specs: str = "",
    custom_frame_mappings_pts: torch.Tensor | None = None,
    custom_frame_mappings_duration: torch.Tensor | None = None,
    custom_frame_mappings_keyframe_indices: torch.Tensor | None = None,
    output_dtype: str = "uint8",
) -> None:
    return


@register_fake("torchcodec_ns::add_audio_stream")
def add_audio_stream_abstract(
    decoder: torch.Tensor,
    *,
    stream_index: int | None = None,
    sample_rate: int | None = None,
    num_channels: int | None = None,
) -> None:
    return


@register_fake("torchcodec_ns::seek_to_pts")
def seek_abstract(decoder: torch.Tensor, seconds: float) -> None:
    return


@register_fake("torchcodec_ns::get_next_frame")
def get_next_frame_abstract(
    decoder: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    # Images are 3 dimensions: height, width, channels.
    # The exact permutation depends on the constructor options passed in.
    image_size = [get_ctx().new_dynamic_size() for _ in range(3)]
    return (
        torch.empty(image_size),
        torch.empty([], dtype=torch.float),
        torch.empty([], dtype=torch.float),
    )


@register_fake("torchcodec_ns::get_frame_at_pts")
def get_frame_at_pts_abstract(
    decoder: torch.Tensor, seconds: float
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    image_size = [get_ctx().new_dynamic_size() for _ in range(3)]
    return (
        torch.empty(image_size),
        torch.empty([], dtype=torch.float),
        torch.empty([], dtype=torch.float),
    )


@register_fake("torchcodec_ns::get_frames_by_pts")
def get_frames_by_pts_abstract(
    decoder: torch.Tensor,
    *,
    timestamps: torch.Tensor | list[float],
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    image_size = [get_ctx().new_dynamic_size() for _ in range(4)]
    return (
        torch.empty(image_size),
        torch.empty([], dtype=torch.float),
        torch.empty([], dtype=torch.float),
    )


@register_fake("torchcodec_ns::get_frame_at_index")
def get_frame_at_index_abstract(
    decoder: torch.Tensor, *, frame_index: int
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    image_size = [get_ctx().new_dynamic_size() for _ in range(3)]
    return (
        torch.empty(image_size),
        torch.empty([], dtype=torch.float),
        torch.empty([], dtype=torch.float),
    )


@register_fake("torchcodec_ns::get_frames_at_indices")
def get_frames_at_indices_abstract(
    decoder: torch.Tensor, *, frame_indices: torch.Tensor | list[int]
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    image_size = [get_ctx().new_dynamic_size() for _ in range(4)]
    return (
        torch.empty(image_size),
        torch.empty([], dtype=torch.float),
        torch.empty([], dtype=torch.float),
    )


@register_fake("torchcodec_ns::get_frames_in_range")
def get_frames_in_range_abstract(
    decoder: torch.Tensor,
    *,
    start: int,
    stop: int,
    step: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    image_size = [get_ctx().new_dynamic_size() for _ in range(4)]
    return (
        torch.empty(image_size),
        torch.empty([], dtype=torch.float),
        torch.empty([], dtype=torch.float),
    )


@register_fake("torchcodec_ns::get_frames_by_pts_in_range")
def get_frames_by_pts_in_range_abstract(
    decoder: torch.Tensor,
    *,
    start_seconds: float,
    stop_seconds: float,
    fps: float | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    image_size = [get_ctx().new_dynamic_size() for _ in range(4)]
    return (
        torch.empty(image_size),
        torch.empty([], dtype=torch.float),
        torch.empty([], dtype=torch.float),
    )


@register_fake("torchcodec_ns::get_frames_by_pts_in_range_audio")
def get_frames_by_pts_in_range_audio_abstract(
    decoder: torch.Tensor,
    *,
    start_seconds: float,
    stop_seconds: float | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    image_size = [get_ctx().new_dynamic_size() for _ in range(4)]
    return (torch.empty(image_size), torch.empty([], dtype=torch.float))


@register_fake("torchcodec_ns::_get_key_frame_indices")
def get_key_frame_indices_abstract(decoder: torch.Tensor) -> torch.Tensor:
    return torch.empty([], dtype=torch.int)


@register_fake("torchcodec_ns::get_json_metadata")
def get_json_metadata_abstract(decoder: torch.Tensor) -> str:
    return ""


@register_fake("torchcodec_ns::get_container_json_metadata")
def get_container_json_metadata_abstract(decoder: torch.Tensor) -> str:
    return ""


@register_fake("torchcodec_ns::get_stream_json_metadata")
def get_stream_json_metadata_abstract(decoder: torch.Tensor, stream_idx: int) -> str:
    return ""


@register_fake("torchcodec_ns::_test_frame_pts_equality")
def _test_frame_pts_equality_abstract(
    decoder: torch.Tensor,
    *,
    frame_index: int,
    pts_seconds_to_test: float,
) -> bool:
    return False


@register_fake("torchcodec_ns::_get_json_ffmpeg_library_versions")
def _get_json_ffmpeg_library_versions_abstract() -> str:
    return ""


@register_fake("torchcodec_ns::scan_all_streams_to_update_metadata")
def scan_all_streams_to_update_metadata_abstract(decoder: torch.Tensor) -> None:
    return


def get_ffmpeg_library_versions():
    versions_json = _get_json_ffmpeg_library_versions()
    return json.loads(versions_json)


@register_fake("torchcodec_ns::_get_backend_details")
def _get_backend_details_abstract(decoder: torch.Tensor) -> str:
    return ""


@register_fake("torchcodec_ns::create_streaming_encoder")
def _create_streaming_encoder_abstract() -> torch.Tensor:
    return torch.empty([], dtype=torch.long)


@register_fake("torchcodec_ns::streaming_encoder_close")
def streaming_encoder_close_abstract(encoder: torch.Tensor) -> None:
    return


@register_fake("torchcodec_ns::streaming_encoder_add_video_stream")
def streaming_encoder_add_video_stream_abstract(
    encoder: torch.Tensor,
    height: int,
    width: int,
    frame_rate: float,
    device: str = "cpu",
    codec: str | None = None,
    pixel_format: str | None = None,
    crf: float | None = None,
    preset: str | None = None,
    extra_options: list[str] | None = None,
) -> int:
    return 0


@register_fake("torchcodec_ns::streaming_encoder_add_audio_stream")
def streaming_encoder_add_audio_stream_abstract(
    encoder: torch.Tensor,
    sample_rate: int,
    num_channels: int,
    bit_rate: int | None = None,
    output_num_channels: int | None = None,
    output_sample_rate: int | None = None,
) -> int:
    return 0


@register_fake("torchcodec_ns::streaming_encoder_open_file")
def streaming_encoder_open_file_abstract(encoder: torch.Tensor, filename: str) -> None:
    return


@register_fake("torchcodec_ns::streaming_encoder_open_file_like")
def streaming_encoder_open_file_like_abstract(
    encoder: torch.Tensor, format: str, file_like_context: int
) -> None:
    return


@register_fake("torchcodec_ns::streaming_encoder_add_frames")
def streaming_encoder_add_frames_abstract(
    encoder: torch.Tensor, frames: torch.Tensor, stream_index: int
) -> None:
    return


@register_fake("torchcodec_ns::streaming_encoder_add_samples")
def streaming_encoder_add_samples_abstract(
    encoder: torch.Tensor, samples: torch.Tensor, stream_index: int
) -> None:
    return


@register_fake("torchcodec_ns::set_nvdec_cache_capacity")
def set_nvdec_cache_capacity_abstract(capacity: int) -> None:
    return


@register_fake("torchcodec_ns::get_nvdec_cache_capacity")
def get_nvdec_cache_capacity_abstract() -> int:
    return 0


@register_fake("torchcodec_ns::_get_nvdec_cache_size")
def _get_nvdec_cache_size_abstract(device_index: int) -> int:
    return 0


@register_fake("torchcodec_ns::_set_cpp_log_level")
def _set_cpp_log_level_abstract(level: int) -> None:
    return


@register_fake("torchcodec_ns::_get_log_level")
def _get_log_level_abstract() -> int:
    return 0


@register_fake("torchcodec_ns::create_wav_decoder_from_file")
def create_wav_decoder_from_file_abstract(filename: str) -> torch.Tensor:
    return torch.empty([], dtype=torch.long)


@register_fake("torchcodec_ns::create_wav_decoder_from_tensor")
def create_wav_decoder_from_tensor_abstract(data: torch.Tensor) -> torch.Tensor:
    return torch.empty([], dtype=torch.long)


@register_fake("torchcodec_ns::_create_wav_decoder_from_file_like")
def _create_wav_decoder_from_file_like_abstract(
    file_like_context: int,
) -> torch.Tensor:
    return torch.empty([], dtype=torch.long)


@register_fake("torchcodec_ns::get_wav_samples_in_range")
def get_wav_samples_in_range_abstract(
    decoder: torch.Tensor, start_seconds: float, stop_seconds: float | None
) -> tuple[torch.Tensor, torch.Tensor]:
    sample_size = [
        get_ctx().new_dynamic_size() for _ in range(2)
    ]  # [channels, samples]
    frames = torch.empty(sample_size)
    pts = torch.empty([], dtype=torch.float64)
    return frames, pts


@register_fake("torchcodec_ns::get_wav_metadata_from_decoder")
def get_wav_metadata_from_decoder_abstract(decoder: torch.Tensor) -> str:
    return ""
