"""Adaptive buffer-ratchet chunk scheduler for chatterbox-fast.

This module is the *meat* of the streaming engine and is deliberately free of
any GPU / torch / chatterbox / FastAPI imports so the no-starvation guarantee
can be validated as a pure simulation (see ``test_scheduler.py``).

The design (docs/design/chatterbox-fast-plan.md §1):

  * Chunk 1 = the first sentence alone — generated and emitted immediately so
    first-audio latency is minimal.
  * While chunk N plays, generate chunk N+1 by greedily accumulating WHOLE
    sentences until the next sentence would push estimated gen-time past
    ``margin × audio_buffered_remaining``. Never split mid-sentence — each
    chunk stays prosodically self-coherent and joins land at natural pauses.
  * Because Chatterbox runs faster than realtime (RTF > 1), every chunk's
    playback buys wall-clock for a larger next chunk, so chunks ratchet up
    ~3× and after 2-3 joins the rest of the paragraph is one big near-full-
    context chunk. Context loss is confined to those 2-3 sentence-boundary
    joins.
  * Drive off the *measured* realtime factor and sec-per-char, not constants —
    track them live and self-correct.

The scheduler is *online*: chunk N+1's boundary depends on the RTF measured
while generating chunk N, so it cannot be precomputed. ``stream_chunks`` runs
the loop with ``generate`` and ``clock`` injected, which is what makes the
simulation possible.
"""

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Callable, Iterator, Sequence

# ── tunables ──────────────────────────────────────────────────────────────


@dataclass
class ChunkConfig:
    """Scheduler tunables. Priors are conservative; live measurement corrects."""

    # Realtime factor prior. 3090 ~3.4×, A6000 ~3.8× (plan §2). Start
    # conservative — under-estimating RTF makes chunks smaller and safer.
    rtf_prior: float = 3.4

    # Safety fraction of the buffer to spend on the next chunk's generation.
    # margin < 1 leaves headroom so estimate error doesn't starve the stream.
    margin: float = 0.8

    # Tighter margin on the FIRST transition (chunk 1 → chunk 2): the buffer is
    # smallest there, so starvation risk is highest (plan §1.5: ~0.6-0.7).
    margin_first: float = 0.65

    # Seconds of audio per character of text. ~0.060 ≈ 16.6 chars/sec speech.
    # Calibrated live per request.
    sec_per_char_prior: float = 0.060

    # EMA weight on the newest measurement when updating rtf / sec_per_char.
    ema_alpha: float = 0.4

    # If the first sentence's estimated audio exceeds this, clause-split it to
    # protect first-audio latency (the ONLY place we split below sentence
    # granularity — plan §1.1).
    max_first_sec: float = 2.0


# ── result record ─────────────────────────────────────────────────────────


@dataclass
class ChunkResult:
    """Telemetry + payload for one emitted chunk."""

    index: int
    text: str
    audio: object  # opaque payload from generate() (wav tensor, fake, …)
    audio_sec: float
    gen_time: float
    est_gen: float  # gen-time the scheduler predicted before generating
    buffer_before: float  # unplayed audio (s) when this chunk's gen started
    buffer_after: float  # unplayed audio (s) once this chunk is emitted
    drained: float  # seconds the buffer ran dry during gen (>0 ⇒ starvation)
    rtf: float  # measured RTF after this chunk
    sec_per_char: float  # measured sec/char after this chunk

    @property
    def starved(self) -> bool:
        return self.drained > 1e-9


# ── text splitting ────────────────────────────────────────────────────────

# Split after sentence-final punctuation when the next non-space looks like a
# new sentence start: a capital/digit (optionally behind an open quote) or an
# inline tag like "[laugh]". Pragmatic, not perfect: abbreviations ("Mr.",
# "e.g.") can over-split — a Phase-2 watch-out, harmless to quality (just an
# extra natural-pause join).
_SENTENCE_END = re.compile(r'(?<=[.!?])["\')\]]*\s+(?=["\'(]*(?:[A-Z0-9]|\[))')

# Clause boundaries for the first-sentence latency fallback only.
_CLAUSE_END = re.compile(r'(?<=[,;:])\s+')


def split_sentences(text: str) -> list[str]:
    """Split into sentence units, preserving punctuation. Whitespace-collapsed."""
    text = " ".join(text.split())
    if not text:
        return []
    return [s for s in (p.strip() for p in _SENTENCE_END.split(text)) if s]


def _split_clauses(sentence: str) -> list[str]:
    return [c for c in (p.strip() for p in _CLAUSE_END.split(sentence)) if c]


def protect_first_audio(units: list[str], cfg: ChunkConfig) -> list[str]:
    """Clause-split the first unit if it's too long to hit the first-audio target.

    Only the *leading* unit is split; the remainder is untouched. If the first
    sentence has no clause boundary we accept the latency rather than split
    mid-clause (quality > latency once we're past the budget).
    """
    if not units:
        return units
    if len(units[0]) * cfg.sec_per_char_prior <= cfg.max_first_sec:
        return units
    pieces = _split_clauses(units[0])
    if len(pieces) <= 1:
        return units  # nothing to split on; keep the long first sentence whole
    return pieces + units[1:]


# ── chunk planning ────────────────────────────────────────────────────────


def _est_gen_time(text: str, *, rtf: float, sec_per_char: float) -> float:
    return (len(text) * sec_per_char) / rtf


def plan_chunk(
    remaining: Sequence[str],
    buffer_remaining: float,
    *,
    margin: float,
    rtf: float,
    sec_per_char: float,
) -> tuple[str, list[str]]:
    """Greedily accumulate whole units until the next would blow the budget.

    Always returns at least one unit (never empty, never splits a unit). With
    ``buffer_remaining == 0`` (the first chunk) the budget is 0, so exactly the
    first unit is taken — which is the latency-critical chunk-1 rule.
    """
    budget = margin * buffer_remaining
    chunk = [remaining[0]]
    i = 1
    while i < len(remaining):
        candidate = " ".join(chunk + [remaining[i]])
        if _est_gen_time(candidate, rtf=rtf, sec_per_char=sec_per_char) > budget:
            break
        chunk.append(remaining[i])
        i += 1
    return " ".join(chunk), list(remaining[i:])


def relieve_leader(
    remaining: list[str],
    buffer_remaining: float,
    *,
    rtf: float,
    sec_per_char: float,
) -> list[str]:
    """Expose a too-big leading sentence's clause boundaries to avoid starvation.

    If generating ``remaining[0]`` alone would overrun the buffer (an audible
    gap), clause-split it in place so ``plan_chunk`` can pack clause-pieces up to
    the buffer — joins then land on commas (natural pauses) instead of producing
    a gap. If the sentence has no clause boundary it's returned unchanged: we do
    NOT split mid-clause (design rule), and the caller accepts + flags the rare
    starvation in telemetry.
    """
    leader = remaining[0]
    if _est_gen_time(leader, rtf=rtf, sec_per_char=sec_per_char) <= buffer_remaining:
        return remaining
    pieces = _split_clauses(leader)
    if len(pieces) <= 1:
        return remaining
    return pieces + remaining[1:]


def _ema(old: float, new: float, alpha: float) -> float:
    return (1 - alpha) * old + alpha * new


# ── the online loop ───────────────────────────────────────────────────────

# generate(text) -> (audio_payload, audio_seconds)
GenerateFn = Callable[[str], "tuple[object, float]"]
ClockFn = Callable[[], float]


def stream_chunks(
    text: str,
    *,
    generate: GenerateFn,
    clock: ClockFn,
    cfg: ChunkConfig | None = None,
) -> Iterator[ChunkResult]:
    """Run the adaptive buffer-ratchet schedule, yielding one ChunkResult per chunk.

    ``generate`` does the actual synthesis (real GPU or a simulation) and
    returns its opaque audio payload plus the audio's duration in seconds.
    ``clock`` returns monotonically increasing seconds; ``gen_time`` is measured
    as the clock delta around ``generate``.

    Buffer model (all in audio-seconds):
      * Playback begins when chunk 1 arrives, so no drain occurs during chunk 1.
      * For chunk N≥2, the client plays ``gen_time`` seconds of buffered audio
        while we generate it, then we add this chunk's audio. If ``gen_time``
        exceeds the buffer the stream starved (``drained`` > 0).
    """
    cfg = cfg or ChunkConfig()

    units = protect_first_audio(split_sentences(text), cfg)
    rtf = cfg.rtf_prior
    sec_per_char = cfg.sec_per_char_prior
    buffer_remaining = 0.0
    remaining: list[str] = units
    index = 0

    while remaining:
        first = index == 0
        if not first:
            # If the next indivisible sentence would starve the buffer, expose
            # its clause boundaries so plan_chunk can pack to the buffer.
            remaining = relieve_leader(
                remaining, buffer_remaining, rtf=rtf, sec_per_char=sec_per_char
            )
        margin = cfg.margin_first if first else cfg.margin
        chunk_text, remaining = plan_chunk(
            remaining, buffer_remaining, margin=margin, rtf=rtf, sec_per_char=sec_per_char
        )
        est_gen = _est_gen_time(chunk_text, rtf=rtf, sec_per_char=sec_per_char)

        t0 = clock()
        audio, audio_sec = generate(chunk_text)
        gen_time = clock() - t0

        # Starvation: did the buffer run dry while we generated this chunk?
        drained = 0.0 if first else max(0.0, gen_time - buffer_remaining)

        # Live self-correction.
        if gen_time > 0:
            rtf = _ema(rtf, audio_sec / gen_time, cfg.ema_alpha)
        if chunk_text:
            sec_per_char = _ema(sec_per_char, audio_sec / len(chunk_text), cfg.ema_alpha)

        buffer_before = buffer_remaining
        if first:
            buffer_remaining = audio_sec
        else:
            buffer_remaining = max(0.0, buffer_remaining - gen_time) + audio_sec

        yield ChunkResult(
            index=index,
            text=chunk_text,
            audio=audio,
            audio_sec=audio_sec,
            gen_time=gen_time,
            est_gen=est_gen,
            buffer_before=buffer_before,
            buffer_after=buffer_remaining,
            drained=drained,
            rtf=rtf,
            sec_per_char=sec_per_char,
        )
        index += 1