from __future__ import annotations

import asyncio
import contextlib
import queue
import threading
import time
from collections.abc import AsyncIterator
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable

import av
import numpy as np

from livekit import rtc

from ...log import logger
from .. import io

if TYPE_CHECKING:
    from ..agent_session import AgentSession

# the recorder currently assume the input is a continous uninterrupted audio stream


WRITE_INTERVAL = 2.5


class RecorderIO:
    def __init__(
        self,
        *,
        agent_session: AgentSession,
        sample_rate: int = 48000,
        loop: asyncio.AbstractEventLoop | None = None,
    ) -> None:
        self._in_record: RecorderAudioInput | None = None
        self._out_record: RecorderAudioOutput | None = None

        self._in_q: queue.Queue[list[rtc.AudioFrame] | None] = queue.Queue()
        self._out_q: queue.Queue[list[rtc.AudioFrame] | None] = queue.Queue()
        self._session = agent_session
        self._sample_rate = sample_rate
        self._started = False
        self._loop = loop or asyncio.get_event_loop()
        self._lock = asyncio.Lock()
        self._close_fut: asyncio.Future[None] = self._loop.create_future()
        self._output_path: Path | None = None

    async def start(self, *, output_path: str | Path) -> None:
        async with self._lock:
            if self._started:
                return

            if not self._in_record or not self._out_record:
                raise RuntimeError(
                    "RecorderIO not properly initialized: both `record_input()` and `record_output()` "
                    "must be called before starting the recorder."
                )

            self._output_path = Path(output_path)
            self._started = True
            self._close_fut = self._loop.create_future()
            self._forward_atask = asyncio.create_task(self._forward_task())

            thread = threading.Thread(
                target=self._encode_thread, daemon=True, name="recorder_io_encode_thread"
            )
            thread.start()

    async def aclose(self) -> None:
        async with self._lock:
            if not self._started:
                return

            self._in_q.put_nowait(None)
            self._out_q.put_nowait(None)
            await asyncio.shield(self._close_fut)
            self._started = False

    def record_input(self, audio_input: io.AudioInput) -> RecorderAudioInput:
        self._in_record = RecorderAudioInput(recording_io=self, source=audio_input)
        return self._in_record

    def record_output(self, audio_output: io.AudioOutput) -> RecorderAudioOutput:
        self._out_record = RecorderAudioOutput(
            recording_io=self, audio_output=audio_output, write_fnc=self._write_cb
        )
        return self._out_record

    @property
    def recording(self) -> bool:
        return self._started

    @property
    def output_path(self) -> Path | None:
        return self._output_path

    @property
    def recording_started_at(self) -> float | None:
        in_t = self._in_record.started_wall_time if self._in_record else None
        out_t = self._out_record.started_wall_time if self._out_record else None

        if in_t is None:
            return out_t

        if out_t is None:
            return in_t

        return min(in_t, out_t)

    def _write_cb(self, buf: list[rtc.AudioFrame]) -> None:
        assert self._in_record is not None

        input_buf = self._in_record.take_buf()
        self._in_q.put_nowait(input_buf)
        self._out_q.put_nowait(buf)

    async def _forward_task(self) -> None:
        assert self._in_record is not None
        assert self._out_record is not None

        # Forward the input audio to the encoder every 5s.
        while True:
            await asyncio.sleep(WRITE_INTERVAL)
            if self._out_record.has_pending_data:
                # if the output is currenetly playing audio, wait for it to stay in sync
                continue  # always wait for the complete output

            input_buf = self._in_record.take_buf()
            self._in_q.put_nowait(input_buf)
            self._out_q.put_nowait([])

    def _encode_thread(self) -> None:
        GROW_FACTOR = 1.5
        INV_INT16 = 1.0 / 32768.0

        assert self._output_path is not None
        self._output_path.parent.mkdir(parents=True, exist_ok=True)

        container = av.open(self._output_path, mode="w", format="ogg")
        stream: av.AudioStream = container.add_stream(
            "opus", rate=self._sample_rate, layout="stereo"
        )  # type: ignore

        in_resampler: rtc.AudioResampler | None = None
        out_resampler: rtc.AudioResampler | None = None

        capacity = self._sample_rate * 6  # 6s, 1ch
        stereo_buf = np.zeros((2, capacity), dtype=np.float32)

        def remix_and_resample(frames: list[rtc.AudioFrame], channel_idx: int) -> int:
            total_samples = sum(f.samples_per_channel * f.num_channels for f in frames)

            nonlocal capacity, stereo_buf
            if total_samples > capacity:
                while capacity < total_samples:
                    capacity = int(capacity * GROW_FACTOR)

                stereo_buf.resize((2, capacity), refcheck=False)

            pos = 0
            dest = stereo_buf[channel_idx]
            for f in frames:
                count = f.samples_per_channel * f.num_channels
                arr_i16 = np.frombuffer(f.data, dtype=np.int16, count=count).reshape(
                    -1, f.num_channels
                )
                slice_ = dest[pos : pos + f.samples_per_channel]
                np.sum(arr_i16, axis=1, dtype=np.float32, out=slice_)
                slice_ *= INV_INT16 / f.num_channels
                pos += f.samples_per_channel

            return pos

        with container:
            while True:
                input_buf = self._in_q.get()
                output_buf = self._out_q.get()

                if input_buf is None or output_buf is None:
                    break

                # lazy creation of the resamplers
                if in_resampler is None and len(input_buf):
                    input_rate, num_channels = input_buf[0].sample_rate, input_buf[0].num_channels
                    in_resampler = rtc.AudioResampler(
                        input_rate=input_rate,
                        output_rate=self._sample_rate,
                        num_channels=num_channels,
                    )

                if out_resampler is None and len(output_buf):
                    input_rate, num_channels = output_buf[0].sample_rate, output_buf[0].num_channels
                    out_resampler = rtc.AudioResampler(
                        input_rate=input_rate,
                        output_rate=self._sample_rate,
                        num_channels=num_channels,
                    )

                input_resampled = []
                for frame in input_buf:
                    assert in_resampler is not None
                    input_resampled.extend(in_resampler.push(frame))

                output_resampled = []
                for frame in output_buf:
                    assert out_resampler is not None
                    output_resampled.extend(out_resampler.push(frame))

                if output_buf:
                    assert out_resampler is not None
                    # the output is sent per-segment. Always flush when the playback is done
                    output_resampled.extend(out_resampler.flush())

                len_left = remix_and_resample(input_resampled, 0)
                len_right = remix_and_resample(output_resampled, 1)

                if len_left != len_right:
                    diff = abs(len_right - len_left)
                    if len_left < len_right:
                        logger.warning(
                            f"Input is shorter by {diff} samples; silence has been prepended to "
                            "align the input channel. The resulting recording may not accurately "
                            "reflect the original audio."
                        )
                        stereo_buf[0, diff : diff + len_left] = stereo_buf[0, :len_left]
                        stereo_buf[0, :diff] = 0.0
                        len_left = len_right
                    else:
                        stereo_buf[1, diff : diff + len_right] = stereo_buf[1, :len_right]
                        stereo_buf[1, :diff] = 0.0
                        len_right = len_left

                max_len = max(len_left, len_right)
                if max_len <= 0:
                    continue

                stereo_slice = stereo_buf[:, :max_len]
                av_frame = av.AudioFrame.from_ndarray(stereo_slice, format="fltp", layout="stereo")
                av_frame.sample_rate = self._sample_rate

                for packet in stream.encode(av_frame):
                    container.mux(packet)

            for packet in stream.encode(None):
                container.mux(packet)

        with contextlib.suppress(RuntimeError):
            self._loop.call_soon_threadsafe(self._close_fut.set_result, None)


class RecorderAudioInput(io.AudioInput):
    def __init__(self, *, recording_io: RecorderIO, source: io.AudioInput) -> None:
        super().__init__(label="RecorderIO", source=source)
        self.__audio_input = source
        self.__recording_io = recording_io
        self.__acc_frames: list[rtc.AudioFrame] = []
        self.__started_time: None | float = None

    @property
    def started_wall_time(self) -> float | None:
        return self.__started_time

    def take_buf(self) -> list[rtc.AudioFrame]:
        frames = self.__acc_frames
        self.__acc_frames = []
        return frames

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

    async def __anext__(self) -> rtc.AudioFrame:
        frame = await self.__audio_input.__anext__()

        if self.__recording_io.recording:
            if self.__started_time is None:
                self.__started_time = time.time()

            self.__acc_frames.append(frame)

        return frame


class RecorderAudioOutput(io.AudioOutput):
    def __init__(
        self,
        *,
        recording_io: RecorderIO,
        audio_output: io.AudioOutput | None = None,
        write_fnc: Callable[[list[rtc.AudioFrame]], Any],
    ) -> None:
        super().__init__(
            label="RecorderIO",
            next_in_chain=audio_output,
            sample_rate=None,
            capabilities=io.AudioOutputCapabilities(pause=True),  # depends on the next_in_chain
        )
        self.__recording_io = recording_io
        self.__write = write_fnc
        self.__acc_frames: list[rtc.AudioFrame] = []
        self.__started_time: None | float = None

    @property
    def started_wall_time(self) -> float | None:
        return self.__started_time

    @property
    def recorder_io(self) -> RecorderIO:
        return self.__recording_io

    @property
    def has_pending_data(self) -> bool:
        return len(self.__acc_frames) > 0

    def on_playback_finished(
        self,
        *,
        playback_position: float,
        interrupted: bool,
        synchronized_transcript: str | None = None,
    ) -> None:
        super().on_playback_finished(
            playback_position=playback_position,
            interrupted=interrupted,
            synchronized_transcript=synchronized_transcript,
        )

        if not self.__recording_io.recording:
            return

        buf = []
        acc_dur = 0.0
        for frame in self.__acc_frames:
            if frame.duration + acc_dur > playback_position:
                duration_needed = playback_position - acc_dur
                samples_needed = int(duration_needed * frame.sample_rate)  #  * frame.num_channels
                truncated_frame = rtc.AudioFrame(
                    data=frame.data[: samples_needed * frame.num_channels],
                    num_channels=frame.num_channels,
                    samples_per_channel=samples_needed,
                    sample_rate=frame.sample_rate,
                )
                buf.append(truncated_frame)
                break

            acc_dur += frame.duration
            buf.append(frame)

        if buf:
            self.__write(buf)

        self.__acc_frames = []

    async def capture_frame(self, frame: rtc.AudioFrame) -> None:
        await super().capture_frame(frame)

        if self.__recording_io.recording:
            if self.__started_time is None:
                self.__started_time = time.time()

            self.__acc_frames.append(frame)

        if self.next_in_chain:
            await self.next_in_chain.capture_frame(frame)

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

        if self.next_in_chain:
            self.next_in_chain.flush()

    def clear_buffer(self) -> None:
        if self.next_in_chain:
            self.next_in_chain.clear_buffer()
