# 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
from pathlib import Path

import torch
from torch import Tensor

from torchcodec import _core as core, AudioSamples
from torchcodec._core._decoder_utils import create_audio_decoder


class AudioDecoder:
    """A single-stream audio decoder.

    This can be used to decode audio from pure audio files (e.g. mp3, wav,
    etc.), or from videos that contain audio streams (e.g. mp4 videos).

    .. tip::

       If you are decoding WAV files, consider using
       :class:`~torchcodec.decoders.WavDecoder` instead for significantly faster
       decoding.

    Returned samples are float samples normalized in [-1, 1]

    Args:
        source (str, ``Pathlib.path``, bytes, ``torch.Tensor`` or file-like
            object): The source of the video or audio:

            - If ``str``: a local path or a URL to a video or audio file.
            - If ``Pathlib.path``: a path to a local video or audio file.
            - If ``bytes`` object or ``torch.Tensor``: the raw encoded audio data.
            - If file-like object: we read video data from the object on demand. The object must
              expose the methods `read(self, size: int) -> bytes` and
              `seek(self, offset: int, whence: int) -> int`. Read more in:
              :ref:`sphx_glr_generated_examples_decoding_file_like.py`.
        stream_index (int, optional): Specifies which stream in the file to decode samples from.
            Note that this index is absolute across all media types. If left unspecified, then
            the :term:`best stream` is used.
        sample_rate (int, optional): The desired output sample rate of the decoded samples.
            By default, the sample rate of the source is used.
        num_channels (int, optional): The desired number of channels of the decoded samples.
            By default, the number of channels of the source is used.

    Attributes:
        metadata (AudioStreamMetadata): Metadata of the audio stream.
        stream_index (int): The stream index that this decoder is retrieving samples from. If a
            stream index was provided at initialization, this is the same value. If it was left
            unspecified, this is the :term:`best stream`.
    """

    def __init__(
        self,
        source: str | Path | io.RawIOBase | io.BufferedReader | bytes | Tensor,
        *,
        stream_index: int | None = None,
        sample_rate: int | None = None,
        num_channels: int | None = None,
    ):
        torch._C._log_api_usage_once("torchcodec.decoders.AudioDecoder")

        (
            self._decoder,
            self.stream_index,
            self.metadata,
        ) = create_audio_decoder(
            source=source,
            seek_mode="approximate",
            stream_index=stream_index,
            sample_rate=sample_rate,
            num_channels=num_channels,
        )

        self._desired_sample_rate = (
            sample_rate if sample_rate is not None else self.metadata.sample_rate
        )

    def get_all_samples(self) -> AudioSamples:
        """Returns all the audio samples from the source.

        To decode samples in a specific range, use
        :meth:`~torchcodec.decoders.AudioDecoder.get_samples_played_in_range`.

        Returns:
            AudioSamples: The samples within the file.
        """
        return self.get_samples_played_in_range()

    def get_samples_played_in_range(
        self, start_seconds: float = 0.0, stop_seconds: float | None = None
    ) -> AudioSamples:
        """Returns audio samples in the given range.

        Samples are in the half open range [start_seconds, stop_seconds).

        To decode all the samples from beginning to end, you can call this
        method while leaving ``start_seconds`` and ``stop_seconds`` to their
        default values, or use
        :meth:`~torchcodec.decoders.AudioDecoder.get_all_samples` as a more
        convenient alias.

        Args:
            start_seconds (float): Time, in seconds, of the start of the
                range. Default: 0.
            stop_seconds (float or None): Time, in seconds, of the end of the
                range. As a half open range, the end is excluded. Default: None,
                which decodes samples until the end.

        Returns:
            AudioSamples: The samples within the specified range.
        """
        if stop_seconds is not None and not start_seconds <= stop_seconds:
            raise ValueError(
                f"Invalid start seconds: {start_seconds}. It must be less than or equal to stop seconds ({stop_seconds})."
            )
        frames, first_pts = core.get_frames_by_pts_in_range_audio(
            self._decoder,
            start_seconds=start_seconds,
            stop_seconds=stop_seconds,
        )
        first_pts = first_pts.item()

        # x = frame boundaries
        #
        #            first_pts                                    last_pts
        #                v                                            v
        # ....x..........x..........x...........x..........x..........x.....
        #                    ^                                 ^
        #               start_seconds                      stop_seconds
        #
        # We want to return the samples in [start_seconds, stop_seconds). But
        # because the core API is based on frames, the `frames` tensor contains
        # the samples in [first_pts, last_pts)
        # So we do some basic math to figure out the position of the view that
        # we'll return.

        sample_rate = self._desired_sample_rate
        # TODO: metadata's sample_rate should probably not be Optional
        assert sample_rate is not None  # mypy.

        if first_pts < start_seconds:
            offset_beginning = round((start_seconds - first_pts) * sample_rate)
            output_pts_seconds = start_seconds
        else:
            # In normal cases we'll have first_pts <= start_pts, but in some
            # edge cases it's possible to have first_pts > start_seconds,
            # typically if the stream's first frame's pts isn't exactly 0.
            offset_beginning = 0
            output_pts_seconds = first_pts

        num_samples = frames.shape[1]
        last_pts = first_pts + num_samples / sample_rate
        if stop_seconds is not None and stop_seconds < last_pts:
            offset_end = num_samples - round((last_pts - stop_seconds) * sample_rate)
        else:
            offset_end = num_samples

        data = frames[:, offset_beginning:offset_end]
        return AudioSamples(
            data=data,
            pts_seconds=output_pts_seconds,
            duration_seconds=data.shape[1] / sample_rate,
            sample_rate=sample_rate,
        )
