# SPDX-License-Identifier: Apache-2.0
"""Audio in: the two SPH0645 microphones on I2S3 (wiring.md §2), or a WAV file, as frames at one rate.

What is real here: the sources, the frame format, an energy VAD, speech spans, and the "mics cut" test
(the 100 kΩ pull-down on DOUT makes a cut rail read as exact zeros). What is a seam, with no model downloaded:
- `Vad`: Silero VAD (onnx) slots in here, replacing EnergyVad;
- `Recogniser`: sherpa-onnx (keyword spotter + streaming Zipformer ASR) takes speech spans, returns (voice, text);
- `Distiller`: llama.cpp with a JSON grammar turns a block's text into the rung texts store.Block expects.
Pinned model versions, the overlay for I2S3 and a WAV-in pipeline script are still backlog B-35.

Channels: MK1 (SEL → GND) is left = head, MK2 (SEL → 3V) is right = tail.
"""
from __future__ import annotations

import subprocess
import wave
from dataclasses import dataclass
from typing import Iterator, Protocol

import numpy as np

RATE = 16_000         # what the ASR wants. The SPH0645 needs BCLK 1.024–4.096 MHz = 16–64 kHz at 64 BCLK per frame,
                      # so 16 kHz is at its lower limit (VERIFY that the RK3566 I2S3 master clocks it cleanly; else 48 kHz and decimate)
FRAME = 512           # samples per channel per frame: 32 ms, the window Silero VAD takes at 16 kHz
CHANNELS = 2          # head, tail
ALSA_DEVICE = "hw:0,0"   # the I2S3 card once its overlay is on (VERIFY with `arecord -l`)


@dataclass
class Frame:
    t: float              # start, seconds since the stream began
    samples: np.ndarray   # float32, shape (FRAME, 2): head, tail, in −1…1

    @property
    def head(self) -> np.ndarray:
        return self.samples[:, 0]

    @property
    def tail(self) -> np.ndarray:
        return self.samples[:, 1]


class AudioSource(Protocol):
    rate: int
    def frames(self) -> Iterator[Frame]: ...
    def close(self) -> None: ...


def _frames(pcm: Iterator[np.ndarray], rate: int, frame: int) -> Iterator[Frame]:
    """Re-chunk blocks of (n, 2) float32 into whole frames; a short last frame is dropped."""
    buf, n = np.zeros((0, CHANNELS), np.float32), 0
    for block in pcm:
        buf = np.concatenate([buf, block])
        while len(buf) >= frame:
            yield Frame(n / rate, buf[:frame].copy())
            buf, n = buf[frame:], n + frame


class WavSource:
    """A WAV file at RATE: 16- or 32-bit PCM, mono (copied to both channels) or stereo (head, tail).
    No resampling: convert first (`sox in.wav -r 16000 out.wav`), so nothing silently changes the audio."""

    def __init__(self, path, frame: int = FRAME):
        self.path, self.frame = path, frame
        with wave.open(str(path), "rb") as w:
            self.rate, self.channels, self.width = w.getframerate(), w.getnchannels(), w.getsampwidth()
        if self.rate != RATE:
            raise ValueError(f"{path}: {self.rate} Hz; convert to {RATE} Hz first")
        if self.width not in (2, 4) or self.channels not in (1, 2):
            raise ValueError(f"{path}: {self.width * 8}-bit × {self.channels} channels; want 16/32-bit mono or stereo")

    def _pcm(self) -> Iterator[np.ndarray]:
        dt, full = ({2: np.int16, 4: np.int32}[self.width], {2: 32768.0, 4: 2147483648.0}[self.width])
        with wave.open(str(self.path), "rb") as w:
            while True:
                raw = w.readframes(4096)
                if not raw:
                    return
                x = (np.frombuffer(raw, dt).astype(np.float32) / full).reshape(-1, self.channels)
                yield np.repeat(x, 2, axis=1) if self.channels == 1 else x

    def frames(self) -> Iterator[Frame]:
        return _frames(self._pcm(), self.rate, self.frame)

    def close(self) -> None:
        pass


class AlsaSource:
    """The I2S pair through `arecord` (alsa-utils, on Radxa's Debian): no Python audio package to install.
    The SPH0645 sends 18 bits left-justified in a 32-bit slot, so S32_LE / 2³¹ is the right scale. It also has a
    DC offset; a high-pass belongs in the VAD/ASR front end (VERIFY its size at bring-up)."""

    def __init__(self, device: str = ALSA_DEVICE, rate: int = RATE, frame: int = FRAME):
        self.device, self.rate, self.frame = device, rate, frame
        self.proc: subprocess.Popen | None = None

    def command(self) -> list[str]:
        return ["arecord", "-q", "-D", self.device, "-c", str(CHANNELS), "-r", str(self.rate), "-f", "S32_LE", "-t", "raw"]

    def _pcm(self) -> Iterator[np.ndarray]:
        self.proc = subprocess.Popen(self.command(), stdout=subprocess.PIPE)
        step = 4 * CHANNELS * self.frame
        while True:
            raw = self.proc.stdout.read(step)
            if not raw:
                return
            raw = raw[:len(raw) - len(raw) % (4 * CHANNELS)]
            yield (np.frombuffer(raw, np.int32).astype(np.float32) / 2147483648.0).reshape(-1, CHANNELS)

    def frames(self) -> Iterator[Frame]:
        return _frames(self._pcm(), self.rate, self.frame)

    def close(self) -> None:
        if self.proc:
            self.proc.terminate()
            self.proc = None


def mics_cut(f: Frame) -> bool:
    """Both channels exact zeros: the mic slide is off (the eye is dark too). display shows it as the room off."""
    return not np.any(f.samples)


# --- the seams ------------------------------------------------------------------------
class Vad(Protocol):
    def __call__(self, f: Frame) -> float: ...        # speech probability 0–1


@dataclass
class EnergyVad:
    """A stand-in until Silero VAD: speech when the louder channel's RMS is above `dbfs`. Real and model-free,
    good enough to drive the speaking light on a laptop; not good enough to decide what is speech in a room."""
    dbfs: float = -45.0

    def __call__(self, f: Frame) -> float:
        rms = float(np.sqrt(np.mean(f.samples.astype(np.float64) ** 2, axis=0)).max())
        return 1.0 if rms > 0 and 20 * np.log10(rms) > self.dbfs else 0.0


class Recogniser(Protocol):
    """sherpa-onnx goes here: keyword spotter on one core for "goldfish", streaming ASR on speech spans.
    Returns (voice label, text) pairs for Goldfish.hear(); the voice label comes from on-device diarisation."""
    def __call__(self, frames: list[Frame]) -> list[tuple[str, str]]: ...


class Distiller(Protocol):
    """llama.cpp goes here: one block's transcript → {3: attributed, 2: gist, 1: outcomes} for store.Block.texts.
    The output is machine-generated and the note says so."""
    def __call__(self, transcript: list[tuple[str, str]]) -> dict[int, str]: ...


def spans(frames: Iterator[Frame], vad: Vad, threshold: float = 0.5, hang: float = 0.3) -> Iterator[tuple[float, float]]:
    """Speech spans (t0, t1) in stream seconds: frames above threshold, joined across gaps shorter than `hang`."""
    start = end = None
    for f in frames:
        dur = len(f.samples) / RATE
        if vad(f) >= threshold:
            if start is None:
                start = f.t
            end = f.t + dur
        elif start is not None and f.t - end >= hang:
            yield (start, end)
            start = end = None
    if start is not None:
        yield (start, end)