"""Streaming capture conditioning and pause-aware chunking for Blurt.

The class deliberately has no UI or sounddevice dependencies.  ``feed`` may be
called with whatever block size the input device supplies; all public sample
positions are in the 16 kHz output timebase.
"""

from __future__ import annotations

from dataclasses import dataclass
from math import gcd
import logging

import numpy as np
from scipy import signal


OUTPUT_RATE = 16_000
FRAME = OUTPUT_RATE // 50  # 20 ms
MIN_SECONDS = 8.0
TARGET_SECONDS = 20.0
MAX_SECONDS = 45.0
EARLY_PAUSE_SECONDS = 0.8
LATE_PAUSE_SECONDS = 0.5
HIGHPASS_HZ = 70.0
TARGET_RMS = 10 ** (-20.0 / 20.0)
PEAK_CEILING = 10 ** (-1.0 / 20.0)
MAX_GAIN = 10 ** (30.0 / 20.0)


@dataclass(frozen=True)
class AudioChunk:
    """A contiguous section of conditioned audio in the session timebase."""

    samples: np.ndarray
    start: int
    end: int
    speech: bool


def _mono(frames: np.ndarray, channels: int) -> np.ndarray:
    data = np.asarray(frames, dtype=np.float32)
    if data.ndim == 1:
        if channels != 1:
            raise ValueError("one-dimensional frames are valid only for mono capture")
        return data
    if data.ndim != 2 or data.shape[1] != channels:
        raise ValueError(f"expected frames shaped (n, {channels})")
    return data.mean(axis=1, dtype=np.float32)


def _levels(samples: np.ndarray) -> np.ndarray:
    count = len(samples) // FRAME
    if not count:
        return np.empty(0, dtype=np.float32)
    framed = samples[: count * FRAME].reshape(count, FRAME)
    framed = framed - framed.mean(axis=1, keepdims=True)
    rms = np.sqrt(np.mean(framed * framed, axis=1))
    return (20.0 * np.log10(np.maximum(rms, 1e-7))).astype(np.float32)


def _normalize(samples: np.ndarray) -> np.ndarray:
    """Normalize one sealed, already high-pass-filtered region."""
    if not len(samples):
        return np.empty(0, dtype=np.int16)
    filtered = samples
    rms = float(np.sqrt(np.mean(filtered * filtered)))
    peak = float(np.percentile(np.abs(filtered), 99.9))
    gain = min(TARGET_RMS / rms, PEAK_CEILING / peak, MAX_GAIN) if rms and peak else 1.0
    scaled = np.clip(filtered * gain, -1.0, 1.0)
    return np.rint(scaled * 32767.0).astype(np.int16)


class StreamingAudio:
    """Convert native-rate blocks into pause-aware 16 kHz mono chunks.

    Native input is resampled in one-second, phase-aligned batches with 50 ms
    of context on both sides.  Thus callback block size does not affect output
    length, while filter-edge errors stay outside the retained middle.  A final
    partial batch is length-corrected to the rounded session duration.  Startup
    audio is deliberately retained: without evidence of a device click, an
    automatic trim risks deleting the first quiet consonant.
    """

    _GUARD = FRAME * 2
    _CONTEXT_SECONDS = 0.05

    def __init__(self, rate: int, channels: int):
        if rate <= 0 or channels <= 0:
            raise ValueError("rate and channels must be positive")
        self.rate = int(rate)
        self.channels = int(channels)
        self._pending: list[np.ndarray] = []
        self._pending_len = 0
        self._left_native = np.empty(0, dtype=np.float32)
        self._resampled = np.empty(0, dtype=np.float32)
        self._base = 0
        self._start = 0  # absolute 16 kHz sample position
        self._output_total = 0
        self._sos = signal.butter(
            2, HIGHPASS_HZ, btype="highpass", fs=OUTPUT_RATE, output="sos"
        )
        self._filter_zi = None
        self._finished = False

    def feed(self, frames: np.ndarray) -> list[AudioChunk]:
        if self._finished:
            raise RuntimeError("cannot feed a finished stream")
        mono = _mono(frames, self.channels)
        previous_available = self._base + len(self._resampled)
        if len(mono):
            self._pending.append(mono.copy())
            self._pending_len += len(mono)
            self._refresh(final=False)
        # Completed native batches already used right-side resampler context.
        stable = self._base + len(self._resampled)
        if stable == previous_available:
            return []
        return self._emit(stable, finishing=False)

    def finish(self) -> list[AudioChunk]:
        if self._finished:
            return []
        self._finished = True
        self._refresh(final=True)
        chunks = self._emit(self._base + len(self._resampled), finishing=True)
        end = self._base + len(self._resampled)
        if self._start < end:
            chunks.append(self._make_chunk(self._start, end))
            self._start = end
            self._discard_sealed()
        return chunks

    def checkpoint(self) -> np.ndarray:
        """Return a conditioned copy of the currently unsealed audio."""
        preview = self._preview()
        local = self._start - self._base
        return _normalize(preview[local:]).copy()

    def _resample(self, native: np.ndarray) -> np.ndarray:
        if self.rate == OUTPUT_RATE:
            return native.astype(np.float32, copy=True)
        divisor = gcd(self.rate, OUTPUT_RATE)
        return signal.resample_poly(
            native, OUTPUT_RATE // divisor, self.rate // divisor
        ).astype(np.float32)

    def _refresh(self, final: bool) -> None:
        batch = self.rate
        context = round(self.rate * self._CONTEXT_SECONDS)
        while self._pending_len >= batch + (0 if final else context):
            pending = np.concatenate(self._pending)
            take = min(batch, self._pending_len)
            right = min(context, self._pending_len - take)
            left = len(self._left_native)
            source = np.concatenate((self._left_native, pending[: take + right]))
            converted = self._resample(source)
            offset = round(left * OUTPUT_RATE / self.rate)
            wanted = round(take * OUTPUT_RATE / self.rate)
            self._append_filtered(converted[offset : offset + wanted])
            consumed = pending[:take]
            self._left_native = np.concatenate((self._left_native, consumed))[-context:]
            remainder = pending[take:]
            self._pending = [remainder] if len(remainder) else []
            self._pending_len = len(remainder)
        if final and self._pending_len:
            pending = np.concatenate(self._pending)
            left = len(self._left_native)
            converted = self._resample(np.concatenate((self._left_native, pending)))
            offset = round(left * OUTPUT_RATE / self.rate)
            wanted = round(len(pending) * OUTPUT_RATE / self.rate)
            tail = converted[offset : offset + wanted]
            self._append_filtered(np.pad(tail, (0, max(0, wanted - len(tail)))))
            self._pending = []
            self._pending_len = 0

    def _append_filtered(self, samples: np.ndarray) -> None:
        if not len(samples):
            return
        if self._filter_zi is None:
            self._filter_zi = signal.sosfilt_zi(self._sos) * float(samples[0])
        filtered, self._filter_zi = signal.sosfilt(
            self._sos, samples, zi=self._filter_zi
        )
        self._resampled = np.concatenate((self._resampled, filtered.astype(np.float32)))
        self._output_total += len(samples)

    def _preview(self) -> np.ndarray:
        """Resample the outstanding tail without advancing stream state."""
        if not self._pending_len:
            return self._resampled
        context = round(self.rate * self._CONTEXT_SECONDS)
        pending = np.concatenate(self._pending)
        left = len(self._left_native)
        source = np.concatenate((self._left_native, pending))
        converted = self._resample(source)
        offset = round(left * OUTPUT_RATE / self.rate)
        wanted = round(len(pending) * OUTPUT_RATE / self.rate)
        tail = converted[offset : offset + wanted]
        zi = self._filter_zi
        if zi is None and len(tail):
            zi = signal.sosfilt_zi(self._sos) * float(tail[0])
        if zi is not None and len(tail):
            tail = signal.sosfilt(self._sos, tail, zi=zi)[0]
        return np.concatenate((self._resampled, tail.astype(np.float32)))

    def _emit(self, available: int, finishing: bool) -> list[AudioChunk]:
        chunks: list[AudioChunk] = []
        while True:
            cut = self._next_cut(available)
            if cut is None:
                break
            chunks.append(self._make_chunk(self._start, cut))
            self._start = cut
            self._discard_sealed()
        # finish() leaves the sub-minimum tail to its caller so it is emitted
        # once, after every pause/forced boundary has been handled.
        return chunks

    def _next_cut(self, available: int) -> int | None:
        length = available - self._start
        minimum = int(MIN_SECONDS * OUTPUT_RATE)
        maximum = int(MAX_SECONDS * OUTPUT_RATE)
        if length < minimum:
            return None
        if length >= maximum:
            search_end = self._start + maximum
        else:
            search_end = available

        segment = self._resampled[self._start - self._base : search_end - self._base]
        levels = _levels(segment)
        if not len(levels):
            return None
        # The lower quartile follows different microphones/noise floors.  The
        # ceiling prevents ordinary quiet speech from being labelled a pause.
        floor = float(np.percentile(levels, 25))
        quiet_below = min(-45.0, floor + 6.0)
        quiet = levels <= quiet_below
        run = 0
        first_eligible = int(MIN_SECONDS * 50)
        for frame_index, is_quiet in enumerate(quiet):
            run = run + 1 if is_quiet else 0
            age = (frame_index + 1) / 50.0
            fraction = np.clip(
                (age - MIN_SECONDS) / (TARGET_SECONDS - MIN_SECONDS), 0.0, 1.0
            )
            required = EARLY_PAUSE_SECONDS + fraction * (
                LATE_PAUSE_SECONDS - EARLY_PAUSE_SECONDS
            )
            required_frames = int(np.ceil(required * 50))
            if frame_index + 1 >= first_eligible and run >= required_frames:
                # Cut in the middle of the accepted pause.  Both adjacent
                # ranges still retain their half, so coverage remains exact.
                cut_frame = frame_index + 1 - run // 2
                cut = self._start + max(minimum, cut_frame * FRAME)
                logging.info(
                    "Streaming chunk pause cut at %.2fs (age %.2fs, pause %.2fs, "
                    "threshold %.1f dBFS)",
                    cut / OUTPUT_RATE,
                    age,
                    run / 50.0,
                    quiet_below,
                )
                return cut
        if length >= maximum:
            # Prefer the quietest 20 ms frame in the final second when speech
            # has been uninterrupted, while strictly respecting the limit.
            end_frame = maximum // FRAME
            begin_frame = max(first_eligible, end_frame - 50)
            chosen = begin_frame + int(np.argmin(levels[begin_frame:end_frame]))
            cut = self._start + min((chosen + 1) * FRAME, maximum)
            logging.info(
                "Streaming chunk forced cut at %.2fs after %.2fs without a pause",
                cut / OUTPUT_RATE,
                length / OUTPUT_RATE,
            )
            return cut
        return None

    def _make_chunk(self, start: int, end: int) -> AudioChunk:
        raw = self._resampled[start - self._base : end - self._base]
        levels = _levels(raw)
        # -80 dBFS distinguishes actual/near digital silence while preserving
        # quiet but valid speech.  Consumers spool false chunks but skip upload.
        if len(levels):
            speech_level = float(np.max(levels))
        elif len(raw):
            centered = raw - float(np.mean(raw))
            speech_level = float(20 * np.log10(max(np.sqrt(np.mean(centered * centered)), 1e-7)))
        else:
            speech_level = -140.0
        speech = speech_level > -80.0
        floor = float(np.percentile(levels, 10)) if len(levels) else -140.0
        longest = run = 0
        for level in levels:
            run = run + 1 if level < -75.0 else 0
            longest = max(longest, run)
        logging.info(
            "Streaming chunk %.2f-%.2fs speech=%s floor=%.1f dBFS dead=%.2fs",
            start / OUTPUT_RATE,
            end / OUTPUT_RATE,
            speech,
            floor,
            longest / 50.0,
        )
        return AudioChunk(_normalize(raw), start, end, speech)

    def _discard_sealed(self) -> None:
        count = self._start - self._base
        if count > 0:
            self._resampled = self._resampled[count:]
            self._base = self._start