# 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.

from __future__ import annotations

from torchcodec._core.ops import _blocks_convert_frame, _blocks_create_color_converter
from torchcodec._frame import Frame

from ._frame import DecodedFrame

# TODO_API_BREAKDOWN FEAT support output_dtype?
# TODO_API_BREAKDOWN FEAT We need to support rotation metadata!!
# TODO_API_BREAKDOWN FEAT Implement seeking?
# TODO_API_BREAKDOWN FEAT Implement range-getting (start/end time) for decoding?


class ColorConverter:
    """Color-conversion building block: turns a decoded (YUV)
    :class:`DecodedFrame` into an RGB :class:`~torchcodec._frame.Frame`
    (CHW, uint8 -- matching ``VideoDecoder``'s default output).

    Not bound to anything: everything it needs (dims, pixel format, colorspace)
    comes from the frame itself, so one converter can process frames from any
    video. Passive and *not* thread-safe: use one ``ColorConverter`` per thread.

    Note: automatic rotation (from stream side data) is not applied, since this
    block is intentionally stream-agnostic.
    """

    # TODO_API_BREAKDOWN P1: device default should be None
    # TODO_API_BREAKDOWN P1: add checks for coupling between device param of
    # PacketDecoder and ColorConverter. What if one is CPU and the other is
    # CUDA? What if they're different CUDA devices? Maybe we should just error.
    def __init__(self, device="cpu"):
        self._handle = _blocks_create_color_converter(device=device)

    def convert(self, decoded_frame: DecodedFrame) -> Frame:
        data = _blocks_convert_frame(self._handle, decoded_frame._handle)
        # The core op produces HWC; permute to CHW to match VideoDecoder (which
        # also returns a non-contiguous permuted view).
        data = data.permute(2, 0, 1)
        return Frame(
            data=data,
            pts_seconds=decoded_frame.pts_seconds,
            duration_seconds=decoded_frame.duration_seconds,
        )
