# Copyright 2024 LiveKit, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import asyncio
import enum
import io
import struct
import threading
from collections.abc import AsyncIterator
from concurrent.futures import ThreadPoolExecutor
from typing import cast

import av
import av.container

from livekit import rtc

from ...log import logger
from .. import aio
from ..audio import AudioByteStream


def _mime_to_av_format(mime: str | None) -> str | None:
    """Return the libav *container* short‑name for a given MIME‑type.

    If *mime* is *None* or not recognised, return *None* so that PyAV will
    fall back to auto‑detection.
    """

    if not mime:
        return None

    mime = mime.lower()
    _TABLE: dict[str, str] = {
        "audio/mpeg": "mp3",
        "audio/mp3": "mp3",
        "audio/x-mpeg": "mp3",
        "audio/aac": "aac",
        "audio/x-aac": "aac",
        "audio/flac": "flac",
        "audio/x-flac": "flac",
        "audio/wav": "wav",
        "audio/wave": "wav",
        "audio/x-wav": "wav",
        "audio/opus": "ogg",
        "audio/ogg": "ogg",
        "audio/webm": "webm",
        "audio/mp4": "mp4",
    }
    return _TABLE.get(mime)


class StreamBuffer:
    """
    A thread-safe buffer that behaves like an IO stream.
    Allows writing from one thread and reading from another.
    """

    _COMPACT_THRESHOLD = 5 * 1024 * 1024  # compact after 5MB consumed

    def __init__(self) -> None:
        self._bio = io.BytesIO()
        self._lock = threading.Lock()
        self._data_available = threading.Condition(self._lock)
        self._eof = False
        self._closed = False
        self._write_pos = 0
        self._read_pos = 0

    def write(self, data: bytes) -> None:
        """Write data to the buffer from a writer thread."""
        with self._data_available:
            self._bio.seek(self._write_pos)
            self._bio.write(data)
            self._write_pos = self._bio.tell()
            self._data_available.notify_all()

    def read(self, size: int = -1) -> bytes:
        """Read data from the buffer in a reader thread."""
        if size == 0:
            return b""

        with self._data_available:
            while True:
                if self._closed:
                    return b""

                available = self._write_pos - self._read_pos
                if available > 0:
                    self._bio.seek(self._read_pos)
                    if size < 0:
                        data = self._bio.read(available)
                    else:
                        data = self._bio.read(min(size, available))
                    self._read_pos = self._bio.tell()

                    if self._read_pos >= self._COMPACT_THRESHOLD:
                        remaining = self._bio.read()
                        self._bio = io.BytesIO(remaining)
                        self._bio.seek(0, io.SEEK_END)
                        self._write_pos = self._bio.tell()
                        self._read_pos = 0

                    return data if data else b""

                if self._eof:
                    return b""

                self._data_available.wait()

    def end_input(self) -> None:
        """Signal that no more data will be written."""
        with self._data_available:
            self._eof = True
            self._data_available.notify_all()

    def close(self) -> None:
        with self._data_available:
            self._closed = True
            self._data_available.notify_all()
            self._bio.close()


class _WavState(enum.IntEnum):
    RIFF_HEADER = 0
    CHUNK_HEADER = 1
    FMT_DATA = 2
    SKIP_CHUNK_DATA = 3
    STREAMING = 4


class _WavInlineDecoder:
    """Incremental WAV decoder that runs entirely on the event loop (no thread).

    Processes WAV bytes via a state machine:
    RIFF_HEADER → CHUNK_HEADER → FMT_DATA/SKIP_CHUNK_DATA → STREAMING.
    Once in STREAMING state, subsequent push() calls feed bytes directly to
    AudioByteStream → optional resampler → output channel.

    Each push() may contain a complete WAV file (with its own headers). When a
    new RIFF header is detected while already streaming, the current stream is
    flushed and the state machine resets to parse the new file's headers.
    """

    _RIFF_MAGIC = b"RIFF"

    def __init__(
        self,
        output_ch: aio.Chan[rtc.AudioFrame],
        sample_rate: int | None,
    ) -> None:
        self._output_ch = output_ch
        self._sample_rate = sample_rate

        self._state = _WavState.RIFF_HEADER
        self._hdr_buf = bytearray()
        self._need = 12  # bytes needed for current header state
        self._skip_remaining = 0
        self._chunk_size = 0

        # set after fmt is parsed
        self._bstream: AudioByteStream | None = None
        self._resampler: rtc.AudioResampler | None = None
        self._wave_channels = 0
        self._wave_rate = 0

    def push(self, data: bytes) -> None:
        if self._state == _WavState.STREAMING:
            if len(data) >= 4 and data[:4] == self._RIFF_MAGIC:
                self._flush_current()
                self._reset_state()
            else:
                self._push_pcm(data)
                return

        buf = memoryview(data)
        pos = 0
        while pos < len(buf):
            if self._state == _WavState.RIFF_HEADER:
                pos = self._consume_riff(buf, pos)
            elif self._state == _WavState.CHUNK_HEADER:
                pos = self._consume_chunk_header(buf, pos)
            elif self._state == _WavState.FMT_DATA:
                pos = self._consume_fmt_data(buf, pos)
            elif self._state == _WavState.SKIP_CHUNK_DATA:
                pos = self._consume_skip(buf, pos)
            elif self._state == _WavState.STREAMING:
                # remainder after headers goes straight to PCM path
                self._push_pcm(bytes(buf[pos:]))
                return

    def flush(self) -> None:
        self._flush_current()

    def _flush_current(self) -> None:
        """Flush AudioByteStream and resampler for the current WAV segment."""
        if self._bstream is not None:
            remaining = self._bstream.flush()
            if self._resampler is not None:
                for frame in remaining:
                    for resampled in self._resampler.push(frame):
                        self._output_ch.send_nowait(resampled)
                for frame in self._resampler.flush():
                    if frame.samples_per_channel > 0:
                        self._output_ch.send_nowait(frame)
            else:
                for frame in remaining:
                    self._output_ch.send_nowait(frame)

    def _reset_state(self) -> None:
        """Reset the state machine to parse a new WAV file."""
        self._state = _WavState.RIFF_HEADER
        self._hdr_buf.clear()
        self._need = 12
        self._skip_remaining = 0
        self._chunk_size = 0
        self._bstream = None
        self._resampler = None
        self._wave_channels = 0
        self._wave_rate = 0

    # -- state handlers -------------------------------------------------------

    def _consume_riff(self, buf: memoryview, pos: int) -> int:
        take = min(self._need - len(self._hdr_buf), len(buf) - pos)
        self._hdr_buf.extend(buf[pos : pos + take])
        pos += take
        if len(self._hdr_buf) < self._need:
            return pos

        if self._hdr_buf[:4] != b"RIFF" or self._hdr_buf[8:12] != b"WAVE":
            raise ValueError(f"Invalid WAV file: missing RIFF/WAVE: {bytes(self._hdr_buf)!r}")
        self._hdr_buf.clear()
        self._need = 8
        self._state = _WavState.CHUNK_HEADER
        return pos

    def _consume_chunk_header(self, buf: memoryview, pos: int) -> int:
        take = min(self._need - len(self._hdr_buf), len(buf) - pos)
        self._hdr_buf.extend(buf[pos : pos + take])
        pos += take
        if len(self._hdr_buf) < self._need:
            return pos

        chunk_id, chunk_size = struct.unpack("<4sI", bytes(self._hdr_buf[:8]))
        self._hdr_buf.clear()
        self._chunk_size = chunk_size

        if chunk_id == b"fmt ":
            self._need = chunk_size
            self._state = _WavState.FMT_DATA
        elif chunk_id == b"data":
            self._init_streaming()
            self._state = _WavState.STREAMING
        else:
            self._skip_remaining = chunk_size
            self._state = _WavState.SKIP_CHUNK_DATA
        return pos

    def _consume_fmt_data(self, buf: memoryview, pos: int) -> int:
        take = min(self._need - len(self._hdr_buf), len(buf) - pos)
        self._hdr_buf.extend(buf[pos : pos + take])
        pos += take
        if len(self._hdr_buf) < self._need:
            return pos

        fmt = bytes(self._hdr_buf[: self._chunk_size])
        audio_format, channels, rate = struct.unpack("<HHI", fmt[:8])
        if len(fmt) >= 16:
            bits_per_sample = struct.unpack("<H", fmt[14:16])[0]
            if bits_per_sample != 16:
                raise ValueError(
                    f"Unsupported WAV bits per sample: {bits_per_sample}"
                    " (only 16-bit PCM supported)"
                )
        if audio_format != 1:
            raise ValueError(f"Unsupported WAV audio format: {audio_format}")

        self._wave_channels = channels
        self._wave_rate = rate
        self._hdr_buf.clear()
        self._need = 8
        self._state = _WavState.CHUNK_HEADER
        return pos

    def _consume_skip(self, buf: memoryview, pos: int) -> int:
        take = min(self._skip_remaining, len(buf) - pos)
        self._skip_remaining -= take
        pos += take
        if self._skip_remaining == 0:
            self._hdr_buf.clear()
            self._need = 8
            self._state = _WavState.CHUNK_HEADER
        return pos

    # -- streaming helpers ----------------------------------------------------

    def _init_streaming(self) -> None:
        if self._wave_rate == 0:
            raise ValueError("Invalid WAV file: data chunk before fmt chunk")

        self._bstream = AudioByteStream(
            sample_rate=self._wave_rate, num_channels=self._wave_channels
        )
        if self._sample_rate is not None and self._sample_rate != self._wave_rate:
            self._resampler = rtc.AudioResampler(
                input_rate=self._wave_rate,
                output_rate=self._sample_rate,
                num_channels=self._wave_channels,
            )

    def _push_pcm(self, data: bytes) -> None:
        assert self._bstream is not None
        if self._resampler is not None:
            for frame in self._bstream.push(data):
                for resampled in self._resampler.push(frame):
                    self._output_ch.send_nowait(resampled)
        else:
            for frame in self._bstream.push(data):
                self._output_ch.send_nowait(frame)


class AudioStreamDecoder:
    """A class that can be used to decode audio stream into PCM AudioFrames.

    Decoders are stateful, and it should not be reused across multiple streams. Each decoder
    is designed to decode a single stream.
    """

    def __init__(
        self,
        *,
        sample_rate: int | None = 48000,
        num_channels: int | None = 1,
        format: str | None = None,
    ):
        self._sample_rate = sample_rate

        self._layout = "mono"
        if num_channels == 2:
            self._layout = "stereo"

        self._mime_type = format.lower() if format else None
        self._av_format = _mime_to_av_format(self._mime_type)
        self._is_wav = self._av_format == "wav"

        self._output_ch = aio.Chan[rtc.AudioFrame]()
        self._closed = False
        self._started = False
        self._loop = asyncio.get_event_loop()

        # lazy-initialized only for non-WAV codecs
        self._input_buf: StreamBuffer | None = None
        self._executor: ThreadPoolExecutor | None = None

        # lazy-initialized only for WAV
        self._wav_decoder: _WavInlineDecoder | None = None

    def push(self, chunk: bytes) -> None:
        if self._is_wav:
            if self._wav_decoder is None:
                self._wav_decoder = _WavInlineDecoder(self._output_ch, self._sample_rate)
            try:
                self._wav_decoder.push(chunk)
            except Exception:
                if not self._closed:
                    logger.exception("error decoding WAV audio")
                    self._output_ch.close()
                    self._closed = True
                return
            self._started = True
            return

        if self._input_buf is None:
            self._input_buf = StreamBuffer()
            self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="AudioDecoder")
        self._input_buf.write(chunk)
        if not self._started:
            self._started = True
            self._loop.run_in_executor(self._executor, self._decode_loop)

    def end_input(self) -> None:
        if self._is_wav:
            if self._wav_decoder is not None and not self._closed:
                try:
                    self._wav_decoder.flush()
                except Exception:
                    logger.exception("error flushing WAV audio")
            if not self._closed:
                self._output_ch.close()
            return

        if self._input_buf is not None:
            self._input_buf.end_input()
        if not self._started:
            self._output_ch.close()

    def _decode_loop(self) -> None:
        container: av.container.InputContainer | None = None
        resampler: av.AudioResampler | None = None
        try:
            # open container in low-latency streaming mode
            container = av.open(
                self._input_buf,
                mode="r",
                format=self._av_format,
                buffer_size=256,
                options={
                    "probesize": "32",
                    "analyzeduration": "0",
                    "fflags": "nobuffer+flush_packets",
                    "flags": "low_delay",
                    "reorder_queue_size": "0",
                    "max_delay": "0",
                    "avioflags": "direct",
                },
            )
            # explicitly disable internal buffering flags on the FFmpeg container
            container.flags |= cast(
                int, av.container.Flags.no_buffer.value | av.container.Flags.flush_packets.value
            )

            if len(container.streams.audio) == 0:
                raise ValueError("no audio stream found")

            audio_stream = container.streams.audio[0]

            # Set up resampler only if needed
            if self._sample_rate is not None or self._layout is not None:
                resampler = av.AudioResampler(
                    format="s16", layout=self._layout, rate=self._sample_rate
                )

            for frame in container.decode(audio_stream):
                if self._closed:
                    return

                if resampler:
                    frames = resampler.resample(frame)
                else:
                    frames = [frame]

                for f in frames:
                    self._emit_av_frame(f)

            # flush the resampler to get any remaining buffered samples
            if resampler and not self._closed:
                for f in resampler.resample(None):
                    self._emit_av_frame(f)

        except Exception:
            logger.exception("error decoding audio")
        finally:
            self._loop.call_soon_threadsafe(self._output_ch.close)
            if container:
                container.close()

    def _emit_av_frame(self, f: av.AudioFrame) -> None:
        self._loop.call_soon_threadsafe(
            self._output_ch.send_nowait,
            rtc.AudioFrame(
                data=f.to_ndarray().tobytes(),
                num_channels=len(f.layout.channels),
                sample_rate=int(f.sample_rate),
                samples_per_channel=f.samples,
            ),
        )

    def __aiter__(self) -> AsyncIterator[rtc.AudioFrame]:
        return self

    async def __anext__(self) -> rtc.AudioFrame:
        return await self._output_ch.__anext__()

    async def aclose(self) -> None:
        if self._closed:
            return

        self.end_input()
        self._closed = True

        if self._input_buf is not None:
            self._input_buf.close()

        if not self._started:
            return

        try:
            async for _ in self._output_ch:
                pass
        finally:
            if self._executor is not None:
                self._executor.shutdown(wait=False, cancel_futures=True)