"""LocalAgreement-2 streaming commit core (ufal/whisper_streaming).

Extracted from LiveTranscriptionService so the same, field-proven commit state
machine drives both the local live dictation loop AND the server's /v1/realtime
WebSocket session. The only thing that differs between the two is the audio
source (local mic snapshot vs. PCM frames over the wire) and how the committed
delta is delivered — both of which stay in the caller.

The committer owns: window selection out of a growing buffer (with hard-cap /
eviction realignment), the LocalAgreement-2 agreement test, re-anchoring, and
seam de-duplication. It is audio-source- and transport-agnostic; the caller
injects a `transcribe_fn(window: np.ndarray, language) -> list[word-dict]`
where each word is {"raw": str, "norm": str, "end": float}.
"""
from __future__ import annotations

from typing import Callable

import numpy as np
from loguru import logger

SAMPLE_RATE = 16000
POLL_INTERVAL = 0.6          # seconds between re-transcriptions
MIN_NEW_AUDIO_S = 0.4        # skip a poll until this much NEW audio arrived
REANCHOR_SECONDS = 16.0      # cut committed audio once the window exceeds this
# Hard cap on the audio length ever fed to the model. Keeps every inference
# well under Windows' ~2 s GPU TDR limit even when nothing commits (silence,
# noise, unstable audio) so the window can never grow without bound.
HARD_CAP_SECONDS = 24.0
FORCE_KEEP_SECONDS = 12.0    # audio kept when force-cutting an uncommitted window
MIN_AUDIO_TO_START_S = 0.4   # need a little audio before the first inference

TranscribeFn = Callable[[np.ndarray, "str | None"], list[dict]]


def trim_trailing_silence(audio: np.ndarray, threshold: float = 0.005) -> np.ndarray:
    """Drop trailing near-silence (keeps a 100 ms tail) so a short final window
    does not make Whisper hallucinate an extra trailing word."""
    if len(audio) < SAMPLE_RATE // 2:
        return audio
    abs_a = np.abs(audio)
    chunk = 800
    end = len(audio)
    min_keep = SAMPLE_RATE // 2
    while end > min_keep:
        start = max(end - chunk, 0)
        if float(np.max(abs_a[start:end])) > threshold:
            return audio[:min(end + 1600, len(audio))]
        end = start
    return audio


def commit_agreement(prev: list[str], cur: list[str], already: int) -> int:
    """LocalAgreement-2: how many NEW words (beyond `already` committed) agree
    between the previous and current hypothesis. Returns the count of newly
    stable words at the front of the uncommitted tail."""
    prev_tail = prev[already:]
    cur_tail = cur[already:]
    n = 0
    for a, b in zip(prev_tail, cur_tail):
        if a == b:
            n += 1
        else:
            break
    return n


class LocalAgreementCommitter:
    """Per-session streaming commit state. Not thread-safe; drive it from one
    task/thread. `anchor_samples` is an ABSOLUTE 16 kHz index from the start of
    the stream; a snapshot's `base` maps it onto the (possibly evicted) buffer."""

    def __init__(self, language: str | None = None) -> None:
        self.language = language
        self.anchor_samples = 0
        self.committed: list[dict] = []   # committed words in the CURRENT anchor
        self.prev_norm: list[str] = []    # previous hypothesis (normalized)
        self.last_infer_abs = 0           # absolute 16k index at last inference
        # Normalized text of the last word committed before the current anchor,
        # used to drop a duplicate first word at a re-anchor seam (Whisper word
        # end timestamps are approximate, so the cut can re-include one word).
        self.seam_norm = ""

    def reset(self) -> None:
        self.anchor_samples = 0
        self.committed = []
        self.prev_norm = []
        self.last_infer_abs = 0
        self.seam_norm = ""

    def window(self, full: np.ndarray, base: int) -> np.ndarray | None:
        """Slice the current window out of a snapshot, mapping the absolute
        anchor through the snapshot's base offset and bounding its length."""
        total = base + len(full)  # absolute samples available so far
        # Anchor scrolled off the front (deque eviction on a very long session):
        # realign to the buffer start. The evicted region was uncommitted, so
        # nothing already typed is lost or duplicated.
        if self.anchor_samples < base:
            self.anchor_samples = base
            self.committed = []
            self.prev_norm = []
            self.seam_norm = ""
        # Hard cap: never feed the model more than HARD_CAP_SECONDS. When the
        # window is over the cap AND nothing is committed to cut at (silence /
        # unstable audio), force-drop the oldest audio down to FORCE_KEEP.
        cap = int(HARD_CAP_SECONDS * SAMPLE_RATE)
        if total - self.anchor_samples > cap and not self.committed:
            self.anchor_samples = total - int(FORCE_KEEP_SECONDS * SAMPLE_RATE)
            self.prev_norm = []
            logger.debug("Live window hard-capped (no commits); anchor -> {}",
                         self.anchor_samples)
        start = self.anchor_samples - base
        if start < 0 or start >= len(full):
            return None
        return full[start:]

    def process(
        self, full: np.ndarray, base: int, transcribe_fn: TranscribeFn, *,
        final: bool,
    ) -> str:
        """Run one window through inference + LocalAgreement commit. Returns the
        newly committed delta text (may be ""). Sets last_infer_abs when it
        actually infers so callers can throttle on new-audio arrival."""
        total = base + len(full)
        window = self.window(full, base)
        floor = int((0.1 if final else MIN_AUDIO_TO_START_S) * SAMPLE_RATE)
        if window is None or len(window) < floor:
            return ""
        self.last_infer_abs = total
        if final:
            window = trim_trailing_silence(window)

        words = transcribe_fn(window, self.language)
        if not words:
            return ""
        # Drop a duplicate first word left at a re-anchor seam.
        if self.seam_norm and words[0]["norm"] == self.seam_norm:
            words = words[1:]
        cur_norm = [w["norm"] for w in words]
        self.seam_norm = ""  # consumed (only the first word after a cut)
        already = len(self.committed)
        if already > len(words):
            # a re-transcription returned fewer words than already committed
            # (segmentation jitter) — keep prior commits, wait for recovery.
            logger.debug("Live hypothesis shorter than committed; skipping")
            self.prev_norm = cur_norm
            return ""

        if final:
            newly = words[already:]            # commit the whole remaining tail
        else:
            n = commit_agreement(self.prev_norm, cur_norm, already)
            newly = words[already:already + n]
        self.prev_norm = cur_norm

        delta = ""
        if newly:
            self.committed.extend(newly)
            delta = "".join(w["raw"] for w in newly)

        if not final:
            self.maybe_reanchor(len(window))
        return delta

    def maybe_reanchor(self, window_samples: int) -> None:
        """Once the window is long and words are committed, cut the audio at the
        last committed word so the next inference stays short (bounded latency).
        Cutting exactly at a committed (already-typed) boundary means no word is
        re-typed or lost."""
        if window_samples < int(REANCHOR_SECONDS * SAMPLE_RATE):
            return
        if not self.committed:
            return
        cut_time = self.committed[-1]["end"]  # seconds into the window
        cut_samples = int(cut_time * SAMPLE_RATE)
        if cut_samples <= 0:
            return
        self.seam_norm = self.committed[-1]["norm"]  # dedupe seam next window
        self.anchor_samples += cut_samples
        self.committed = []
        self.prev_norm = []
        logger.debug("Live re-anchored at +{:.1f}s (anchor now {})",
                     cut_time, self.anchor_samples)
