"""Pipecat local-audio transport backed by sounddevice instead of PyAudio.

Pipecat ships `pipecat.transports.local.audio.LocalAudioTransport`, but it needs
PyAudio, which has no macOS wheel and must be compiled against a Homebrew
portaudio. This machine blocks both, so we talk to the same portaudio through
sounddevice, whose wheel bundles a prebuilt dylib. The frame contract is
identical to the upstream transport.
"""

import asyncio
from concurrent.futures import ThreadPoolExecutor

import sounddevice as sd
from loguru import logger

from pipecat.frames.frames import InputAudioRawFrame, OutputAudioRawFrame, StartFrame
from pipecat.processors.frame_processor import FrameProcessor
from pipecat.transports.base_input import BaseInputTransport
from pipecat.transports.base_output import BaseOutputTransport
from pipecat.transports.base_transport import BaseTransport, TransportParams


class SoundDeviceTransportParams(TransportParams):
    """Configuration for the sounddevice transport.

    Parameters:
        input_device: sounddevice device index or name substring. None uses the default.
        output_device: sounddevice device index or name substring. None uses the default.
    """

    input_device: int | str | None = None
    output_device: int | str | None = None


class SoundDeviceInputTransport(BaseInputTransport):
    """Captures microphone audio and pushes it into the pipeline."""

    _params: SoundDeviceTransportParams

    def __init__(self, params: SoundDeviceTransportParams):
        super().__init__(params)
        self._in_stream: sd.RawInputStream | None = None
        self._sample_rate = 0

    async def start(self, frame: StartFrame):
        await super().start(frame)

        if self._in_stream:
            return

        self._sample_rate = self._params.audio_in_sample_rate or frame.audio_in_sample_rate
        blocksize = int(self._sample_rate / 100) * 2  # 20ms

        self._in_stream = sd.RawInputStream(
            samplerate=self._sample_rate,
            blocksize=blocksize,
            device=self._params.input_device,
            channels=self._params.audio_in_channels,
            dtype="int16",
            callback=self._audio_in_callback,
        )
        self._in_stream.start()

        device_name = sd.query_devices(self._in_stream.device, "input")["name"]
        logger.info(f"Microphone: {device_name} @ {self._sample_rate} Hz")

        await self.set_transport_ready(frame)

    async def cleanup(self):
        await super().cleanup()
        if self._in_stream:
            self._in_stream.stop()
            self._in_stream.close()
            self._in_stream = None

    def _audio_in_callback(self, indata, frame_count, time_info, status):
        if status:
            logger.trace(f"Audio input status: {status}")

        frame = InputAudioRawFrame(
            audio=bytes(indata),
            sample_rate=self._sample_rate,
            num_channels=self._params.audio_in_channels,
        )

        asyncio.run_coroutine_threadsafe(self.push_audio_frame(frame), self.get_event_loop())


class SoundDeviceOutputTransport(BaseOutputTransport):
    """Plays pipeline audio out through the speakers."""

    _params: SoundDeviceTransportParams

    def __init__(self, params: SoundDeviceTransportParams):
        super().__init__(params)
        self._out_stream: sd.RawOutputStream | None = None
        self._sample_rate = 0
        # Writes are serialized by the pipeline, so one worker is enough.
        self._executor = ThreadPoolExecutor(max_workers=1)

    async def start(self, frame: StartFrame):
        await super().start(frame)

        if self._out_stream:
            return

        self._sample_rate = self._params.audio_out_sample_rate or frame.audio_out_sample_rate

        self._out_stream = sd.RawOutputStream(
            samplerate=self._sample_rate,
            device=self._params.output_device,
            channels=self._params.audio_out_channels,
            dtype="int16",
        )
        self._out_stream.start()

        device_name = sd.query_devices(self._out_stream.device, "output")["name"]
        logger.info(f"Speaker: {device_name} @ {self._sample_rate} Hz")

        await self.set_transport_ready(frame)

    async def cleanup(self):
        await super().cleanup()
        if self._out_stream:
            self._out_stream.stop()
            self._out_stream.close()
            self._out_stream = None

    async def write_audio_frame(self, frame: OutputAudioRawFrame) -> bool:
        if not self._out_stream:
            return False
        await self.get_event_loop().run_in_executor(
            self._executor, self._out_stream.write, frame.audio
        )
        return True


class SoundDeviceTransport(BaseTransport):
    """Local microphone + speaker transport."""

    def __init__(self, params: SoundDeviceTransportParams):
        super().__init__()
        self._params = params
        self._input: SoundDeviceInputTransport | None = None
        self._output: SoundDeviceOutputTransport | None = None

    def input(self) -> FrameProcessor:
        if not self._input:
            self._input = SoundDeviceInputTransport(self._params)
        return self._input

    def output(self) -> FrameProcessor:
        if not self._output:
            self._output = SoundDeviceOutputTransport(self._params)
        return self._output


def list_devices() -> str:
    """Render the audio devices sounddevice can see, for `--list-devices`."""
    lines = []
    for index, device in enumerate(sd.query_devices()):
        capability = []
        if device["max_input_channels"]:
            capability.append(f"in:{device['max_input_channels']}")
        if device["max_output_channels"]:
            capability.append(f"out:{device['max_output_channels']}")
        lines.append(f"  [{index}] {device['name']} ({', '.join(capability)})")
    return "\n".join(lines)