"""WASAPI audio capture engine (Windows only, via PyAudioWPatch)."""

from typing import Any, Dict, List, Optional

import numpy as np

from wled_controller.core.audio.base import (
    AudioCaptureEngine,
    AudioCaptureStreamBase,
    AudioDeviceInfo,
)
from wled_controller.utils import get_logger

logger = get_logger(__name__)


class WasapiCaptureStream(AudioCaptureStreamBase):
    """Audio capture stream using PyAudioWPatch (WASAPI)."""

    def __init__(self, device_index: int, is_loopback: bool, config: Dict[str, Any]):
        super().__init__(device_index, is_loopback, config)
        self._pa = None
        self._stream = None
        self._channels = config.get("channels", 2)
        self._sample_rate = config.get("sample_rate", 44100)
        self._chunk_size = config.get("chunk_size", 2048)

    @property
    def channels(self) -> int:
        return self._channels

    @property
    def sample_rate(self) -> int:
        return self._sample_rate

    @property
    def chunk_size(self) -> int:
        return self._chunk_size

    def initialize(self) -> None:
        if self._initialized:
            return

        try:
            import pyaudiowpatch as pyaudio
        except ImportError:
            raise RuntimeError("PyAudioWPatch is not installed — WASAPI engine unavailable")

        self._pa = pyaudio.PyAudio()

        if self.is_loopback:
            loopback_device = self._find_loopback_device(self._pa, self.device_index)
            if loopback_device is None:
                self._pa.terminate()
                self._pa = None
                raise RuntimeError(
                    f"No loopback device found for output device {self.device_index}"
                )
            device_idx = loopback_device["index"]
            self._channels = loopback_device["maxInputChannels"]
            self._sample_rate = int(loopback_device["defaultSampleRate"])
        else:
            device_idx = self.device_index if self.device_index >= 0 else None
            if device_idx is not None:
                dev_info = self._pa.get_device_info_by_index(device_idx)
                self._channels = max(1, dev_info["maxInputChannels"])
                self._sample_rate = int(dev_info["defaultSampleRate"])

        self._stream = self._pa.open(
            format=pyaudio.paFloat32,
            channels=self._channels,
            rate=self._sample_rate,
            input=True,
            input_device_index=device_idx,
            frames_per_buffer=self._chunk_size,
        )
        self._initialized = True
        logger.info(
            f"WASAPI stream opened: device={device_idx} loopback={self.is_loopback} "
            f"channels={self._channels} sr={self._sample_rate}"
        )

    def cleanup(self) -> None:
        if self._stream is not None:
            try:
                self._stream.stop_stream()
                self._stream.close()
            except Exception as e:
                logger.debug("WASAPI stream cleanup: %s", e)
                pass
            self._stream = None
        if self._pa is not None:
            try:
                self._pa.terminate()
            except Exception as e:
                logger.debug("PyAudio terminate during cleanup: %s", e)
                pass
            self._pa = None
        self._initialized = False

    def read_chunk(self) -> Optional[np.ndarray]:
        if self._stream is None:
            return None
        try:
            raw_data = self._stream.read(self._chunk_size, exception_on_overflow=False)
            return np.frombuffer(raw_data, dtype=np.float32)
        except Exception as e:
            logger.warning(f"WASAPI read error: {e}")
            return None

    @staticmethod
    def _find_loopback_device(pa, output_device_index: int) -> Optional[dict]:
        """Find the PyAudioWPatch loopback device for a given output device."""
        try:
            first_loopback = None
            for loopback in pa.get_loopback_device_info_generator():
                if first_loopback is None:
                    first_loopback = loopback

                if output_device_index < 0:
                    return loopback

                target_info = pa.get_device_info_by_index(output_device_index)
                if target_info["name"] in loopback["name"]:
                    return loopback

            return first_loopback
        except Exception as e:
            logger.error(f"Error finding loopback device: {e}")
            return None


class WasapiEngine(AudioCaptureEngine):
    """WASAPI audio capture engine (Windows only)."""

    ENGINE_TYPE = "wasapi"
    ENGINE_PRIORITY = 10

    @classmethod
    def is_available(cls) -> bool:
        try:
            import pyaudiowpatch  # noqa: F401
            return True
        except ImportError as e:
            logger.debug("WASAPI engine unavailable (pyaudiowpatch not installed): %s", e)
            return False

    @classmethod
    def get_default_config(cls) -> Dict[str, Any]:
        return {
            "sample_rate": 44100,
            "chunk_size": 2048,
        }

    @classmethod
    def enumerate_devices(cls) -> List[AudioDeviceInfo]:
        try:
            import pyaudiowpatch as pyaudio
        except ImportError as e:
            logger.debug("Cannot enumerate WASAPI devices (pyaudiowpatch not installed): %s", e)
            return []

        pa = None
        try:
            pa = pyaudio.PyAudio()
            wasapi_info = pa.get_host_api_info_by_type(pyaudio.paWASAPI)
            wasapi_idx = wasapi_info["index"]

            result = []
            loopback_names: set = set()
            device_count = pa.get_device_count()

            # First pass: collect input devices.  PyAudioWPatch creates
            # dedicated loopback input endpoints for output devices; these
            # show up as input devices whose name already contains
            # "[Loopback]".  We mark them as loopback and remember the name
            # so the second pass won't duplicate them.
            for i in range(device_count):
                dev = pa.get_device_info_by_index(i)
                if dev["hostApi"] != wasapi_idx:
                    continue
                if dev["maxInputChannels"] <= 0:
                    continue

                name = dev["name"]
                is_loopback = "[Loopback]" in name
                if is_loopback:
                    loopback_names.add(name)

                result.append(AudioDeviceInfo(
                    index=i,
                    name=name,
                    is_input=True,
                    is_loopback=is_loopback,
                    channels=dev["maxInputChannels"],
                    default_samplerate=dev["defaultSampleRate"],
                ))

            # Second pass: add loopback entries for output devices that
            # don't already have a dedicated loopback input endpoint.
            for i in range(device_count):
                dev = pa.get_device_info_by_index(i)
                if dev["hostApi"] != wasapi_idx:
                    continue
                if dev["maxOutputChannels"] <= 0:
                    continue

                loopback_name = f"{dev['name']} [Loopback]"
                if loopback_name in loopback_names:
                    continue  # already covered by a dedicated loopback endpoint

                result.append(AudioDeviceInfo(
                    index=i,
                    name=loopback_name,
                    is_input=False,
                    is_loopback=True,
                    channels=dev["maxOutputChannels"],
                    default_samplerate=dev["defaultSampleRate"],
                ))

            return result

        except Exception as e:
            logger.error(f"Failed to enumerate WASAPI devices: {e}", exc_info=True)
            return []
        finally:
            if pa is not None:
                try:
                    pa.terminate()
                except Exception as e:
                    logger.debug("PyAudio terminate in enumerate cleanup: %s", e)
                    pass

    @classmethod
    def create_stream(
        cls,
        device_index: int,
        is_loopback: bool,
        config: Dict[str, Any],
    ) -> WasapiCaptureStream:
        merged = {**cls.get_default_config(), **config}
        return WasapiCaptureStream(device_index, is_loopback, merged)
