"""Speaker attribution over pre-segmented spans.

The passage has already been split into narration and dialogue spans by
deterministic quote analysis (static/js/audiobook-spans.js). The LLM is asked
only WHO speaks each dialogue span and never re-emits any text, so it cannot
drop, duplicate or fuse passages, and its output is a few tokens per line.
"""
from __future__ import annotations

import json
import re
from typing import Any

MAX_SPANS = 400
MAX_SPAN_CHARS = 20_000
MAX_TOTAL_CHARS = 60_000
_ID_RE = re.compile(r"^[ND]\d{1,6}$")
_UNKNOWN = {"", "unknown", "unbekannt", "?", "n/a", "none", "null", "unklar"}
_NARRATOR = {"narrator", "erzähler", "erzaehler", "narration"}


class SpanRequestError(ValueError):
    """The request body is not a valid span-attribution request."""


def validate_spans(raw: Any) -> list[dict]:
    """Validate and normalise the span list at the API boundary."""
    if not isinstance(raw, list) or not raw:
        raise SpanRequestError("spans must be a non-empty list")
    if len(raw) > MAX_SPANS:
        raise SpanRequestError(f"too many spans (max {MAX_SPANS})")
    spans, seen, total = [], set(), 0
    for item in raw:
        if not isinstance(item, dict):
            raise SpanRequestError("each span must be an object")
        sid, kind, text = item.get("id"), item.get("kind"), item.get("text")
        if not isinstance(sid, str) or not _ID_RE.match(sid):
            raise SpanRequestError("span ids must look like N12 or D7")
        if sid in seen:
            raise SpanRequestError(f"duplicate span id {sid}")
        if kind not in ("narration", "dialogue") or (kind == "dialogue") != sid.startswith("D"):
            raise SpanRequestError(f"span {sid} has an inconsistent kind")
        if not isinstance(text, str) or not text.strip():
            raise SpanRequestError(f"span {sid} has no text")
        if len(text) > MAX_SPAN_CHARS:
            raise SpanRequestError(f"span {sid} is too long")
        total += len(text)
        if total > MAX_TOTAL_CHARS:
            raise SpanRequestError("passage is too long")
        seen.add(sid)
        spans.append({"id": sid, "kind": kind, "text": text.strip()})
    if not any(s["kind"] == "dialogue" for s in spans):
        raise SpanRequestError("passage has no dialogue spans")
    return spans


def validate_only_ids(raw: Any, spans: list[dict]) -> list[str]:
    """Ids to attribute when the rest of the window is context only."""
    if not isinstance(raw, list):
        return []
    dialogue = {s["id"] for s in spans if s["kind"] == "dialogue"}
    return [i for i in dict.fromkeys(raw) if isinstance(i, str) and i in dialogue]


_RULES = """\
You attribute speakers in prose fiction for a multi-voice audiobook.

The passage below is ALREADY divided into numbered lines. Lines starting with N are narration; lines starting with D are quoted speech. Do NOT rewrite, split, merge or quote any text. Your only job: for EVERY D line, decide who speaks it.

Answer with ONLY this JSON object (no markdown, no commentary):
{"speakers":[{"id":"D1","speaker":"Name","emotion":"one or two words"}]}
- Exactly one entry per D line, in order, with the exact ids given.
- speaker: the character's name. Reuse the exact spelling from the known-characters list whenever it is the same person.
- emotion: one or two words in the language of the text describing how the line is spoken, or "".

How to decide (apply in this order):
1. Speech tags in neighbouring N lines name the speaker ("sagte Marcian", "fragte sie", "asked Peter"). Resolve pronouns (er/sie/he/she) to the most recently named character of matching gender.
2. A narration line ending in ":" makes its subject the speaker of the next D line.
3. The narration right AFTER a D line often names the speaker, also in impersonal wording ("»Was machst du da?« ertönte es über ihm. Karyla war herübergekommen." -> Karyla).
4. ADDRESSEE vs SPEAKER: a person named next to a quote is very often the one being ADDRESSED, looked at or answered - not the one speaking. Check who actually performs the speech verb or the action framing the quote. "X wandte sich an Y" means X speaks and Y answers next.
5. Turn-taking: in a two-person exchange the speakers alternate strictly, even across many untagged D lines. Continue the alternation from the "Recent dialogue" lines. A short tag between two D lines ("»Die Pause«, sagte Peter, »ist vorbei.«") means the SAME speaker on both sides.
6. Keep the alternation consistent over a whole scene instead of deciding each line in isolation. Also watch for two characters' lines being swapped, especially the protagonist attracting lines that belong to the person talking WITH him or her.
7. Voice announcements ("Marcians Stimme klang fest", "X setzte zum Sprechen an", "ihre ersten Worte waren") give the next D line to that person. A line that names its own speaker as an introduction ("Man nennt mich Andra") belongs to that person.
8. Names: use the exact name the text or the known-characters list gives for that person. If in this scene the person is referred to ONLY by a role or epithet ("der Fremde", "die Baronin", "der junge Oberst", "der Inquisitor"), answer with that label in its plain form ("Der Fremde", "Baronin", "Oberst", "Inquisitor") - do NOT swap in the name of some other known character just because a proper name feels more specific. Use a known character's name for such an epithet only when the text itself ties them together (e.g. "der Inquisitor Marcian", "Marcian, der Fremde", or the same person addressed by name in this passage). Which epithet belongs to which named character is settled afterwards, from the whole book.
9. Answer "Unknown" only when the sole evidence is that a character stands nearby or is important in the scene; a confidently wrong name silently gives one character another's voice, an Unknown is cheap to fix. Turn-taking and clear tags always come first.
10. If a D line is clearly not spoken words at all (a quoted title or term), answer "Narrator".
"""


def build_messages(
    spans: list[dict],
    known_characters: list[str] | None = None,
    recent: str = "",
    following: str = "",
    language: str = "",
    want_reasoning: bool = False,
    only_ids: list[str] | None = None,
) -> tuple[str, str]:
    """Return (system, user) prompts for one window of spans."""
    system = _RULES
    if language:
        system += f"\nThe text language is {language}; keep names and emotion words in that language.\n"
    if want_reasoning:
        system += (
            "\nFirst, inside <think>...</think>, reason in at most 4 short sentences about who speaks the "
            "untagged or ambiguous D lines. Then output the JSON object and nothing else."
        )
    lines = []
    for sp in spans:
        if sp["kind"] == "dialogue":
            lines.append(f'[{sp["id"]}] "{sp["text"]}"')
        else:
            lines.append(f'[{sp["id"]}] {sp["text"]}')
    known = [str(n).strip() for n in (known_characters or []) if str(n).strip()][:60]
    user = ""
    if known:
        user += "Known characters so far: " + ", ".join(known) + "\n\n"
    if recent.strip():
        user += "Recent dialogue (already attributed - continue the same conversation and alternation):\n" + recent.strip() + "\n\n"
    user += "Passage:\n" + "\n".join(lines) + "\n\n"
    if following.strip():
        user += "What follows the passage (context only, do not attribute):\n" + following.strip() + "\n\n"
    ids = [s["id"] for s in spans if s["kind"] == "dialogue" and (not only_ids or s["id"] in only_ids)]
    user += f"Attribute these lines: {', '.join(ids)}\nReturn ONLY the JSON object."
    return system, user


def max_tokens_for(spans: list[dict], want_reasoning: bool = False) -> int:
    n = sum(1 for s in spans if s["kind"] == "dialogue")
    return min(4096, 300 + 60 * n) + (350 if want_reasoning else 0)


def _clean_speaker(value: Any) -> str:
    s = re.sub(r"\s+", " ", str(value or "")).strip().strip("\"'»«„“”‚‘’`").strip()
    s = re.sub(r"[.:;,]+$", "", s).strip()
    low = s.casefold()
    if low in _UNKNOWN:
        return "Unknown"
    if low in _NARRATOR:
        return "Narrator"
    return s[:80]


_OBJ_RE = re.compile(r"\{[^{}]*\}")


def parse_response(raw: str, spans: list[dict]) -> dict:
    """Extract {id: {speaker, emotion}} from the model's answer.

    Tolerates reasoning blocks, markdown fences and truncated JSON by reading
    each {"id": ...} object independently; unknown ids are ignored and missing
    dialogue ids are reported so the caller can retry or mark them Unknown.
    """
    content = re.sub(r"<think>.*?</think>", "", raw or "", flags=re.DOTALL)
    if "<think>" in content and "</think>" not in content:
        content = content.split("<think>", 1)[0]
    wanted = {s["id"] for s in spans if s["kind"] == "dialogue"}
    speakers: dict[str, dict] = {}
    for m in _OBJ_RE.finditer(content):
        try:
            obj = json.loads(m.group(0))
        except Exception:
            continue
        sid = obj.get("id") if isinstance(obj, dict) else None
        if sid in wanted and sid not in speakers:
            speakers[sid] = {
                "speaker": _clean_speaker(obj.get("speaker")),
                "emotion": re.sub(r"\s+", " ", str(obj.get("emotion") or "")).strip()[:60],
            }
    order = [s["id"] for s in spans if s["kind"] == "dialogue"]
    missing = [i for i in order if i not in speakers]
    characters: list[str] = []
    for i in order:
        sp = speakers.get(i, {}).get("speaker")
        if sp and sp not in ("Unknown", "Narrator") and sp not in characters:
            characters.append(sp)
    return {"speakers": speakers, "missing": missing, "characters": characters}