"""Word-level timing of a narrated article, measured by the forced-aligner sidecar (aligner/).

The MP3 is rendered sentence by sentence, so `audio/<id>.timings.json` already says where every sentence sits.
This module asks the sidecar where each *word* of a sentence starts and ends inside that span and records the
answer as `words: [[start_ms, end_ms], ...]` per segment (the file becomes version 2; version 1 readers simply
ignore the key). Never blocks playback: it runs in a background thread, retries with backoff, and when the
sidecar is down it gives up quietly -- the player then interpolates words inside each sentence as before.
"""
from __future__ import annotations

import io
import json
import logging
import os
import threading
import time
from pathlib import Path
from typing import Any, Callable, Sequence

import requests

from core import epub_export

log = logging.getLogger(__name__)

CHUNK_SECONDS = 45.0             # audio per sidecar request: bounded time per call and real progress
CHUNK_MAX_SPANS = 25
ATTEMPTS = 3
CONNECT_TIMEOUT = 3.0
READ_TIMEOUT = 150.0
FAILURE_COOLDOWN = 600.0         # after a failed run, do not retry on every timings request
WORDS_VERSION = 2

_jobs: dict[tuple[str, str], dict[str, Any]] = {}
_lock = threading.Lock()


def aligner_url() -> str:
    return os.environ.get("ALIGNER_URL", "").strip().rstrip("/")


# ── planning ───────────────────────────────────────────────────────────────────

def alignment_jobs(units: Sequence[dict[str, Any]], blocks: Sequence[dict[str, Any]],
                   interview: bool) -> list[dict[str, Any]]:
    """Group sentence units into alignment jobs: [{units: [idx...], start, end, text, lead}] (ms).

    A block whose sentences all have their own recorded utterance gets one job per sentence (tightest span).
    A block whose spans were only anchored (counts disagreed) is aligned as one span and split afterwards.
    `lead` counts leading tokens that were never spoken (an interview "Speaker:" label) and get no time.
    """
    from core.magazine import _INTERVIEW_LABEL_RE
    by_block: dict[int, list[int]] = {}
    for idx, u in enumerate(units):
        by_block.setdefault(u["block"], []).append(idx)
    jobs: list[dict[str, Any]] = []
    for block, idxs in by_block.items():
        lead = 0
        if interview:
            m = _INTERVIEW_LABEL_RE.match(units[idxs[0]]["text"])
            if m:
                lead = len(units[idxs[0]]["text"].split()) - len(m.group(2).split())
        exact = all(units[i].get("nth") is not None for i in idxs)
        groups = [[i] for i in idxs] if exact else [idxs]
        for g in groups:
            text = " ".join(units[i]["text"] for i in g)
            first_lead = lead if g[0] == idxs[0] else 0
            spoken = " ".join(text.split()[first_lead:])
            if not spoken:
                continue
            jobs.append({"units": g, "start": units[g[0]]["start"], "end": units[g[-1]]["end"],
                         "text": spoken, "lead": first_lead})
    return jobs


def _chunks(jobs: Sequence[dict[str, Any]]) -> list[list[dict[str, Any]]]:
    out: list[list[dict[str, Any]]] = []
    cur: list[dict[str, Any]] = []
    for j in jobs:
        if cur and (len(cur) >= CHUNK_MAX_SPANS or (j["end"] - cur[0]["start"]) / 1000 > CHUNK_SECONDS):
            out.append(cur)
            cur = []
        cur.append(j)
    if cur:
        out.append(cur)
    return out


# ── sidecar client ─────────────────────────────────────────────────────────────

class AlignerError(RuntimeError):
    """The sidecar could not be used (down, busy too long, rejected the request)."""


def _post_chunk(url: str, wav: bytes, spans: list[dict[str, Any]], language: str,
                sleep: Callable[[float], None] = time.sleep) -> list[list[dict[str, Any]]]:
    last: Exception | None = None
    for attempt in range(ATTEMPTS):
        try:
            r = requests.post(f"{url}/align/batch", files={"audio": ("chunk.wav", wav, "audio/wav")},
                              data={"spans": json.dumps(spans), "language": language},
                              timeout=(CONNECT_TIMEOUT, READ_TIMEOUT))
            if r.status_code == 200:
                return r.json()["spans"]
            if 400 <= r.status_code < 500 and r.status_code != 429:
                raise AlignerError(f"aligner rejected the request ({r.status_code})")
            last = AlignerError(f"aligner answered {r.status_code}")
        except requests.RequestException as exc:
            last = AlignerError(f"aligner unreachable ({type(exc).__name__})")
        except (ValueError, KeyError) as exc:
            last = AlignerError(f"aligner sent an unreadable answer ({type(exc).__name__})")
        if attempt < ATTEMPTS - 1:
            sleep(2 ** attempt)
    raise last or AlignerError("aligner failed")


def _wav_slice(audio: Any, start_ms: int, end_ms: int) -> bytes:
    buf = io.BytesIO()
    audio[start_ms:end_ms].set_frame_rate(16000).set_channels(1).export(buf, format="wav")
    return buf.getvalue()


# ── the job ────────────────────────────────────────────────────────────────────

def align_words(mp3: Path, units: Sequence[dict[str, Any]], jobs: Sequence[dict[str, Any]], language: str, *,
                url: str, progress: Callable[[int, int], None] | None = None,
                post: Callable[..., Any] = _post_chunk) -> dict[int, list[list[int]]]:
    """Ask the sidecar for the words of every job; returns {unit index: [[start_ms, end_ms], ...]} for the
    units it could time (a unit's list has one pair per whitespace token, label tokens first at zero length)."""
    from pydub import AudioSegment
    audio = AudioSegment.from_file(mp3)
    chunks = _chunks(jobs)
    result: dict[int, list[list[int]]] = {}
    done = 0
    for chunk in chunks:
        lo = max(0, int(chunk[0]["start"]) - 100)
        hi = min(len(audio), int(chunk[-1]["end"]) + 100)
        spans = [{"start": round((j["start"] - lo) / 1000, 3), "end": round(min(j["end"], hi) / 1000 - lo / 1000, 3),
                  "text": j["text"]} for j in chunk]
        answer = post(url, _wav_slice(audio, lo, hi), spans, language)
        for j, words in zip(chunk, answer):
            pairs = [[int(round(lo + w["start"] * 1000)), int(round(lo + w["end"] * 1000))] for w in words]
            pairs = [[j["start"], j["start"]]] * j["lead"] + pairs
            at = 0
            for i in j["units"]:
                n = len(units[i]["text"].split())
                result[i] = pairs[at:at + n]
                at += n
            done += 1
        if progress:
            progress(done, len(jobs))
    return result


def build_segment_words(units: Sequence[dict[str, Any]], per_unit: dict[int, list[list[int]]],
                        n_segments: int) -> dict[int, list[list[int]]]:
    """Per record segment, the words of all its sentences in order -- only for segments fully covered."""
    out: dict[int, list[list[int]]] = {}
    bad: set[int] = set()
    for idx, u in enumerate(units):
        k = u["seg"]
        words = per_unit.get(idx)
        if words is None or len(words) != len(u["text"].split()):
            bad.add(k)
            continue
        out.setdefault(k, []).extend(words)
    return {k: v for k, v in out.items() if k not in bad and 0 <= k < n_segments}


def write_words(mp3: Path, segment_words: dict[int, list[list[int]]]) -> bool:
    """Add `words` to the record (atomic replace; everything else in the file is kept as it was)."""
    path = mp3.with_suffix(".timings.json")
    try:
        rec = json.loads(path.read_text("utf-8"))
        segs = rec["segments"]
        for k, words in segment_words.items():
            segs[k]["words"] = words
        rec["version"] = WORDS_VERSION
        rec["words_engine"] = "mms_fa"
        tmp = path.with_suffix(".tmp")
        tmp.write_text(json.dumps(rec), "utf-8")
        tmp.replace(path)
        return True
    except (OSError, ValueError, KeyError, IndexError):
        log.warning("Could not write word timings for %s", mp3.name)
        return False


# ── background job registry ────────────────────────────────────────────────────

def status(doc_id: str, article_id: str) -> dict[str, Any]:
    with _lock:
        return dict(_jobs.get((doc_id, article_id)) or {"state": "idle", "done": 0, "total": 0})


def should_autostart(doc_id: str, article_id: str) -> bool:
    if not aligner_url():
        return False
    st = status(doc_id, article_id)
    if st["state"] == "running":
        return False
    # Whatever the last outcome, do not re-run on every request: a run that finished without covering every
    # sentence (or after the audio was re-rendered) is retried after the cooldown.
    return time.time() - st.get("finished", 0) >= FAILURE_COOLDOWN


def start(doc_id: str, article_id: str, run: Callable[[Callable[[int, int], None]], dict[str, Any]]) -> bool:
    """Run `run(progress)` in a daemon thread unless this article is already being aligned."""
    key = (doc_id, article_id)
    with _lock:
        if (_jobs.get(key) or {}).get("state") == "running":
            return False
        _jobs[key] = {"state": "running", "done": 0, "total": 0, "started": time.time()}

    def progress(done: int, total: int) -> None:
        with _lock:
            _jobs[key].update(done=done, total=total)

    def target() -> None:
        try:
            info = run(progress)
            end = {"state": "done", **info}
        except AlignerError as exc:
            log.warning("Word alignment for %s skipped: %s", article_id, exc)
            end = {"state": "failed", "error": str(exc)}
        except Exception as exc:                                  # never let a worker thread die noisily
            log.exception("Word alignment for %s failed", article_id)
            end = {"state": "failed", "error": type(exc).__name__}
        with _lock:
            _jobs[key].update(end, finished=time.time())

    threading.Thread(target=target, name=f"align-{article_id}", daemon=True).start()
    return True