"""Playback and recording with Qt Multimedia (bundled with PySide6; WASAPI on Windows,
PulseAudio/PipeWire on Linux). Audio is exchanged as float32 numpy arrays (frames, channels).

Set ACKIT_NO_AUDIO=1 to disable devices (headless tests); `Recorder.start_simulated` feeds a recording from
an array in real time, so the recording logic is testable without a microphone.
"""
import os

import numpy as np
from PySide6.QtCore import QBuffer, QByteArray, QIODevice, QObject, QTimer, Signal
from PySide6.QtMultimedia import QAudio, QAudioFormat, QAudioSink, QAudioSource, QMediaDevices

from ackit import audio


def audio_enabled():
    return os.environ.get("ACKIT_NO_AUDIO") != "1"


def output_devices():
    return list(QMediaDevices.audioOutputs()) if audio_enabled() else []


def input_devices():
    return list(QMediaDevices.audioInputs()) if audio_enabled() else []


def find_device(devices, description):
    return next((d for d in devices if d.description() == description), None)


def _to_int16(x):
    return (np.clip(x, -1.0, 1.0) * 32767.0).astype("<i2").tobytes()


class Player(QObject):
    """Plays one buffer at a time. Signals: position(seconds), playing(bool), error(str)."""
    position = Signal(float)
    playing = Signal(bool)
    error = Signal(str)

    def __init__(self, parent=None):
        super().__init__(parent)
        self.device = None
        self.sink = None
        self.buffer = None
        self.data = QByteArray()
        self.fmt = None
        self.base = 0.0
        self.duration = 0.0
        self.loop = False
        self.volume = 1.0
        self.label = None
        self.timer = QTimer(self)
        self.timer.setInterval(40)
        self.timer.timeout.connect(self._tick)

    def set_device(self, device):
        self.stop()
        self.device = device

    def _device(self):
        if self.device is not None and not self.device.isNull():
            return self.device
        d = QMediaDevices.defaultAudioOutput() if audio_enabled() else None
        return d if d is not None and not d.isNull() else None

    def play(self, x, rate, start=0.0, loop=False, label=None):
        """x: float32 (n, ch). Converted to the device's format (rate, channels, 16-bit)."""
        self.stop()
        dev = self._device()
        if dev is None:
            self.error.emit("No audio output device available.")
            return False
        pref = dev.preferredFormat()
        fmt = QAudioFormat()
        fmt.setSampleRate(pref.sampleRate() or 48000)
        fmt.setChannelCount(min(2, pref.channelCount() or 2))
        fmt.setSampleFormat(QAudioFormat.SampleFormat.Int16)
        if not dev.isFormatSupported(fmt):
            fmt.setChannelCount(2)
            fmt.setSampleRate(48000)
        y = audio.to_channels(audio.resample(x, rate, fmt.sampleRate()), fmt.channelCount())
        self.fmt, self.loop, self.label = fmt, loop, label
        self.data = QByteArray(_to_int16(y))
        self.duration = len(y) / float(fmt.sampleRate())
        self.sink = QAudioSink(dev, fmt, self)
        self.sink.setVolume(self.volume)
        self.sink.stateChanged.connect(self._state)
        self._start_at(start)
        return True

    def _start_at(self, seconds):
        seconds = max(0.0, min(seconds, self.duration))
        if self.buffer is not None:
            self.buffer.close()
        self.buffer = QBuffer(self)
        self.buffer.setData(self.data)
        self.buffer.open(QIODevice.OpenModeFlag.ReadOnly)
        frame = self.fmt.bytesPerFrame()
        self.buffer.seek(int(seconds * self.fmt.sampleRate()) * frame)
        self.base = seconds
        self.sink.start(self.buffer)
        self.timer.start()
        self.playing.emit(True)

    def seek(self, seconds):
        if self.sink is not None:
            self.sink.stop()
            self._start_at(seconds)

    def set_volume(self, v):
        self.volume = float(v)
        if self.sink is not None:
            self.sink.setVolume(self.volume)

    def is_playing(self):
        return self.sink is not None and self.sink.state() == QAudio.State.ActiveState

    def current_time(self):
        if self.sink is None:
            return 0.0
        return self.base + self.sink.processedUSecs() / 1e6

    def _tick(self):
        self.position.emit(min(self.current_time(), self.duration))

    def _state(self, st):
        if st == QAudio.State.IdleState:  # reached the end of the buffer
            if self.loop:
                self.sink.stop()
                self._start_at(0.0)
            else:
                self.stop()

    def stop(self):
        self.timer.stop()
        if self.sink is not None:
            try:
                self.sink.stateChanged.disconnect(self._state)
            except (RuntimeError, TypeError):
                pass
            self.sink.stop()
            self.sink.deleteLater()
            self.sink = None
            self.playing.emit(False)
        if self.buffer is not None:
            self.buffer.close()
            self.buffer = None


class Recorder(QObject):
    """Records from an input device. Signals: level(peak 0..1), elapsed(seconds), error(str)."""
    level = Signal(float)
    elapsed = Signal(float)
    error = Signal(str)

    def __init__(self, parent=None):
        super().__init__(parent)
        self.device = None
        self.source = None
        self.io = None
        self.chunks = []
        self.rate = 48000
        self.channels = 1
        self.fmt_float = False
        self._sim = None

    def set_device(self, device):
        self.device = device

    def is_recording(self):
        return self.source is not None or self._sim is not None

    def start(self):
        self.chunks = []
        dev = self.device if (self.device is not None and not self.device.isNull()) else (
            QMediaDevices.defaultAudioInput() if audio_enabled() else None)
        if dev is None or dev.isNull():
            self.error.emit("No audio input device available.")
            return False
        fmt = QAudioFormat()
        fmt.setSampleRate(48000)
        fmt.setChannelCount(1)
        fmt.setSampleFormat(QAudioFormat.SampleFormat.Int16)
        if not dev.isFormatSupported(fmt):
            fmt = dev.preferredFormat()
        self.rate, self.channels = fmt.sampleRate(), fmt.channelCount()
        self.fmt_float = fmt.sampleFormat() == QAudioFormat.SampleFormat.Float
        self.source = QAudioSource(dev, fmt, self)
        self.io = self.source.start()
        if self.io is None:
            self.error.emit("Could not open the input device.")
            self.source = None
            return False
        self.io.readyRead.connect(self._read)
        return True

    def _read(self):
        if self.io is None:
            return
        self._ingest(bytes(self.io.readAll()))

    def _ingest(self, raw):
        if not raw:
            return
        if self.fmt_float:
            x = np.frombuffer(raw, dtype="<f4")
        else:
            x = np.frombuffer(raw[: len(raw) // 2 * 2], dtype="<i2").astype(np.float32) / 32768.0
        x = x[: len(x) // self.channels * self.channels].reshape(-1, self.channels)
        self.chunks.append(x)
        self.level.emit(float(np.max(np.abs(x))) if len(x) else 0.0)
        self.elapsed.emit(sum(len(c) for c in self.chunks) / float(self.rate))

    # simulated input (tests / no microphone) ----------------------------------------------------
    def start_simulated(self, x, rate, realtime=True):
        """Feed `x` as if it came from a microphone (20 ms blocks)."""
        self.chunks = []
        self.rate, self.channels, self.fmt_float = rate, x.shape[1], True
        blocks = [x[i:i + rate // 50] for i in range(0, len(x), rate // 50)]
        if not realtime:
            for b in blocks:
                self._ingest(np.ascontiguousarray(b, dtype="<f4").tobytes())
            return True
        self._sim = (QTimer(self), iter(blocks))
        def step():
            b = next(self._sim[1], None) if self._sim else None
            if b is None:
                if self._sim:
                    self._sim[0].stop()
                return
            self._ingest(np.ascontiguousarray(b, dtype="<f4").tobytes())
        self._sim[0].timeout.connect(step)
        self._sim[0].start(20)
        return True

    def stop(self):
        """-> (float32 array (n, ch), rate) of everything recorded."""
        if self.source is not None:
            self._read()
            self.source.stop()
            self.source.deleteLater()
            self.source = None
            self.io = None
        if self._sim is not None:
            self._sim[0].stop()
            self._sim = None
        x = np.concatenate(self.chunks) if self.chunks else np.zeros((0, self.channels), np.float32)
        return x, self.rate