"""Shared plumbing for the Sanas SDK examples.

Keeps the non-SDK mechanics (activation-free API-key login is inline, pipeline
readiness, real-time pacing, the synchronous feed+drain loop) out of the example
scripts so each one stays focused on the SDK calls it demonstrates.

The SDK's ``process_frame`` is synchronous: it returns the processed frame
directly (see IAudioProcessor::ProcessFrame). The model buffers ~50-100 ms
internally, so the output lags the input; after the real audio is exhausted we
push silence to pull the buffered tail back out.
"""

import array
import math
import threading
import time

import sanas

from wav_utils import read_wav

_CHUNK_MS = 20


def load_samples(path):
    """Read a WAV into interleaved float32 samples.

    Returns (samples, sample_rate, channels) where samples is an
    ``array('f')`` suitable for slicing into AudioFrame(samples=...).
    """
    float32_bytes, sample_rate, channels = read_wav(path)
    samples = array.array("f")
    samples.frombytes(float32_bytes)
    return samples, sample_rate, channels


def sleep_until(deadline):
    """Sleep until *deadline* (a time.monotonic() value), busy-waiting the last
    1 ms for tighter pacing than a bare time.sleep()."""
    remaining = deadline - time.monotonic()
    if remaining > 0.001:
        time.sleep(remaining - 0.001)
    while time.monotonic() < deadline:
        pass


class PipelineWaiter:
    """Bridges AudioAttributes.audio_pipeline_state_notify to a waitable value.

    The pipeline initializes asynchronously; callers must wait for
    PipelineState.RUNNING before feeding frames. Pass ``waiter.callback`` as
    ``audio_pipeline_state_notify`` and then call ``wait_until_running()``.

    A plain threading.Event only signals "something happened" and carries no
    value, so this uses a Condition that notifies on every state change while
    keeping the latest state in ``last_state`` — waiters wake on any transition
    and then inspect the value to decide whether to proceed, fail, or keep
    waiting.
    """

    def __init__(self, on_state=None):
        self._cond = threading.Condition()
        self._on_state = on_state
        self.last_state = None

    def callback(self, state):
        if self._on_state is not None:
            self._on_state(state)
        with self._cond:
            self.last_state = state
            self._cond.notify_all()

    def wait_until_running(self, timeout=30.0):
        """Block until RUNNING; raise on NOT_RUNNING or timeout."""
        deadline = None if timeout is None else time.monotonic() + timeout
        with self._cond:
            while self.last_state != sanas.PipelineState.RUNNING:
                if self.last_state == sanas.PipelineState.NOT_RUNNING:
                    raise RuntimeError(
                        f"pipeline did not start (last state: {self.last_state})"
                    )
                remaining = None if deadline is None else deadline - time.monotonic()
                if remaining is not None and remaining <= 0:
                    raise TimeoutError("pipeline did not reach RUNNING in time")
                self._cond.wait(timeout=remaining)
        return self.last_state


def feed_and_drain(
    proc,
    samples,
    sample_rate,
    channels,
    chunk_ms=_CHUNK_MS,
    realtime=True,
    progress=None,
    drain_seconds=None,
    collect_output=True,
):
    """Feed ``samples`` through ``proc`` in fixed chunks and drain the tail.

    Returns a dict with:
      - ``output``: processed interleaved float32 bytes (empty when
        ``collect_output`` is False)
      - ``input_frames``: frames (samples/channel) pushed

    When ``realtime`` is True the feed is paced at the chunk cadence against a
    fixed monotonic schedule (self-correcting, no drift). Set False to push as
    fast as possible (throughput benchmarking).

    ``collect_output`` controls whether the processed frames are retained. When
    False, the output is fed through but not copied/accumulated (``output`` is
    empty) — use this for memory/soak benchmarking where the processed audio is
    not needed, so the run's memory footprint isn't inflated by held buffers.

    ``drain_seconds`` controls how the buffered tail is flushed with silence:
      - None (default): stop after a short run of empty frames — fine for the
        low-latency models (NC/SE, ~50-100 ms tail).
      - a number: pump silence for a fixed window of that many seconds, paced at
        the chunk cadence, ignoring the empty-frame early-exit. Use this for
        high-latency pipelines (e.g. Language Translation, ~3-4 s) so the remote
        has wall-clock time to return its output.
    """
    if sample_rate <= 0 or channels <= 0:
        raise ValueError(
            f"sample_rate and channels must be positive "
            f"(got sample_rate={sample_rate}, channels={channels})"
        )
    chunk_frames = int(sample_rate * chunk_ms / 1000)
    if chunk_frames < 1:
        raise ValueError(
            f"chunk_ms={chunk_ms} at sample_rate={sample_rate} yields a "
            "zero-length chunk; increase chunk_ms or sample_rate"
        )
    chunk_size = chunk_frames * channels
    chunk_seconds = chunk_frames / sample_rate

    out = [] if collect_output else None
    total_frames = len(samples) // channels

    offset = 0
    pushed_frames = 0
    next_deadline = time.monotonic()
    while offset < len(samples):
        chunk = samples[offset : offset + chunk_size]
        n = len(chunk)
        if n < chunk_size:
            # Pad the final short chunk with silence so no input audio is dropped
            # when the length isn't an exact multiple of the chunk size.
            chunk = chunk + array.array("f", bytes((chunk_size - n) * 4))
        frame = proc.process_frame(sanas.AudioFrame(samples=chunk))
        if frame.frame_count and collect_output:
            out.append(bytes(frame.samples))
        offset += chunk_size
        pushed_frames += n // channels
        if progress is not None:
            progress(pushed_frames, total_frames)
        if realtime:
            next_deadline += chunk_seconds
            sleep_until(next_deadline)

    # Drain the buffered tail with silence.
    silence = array.array("f", bytes(chunk_size * 4))
    if drain_seconds is not None:
        # Fixed real-time window: keep pumping silence for the whole duration,
        # paced like the feed, so a high-latency pipeline has wall-clock time to
        # return its output (each paced push pulls back the now-available audio).
        for _ in range(max(1, math.ceil(drain_seconds / chunk_seconds))):
            frame = proc.process_frame(sanas.AudioFrame(samples=silence))
            if frame.frame_count and collect_output:
                out.append(bytes(frame.samples))
            if realtime:
                next_deadline += chunk_seconds
                sleep_until(next_deadline)
    else:
        # Low-latency tail (NC/SE): stop once the pipeline goes quiet.
        empty_streak = 0
        for _ in range(100):
            if empty_streak >= 3:
                break
            frame = proc.process_frame(sanas.AudioFrame(samples=silence))
            if frame.frame_count:
                if collect_output:
                    out.append(bytes(frame.samples))
                empty_streak = 0
            else:
                empty_streak += 1

    return {
        "output": b"".join(out) if collect_output else b"",
        "input_frames": pushed_frames,
    }
