"""General audio utilities for SoundDevice device management and audio streaming."""

from __future__ import annotations

import asyncio
import functools
import logging
import threading
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass
from typing import TYPE_CHECKING, Literal

from rich.text import Text

from agent_cli import constants
from agent_cli.core.utils import (
    InteractiveStopEvent,
    console,
    print_device_index,
    print_with_style,
)

if TYPE_CHECKING:
    from collections.abc import AsyncGenerator, Awaitable, Callable, Generator

    import sounddevice as sd
    from rich.live import Live

    from agent_cli import config

LOGGER = logging.getLogger(__name__)
_AUDIO_SHUTDOWN_TIMEOUT_SECONDS = 0.5


@dataclass
class StreamConfig:
    """Configuration for an audio stream."""

    rate: int
    channels: int
    dtype: str
    device: int | None
    blocksize: int
    kind: Literal["input", "output"]

    def to_stream(self) -> sd.Stream:
        """Create a SoundDevice stream from this configuration."""
        import sounddevice as sd  # noqa: PLC0415

        if self.kind == "input":
            stream_cls = sd.InputStream
        elif self.kind == "output":
            stream_cls = sd.OutputStream
        else:
            msg = f"Invalid stream kind: {self.kind}"
            raise ValueError(msg)

        return stream_cls(
            samplerate=self.rate,
            blocksize=self.blocksize,
            device=self.device,
            channels=self.channels,
            dtype=self.dtype,
        )


class _AudioTee:
    """A thread-safe class to tee a continuous audio stream into multiple asyncio queues.

    This class reads from a single audio stream in a background task and forwards
    the audio chunks to any number of dynamically added consumer queues. It is designed
    to be started once and run for the lifetime of the stream.
    """

    def __init__(
        self,
        stream: sd.InputStream,
        stop_event: InteractiveStopEvent,
        logger: logging.Logger,
    ) -> None:
        """Initialize the AudioTee."""
        self.stream = stream
        self.stop_event = stop_event
        self.logger = logger
        self.queues: list[asyncio.Queue[bytes | None]] = []
        self._task: asyncio.Task | None = None
        self._stop_tee_event = asyncio.Event()
        self._lock = asyncio.Lock()  # For thread-safe modification of the queues list

    async def add_queue(self) -> asyncio.Queue[bytes | None]:
        queue: asyncio.Queue[bytes | None] = asyncio.Queue()
        async with self._lock:
            self.queues.append(queue)
        self.logger.debug("Added a queue to the tee. Total queues: %d", len(self.queues))
        return queue

    async def remove_queue(self, queue: asyncio.Queue[bytes | None]) -> None:
        async with self._lock:
            if queue in self.queues:
                self.queues.remove(queue)
        # Signal the end of the stream for this specific queue consumer
        await queue.put(None)
        self.logger.debug("Removed a queue from the tee. Total queues: %d", len(self.queues))

    async def _run(self) -> None:
        """The main background task that reads from the stream and pushes to all queues."""
        self.logger.debug("Starting continuous audio reading task.")
        try:
            while not self.stop_event.is_set() and not self._stop_tee_event.is_set():
                # sd.InputStream.read() blocks until data is available
                # We run it in a thread to avoid blocking the event loop
                data, _overflow = await asyncio.to_thread(
                    self.stream.read,
                    constants.AUDIO_CHUNK_SIZE,
                )
                chunk = data.tobytes()
                # Lock the queue list while iterating to prevent modification during iteration
                async with self._lock:
                    for queue in self.queues:
                        await queue.put(chunk)
        except Exception:
            self.logger.exception("Error reading audio stream")
        finally:
            # Signal the end of the stream to all remaining consumers
            self.logger.debug("Stopping audio reading task and signaling all consumers.")
            async with self._lock:
                for queue in self.queues:
                    await queue.put(None)

    def start(self) -> None:
        """Start the background reading task."""
        if self._task is None:
            self._task = asyncio.create_task(self._run())

    async def stop(self) -> None:
        """Stop the background reading task gracefully."""
        if self._task and not self._task.done():
            self._stop_tee_event.set()
            await self._task
        self.logger.debug("Audio tee stopped successfully.")


@asynccontextmanager
async def tee_audio_stream(
    stream: sd.InputStream,
    stop_event: InteractiveStopEvent,
    logger: logging.Logger,
) -> AsyncGenerator[_AudioTee, None]:
    """Context manager for an AudioTee."""
    tee = _AudioTee(stream, stop_event, logger)
    tee.start()
    try:
        yield tee
    finally:
        await tee.stop()


async def read_from_queue(
    queue: asyncio.Queue[bytes | None],
    chunk_handler: Callable[[bytes], None] | Callable[[bytes], Awaitable[None]],
    logger: logging.Logger,
) -> None:
    """Read audio chunks from a queue and call a handler."""
    while True:
        chunk = await queue.get()
        if chunk is None:
            break
        if asyncio.iscoroutinefunction(chunk_handler):
            await chunk_handler(chunk)
        else:
            chunk_handler(chunk)
        logger.debug("Processed %d byte(s) of audio from queue", len(chunk))


@contextmanager
def open_audio_stream(
    config: StreamConfig,
) -> Generator[sd.Stream, None, None]:
    """Context manager for a SoundDevice stream that ensures it's properly closed."""
    stream = config.to_stream()
    stream.start()
    try:
        yield stream
    finally:
        # Use abort for teardown to avoid PortAudio/CoreAudio stop deadlocks.
        _run_with_timeout(
            lambda: stream.abort(ignore_errors=True),
            action="audio stream.abort()",
        )
        _run_with_timeout(
            lambda: stream.close(ignore_errors=True),
            action="audio stream.close()",
        )


def _run_with_timeout(
    operation: Callable[[], None],
    *,
    action: str,
    timeout_seconds: float = _AUDIO_SHUTDOWN_TIMEOUT_SECONDS,
) -> bool:
    """Run a blocking operation in a daemon thread with a timeout guard."""
    done = threading.Event()
    errors: list[Exception] = []

    def _invoke() -> None:
        try:
            operation()
        except Exception as exc:  # pragma: no cover - defensive logging path
            errors.append(exc)
        finally:
            done.set()

    thread = threading.Thread(
        target=_invoke,
        name="agent-cli-audio-shutdown",
        daemon=True,
    )
    thread.start()

    if not done.wait(timeout_seconds):
        LOGGER.warning(
            "Timed out after %.2fs waiting for %s; continuing.",
            timeout_seconds,
            action,
        )
        return False

    if errors:
        LOGGER.warning("%s failed: %s", action, errors[0])
        return False

    return True


async def read_audio_stream(
    stream: sd.InputStream,
    stop_event: InteractiveStopEvent,
    chunk_handler: Callable[[bytes], None] | Callable[[bytes], Awaitable[None]],
    logger: logging.Logger,
    *,
    live: Live | None = None,
    quiet: bool = False,
    progress_message: str = "Processing audio",
    progress_style: str = "blue",
) -> None:
    """Core audio reading function - reads chunks and calls handler.

    This is the single source of truth for audio reading logic.
    All other audio functions should use this to avoid duplication.

    Args:
        stream: SoundDevice InputStream
        stop_event: Event to stop reading
        chunk_handler: Function to handle each chunk (sync or async)
        logger: Logger instance
        live: Rich Live display for progress
        quiet: If True, suppress console output
        progress_message: Message to display
        progress_style: Rich style for progress

    """
    try:
        seconds_streamed = 0.0
        while not stop_event.is_set():
            data, _overflow = await asyncio.to_thread(
                stream.read,
                constants.AUDIO_CHUNK_SIZE,
            )
            chunk = data.tobytes()

            # Handle chunk (sync or async)
            if asyncio.iscoroutinefunction(chunk_handler):
                await chunk_handler(chunk)
            else:
                chunk_handler(chunk)

            logger.debug("Processed %d byte(s) of audio", len(chunk))

            # Update progress display
            seconds_streamed += len(chunk) / (constants.AUDIO_RATE * constants.AUDIO_CHANNELS * 2)
            if live and not quiet:
                if stop_event.ctrl_c_pressed:
                    msg = f"Ctrl+C pressed. Stopping {progress_message.lower()}..."
                    live.update(Text(msg, style="yellow"))
                else:
                    live.update(
                        Text(
                            f"{progress_message}... ({seconds_streamed:.1f}s)",
                            style=progress_style,
                        ),
                    )

    except Exception:
        logger.exception("Error reading audio")


def setup_input_stream(
    input_device_index: int | None,
) -> StreamConfig:
    """Get standard audio input stream configuration.

    Args:
        input_device_index: Input device index

    Returns:
        StreamConfig for open_audio_stream

    """
    return StreamConfig(
        dtype=constants.AUDIO_FORMAT_STR,
        channels=constants.AUDIO_CHANNELS,
        rate=constants.AUDIO_RATE,
        kind="input",
        blocksize=constants.AUDIO_CHUNK_SIZE,
        device=input_device_index,
    )


def setup_output_stream(
    output_device_index: int | None,
    *,
    sample_rate: int | None = None,
    sample_width: int | None = None,
    channels: int | None = None,
) -> StreamConfig:
    """Get standard audio output stream configuration.

    Args:
        output_device_index: Output device index
        sample_rate: Custom sample rate (defaults to config)
        sample_width: Custom sample width in bytes (defaults to config)
        channels: Custom channel count (defaults to config)

    Returns:
        StreamConfig for open_audio_stream

    """
    # Map sample width to dtype if necessary. 2 -> int16.
    dtype = constants.AUDIO_FORMAT_STR
    if sample_width == 1:
        dtype = "int8"
    elif sample_width == 4:  # noqa: PLR2004
        dtype = "int32"

    return StreamConfig(
        dtype=dtype,
        channels=channels or constants.AUDIO_CHANNELS,
        rate=sample_rate or constants.AUDIO_RATE,
        kind="output",
        blocksize=constants.AUDIO_CHUNK_SIZE,
        device=output_device_index,
    )


@functools.cache
def _get_all_devices() -> list[dict]:
    """Get information for all audio devices with caching.

    Returns:
        List of device info dictionaries with added 'index' field

    """
    import sounddevice as sd  # noqa: PLC0415

    devices = []
    try:
        query_result = sd.query_devices()
        # sd.query_devices() returns a DeviceList (list-like) of dicts
        for i, device in enumerate(query_result):
            device_info = dict(device)
            device_info["index"] = i
            devices.append(device_info)
    except Exception:
        # Fallback or empty if no devices found/error
        logging.getLogger(__name__).exception("Error querying audio devices")
    return devices


def _get_device_by_index(input_device_index: int) -> dict:
    """Get device info by index from cached device list.

    Args:
        input_device_index: Device index to look up

    Returns:
        Device info dictionary

    Raises:
        ValueError: If device index is not found

    """
    for device in _get_all_devices():
        if device["index"] == input_device_index:
            return device
    msg = f"Device index {input_device_index} not found"
    raise ValueError(msg)


def _list_input_devices() -> None:
    """Print a numbered list of available input devices."""
    console.print("[bold]Available input devices:[/bold]")
    for device in _get_all_devices():
        if device.get("max_input_channels", 0) > 0:
            console.print(f"  [yellow]{device['index']}[/yellow]: {device['name']}")


def _list_output_devices() -> None:
    """Print a numbered list of available output devices."""
    console.print("[bold]Available output devices:[/bold]")
    for device in _get_all_devices():
        if device.get("max_output_channels", 0) > 0:
            console.print(f"  [yellow]{device['index']}[/yellow]: {device['name']}")


def list_all_devices() -> None:
    """Print a numbered list of all available audio devices with their capabilities."""
    _list_input_devices()
    console.print()
    _list_output_devices()


def _in_or_out_device(
    input_device_name: str | None,
    input_device_index: int | None,
    key: str,
    what: str,
) -> tuple[int | None, str | None]:
    """Find an input device by a prioritized, comma-separated list of keywords."""
    if input_device_name is None and input_device_index is None:
        return None, None

    if input_device_index is not None:
        info = _get_device_by_index(input_device_index)
        return input_device_index, info.get("name")
    assert input_device_name is not None
    search_terms = [term.strip().lower() for term in input_device_name.split(",") if term.strip()]

    if not search_terms:
        msg = "Device name string is empty or contains only whitespace."
        raise ValueError(msg)

    devices = []
    for device in _get_all_devices():
        device_info_name = device.get("name")
        if device_info_name and device.get(key, 0) > 0:
            devices.append((device["index"], device_info_name))

    for term in search_terms:
        for index, name in devices:
            if term in name.lower():
                return index, name

    msg = f"No {what} device found matching any of the keywords in {input_device_name!r}"
    raise ValueError(msg)


def _input_device(
    input_device_name: str | None,
    input_device_index: int | None,
) -> tuple[int | None, str | None]:
    """Find an input device by a prioritized, comma-separated list of keywords."""
    return _in_or_out_device(input_device_name, input_device_index, "max_input_channels", "input")


def _output_device(
    input_device_name: str | None,
    input_device_index: int | None,
) -> tuple[int | None, str | None]:
    """Find an output device by a prioritized, comma-separated list of keywords."""
    return _in_or_out_device(
        input_device_name,
        input_device_index,
        "max_output_channels",
        "output",
    )


def setup_devices(
    general_cfg: config.General,
    audio_in_cfg: config.AudioInput | None,
    audio_out_cfg: config.AudioOutput | None,
) -> tuple[int | None, str | None, int | None] | None:
    """Handle device listing and setup."""
    if general_cfg.list_devices:
        list_all_devices()
        return None

    # Setup input device
    input_device_index, input_device_name = _input_device(
        audio_in_cfg.input_device_name if audio_in_cfg else None,
        audio_in_cfg.input_device_index if audio_in_cfg else None,
    )
    if not general_cfg.quiet:
        print_device_index(input_device_index, input_device_name)

    # Setup output device for TTS if enabled
    tts_output_device_index = audio_out_cfg.output_device_index if audio_out_cfg else None
    if (
        audio_out_cfg
        and audio_out_cfg.enable_tts
        and (audio_out_cfg.output_device_name or audio_out_cfg.output_device_index)
    ):
        tts_output_device_index, tts_output_device_name = _output_device(
            audio_out_cfg.output_device_name,
            audio_out_cfg.output_device_index,
        )
        if tts_output_device_index is not None and not general_cfg.quiet:
            msg = f"🔊 TTS output device [bold yellow]{tts_output_device_index}[/bold yellow] ([italic]{tts_output_device_name}[/italic])"
            print_with_style(msg)

    return input_device_index, input_device_name, tts_output_device_index
