UmbraAudio/umbraaudio/analysis/llm_attributor.py

from __future__ import annotations

import json
import logging
import re
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from openai import OpenAI

from umbraaudio.voices.lmstudio_gate import queued_chat_completion
from umbraaudio.analysis.attribution_qa import (
    SPEECH_VERBS as _QA_SPEECH_VERBS,
)
from umbraaudio.analysis.attribution_qa import (
    _detect_addressee as _qa_detect_addressee,
)
from umbraaudio.analysis.attribution_qa import (
    _match_speech_tag as _qa_match_speech_tag,
)
from umbraaudio.analysis.attribution_qa import (
    check_attribution as _qa_check_attribution,
)
from umbraaudio.analysis.attribution_qa import (
    fix_high_confidence as _qa_fix_high_confidence,
)
from umbraaudio.analysis.attribution_qa import (
    verify_conversation_flow as _qa_verify_flow,
)
from umbraaudio.analysis.clarity import evidence_kind_for
from umbraaudio.analysis.context_analyzer import ContextAnalyzer
from umbraaudio.analysis.tone_analyzer import ToneAnalyzer
from umbraaudio.config import LLMConfig

logger = logging.getLogger(__name__)

_PRONOUNS = frozenset(
    {
        "i", "me", "my", "you", "your", "he", "him", "his", "she", "her",
        "they", "them", "their", "we", "us", "our", "it", "its", "one",
    }
)

# Matches a TAG narration line whose subject is a pronoun: "He said,", "She asks."
_PRONOUN_TAG = re.compile(
    r"^\s*(H[ei]|She|They)\s+"
    r"(says|said|asks|asked|replies|replied|whispers|whispered|shouts|shouted|"
    r"mutters|muttered|calls|called|answers|answered)\b",
    re.IGNORECASE,
)

# Words allowed inside a Title Case ROLE name (rule 3g / 15).
_ROLE_WORDS = frozenset(
    {"the", "a", "an", "of", "in", "to", "for", "and", "or", "on", "at", "de"}
)


def _looks_like_role(name: str) -> bool:
    """True when a name plausibly is a Title Case ROLE name, not a personal name."""
    if not name or len(name) > 40:
        return False
    if re.search(r"[()\d]", name):
        return False
    words = [w for w in re.split(r"[\s-]+", name.strip()) if w]
    if not 1 <= len(words) <= 6:
        return False
    for w in words:
        if w.lower() in _ROLE_WORDS:
            continue
        if not w[:1].isupper():
            return False
    return True

# Enhanced prompt with context awareness and confidence scoring
PROMPT_TEMPLATE = """You are a dialogue attribution expert for audiobook production.

Given utterances that need speaker attribution, identify who is speaking each line.

## Character Profiles
{character_profiles}

## Chapter Title
{chapter_title}

## Scene Context
This chapter has {scene_count} scene(s) separated by scene breaks (<<<>>>).
You are attributing dialogue for: {scene_label}

{scene_speakers_context}

## Character Presence in This Scene
{character_presence_context}

## Unnamed Role Mappings (roles that get names later in this chapter)
{role_mappings_context}

## Recent Conversation Flow (context before this batch)
{position_context}

## Mechanical Narration Evidence
For each utterance below, these machine-extracted fields came straight from the
manuscript (the story's own words, not guesses):
- "narr_before": the narration text immediately before this dialogue line, verbatim
- "tag": the speaker named or tagged by that narration ("Kainda", "the lead guard",
  "She"), or null when the narration tags no speaker
- "tag_kind": "name" (exact character name), "role" (role/title), "pronoun"
  (He/She/They), or null
- "addressees": names or titles the dialogue itself addresses at its start, if any

These fields are the STRONGEST signal available. Prefer them over dialogue content.

## Utterances to Attribute
{utterances_json}

## Rules
1. Each utterance is already correctly split — do NOT merge or split them
2. EVERY utterance MUST be attributed to a speaker. Return EXACTLY one entry per utterance index. Do not skip any.
3. ALL these utterances are dialogue (in quotes) — they are spoken by a CHARACTER, almost never "Narrator". Only use "Narrator" if the line is truly un-attributable.
3b. NARRATION-OVERRIDES-CONTENT: The 1-2 narration utterances IMMEDIATELY BEFORE a dialogue line are the STRONGEST signal for speaker identity. Always check these FIRST before looking at dialogue content. When narration says a character IS speaking — using ANY pattern ("X says", "X nods and says", "X checks the screen and says", "X hits a mute button and says", "X is looking at Y and says", "X continues talking with Y", "X pauses and says", "X is telling Y") where X is the speaker's name or role — the following quoted lines are THAT character's dialogue, until another narration tag indicates a different speaker. This rule OVERRIDES any guess based on the dialogue content.
3c. A multi-line speech block from one character continues until dialogue tags, action, or the responder's lines indicate a change.
3d. A single quoted dialogue utterance may contain MULTIPLE sentences (e.g. one paragraph of speech) — every sentence within it belongs to the SAME speaker. Do not split the speaker across sentences within one utterance.
3e. Attribute to a computer/AI/communication-system character ONLY when narration explicitly identifies a computer/AI/communication system as the speaker (e.g. "the ship says", "the computer says", "'[name]' says", "her ear device says", "the comm crackles, '...'") AND such a character exists in the Character Profiles. A bare "a voice is heard" / "the voice calls again" / "someone says" over a radio or in the distance is a HUMAN speaker unless the text says otherwise — use gender cues in the narration ("it sounds female") to pick the human.
3f. When narration says "She says," or "He says," resolve the pronoun to the most recently mentioned character of the appropriate gender in the recent context above. Never resolve a pronoun to a character who is absent from the scene, and never resolve to a named character when the pronoun's antecedent is an unnamed role ("she asks the lead guard... he replies" → the guard speaks).
3g. UNNAMED ROLES: when a speech tag names a role or title WITHOUT a proper name ("the lead guard", "the doctor", "a pilot", "a pirate", "a Warbot", "a guard"), the speaker is that ROLE. If the Character Profiles list contains a matching role speaker (e.g. "Chief of the Guard", "The Doctor", "Warbots", "Pilot (Freighter)"), use its EXACT primary_name. Otherwise use a descriptive Title Case role name (e.g. "Lead Guard"). NEVER attribute such a line to an unrelated named character merely because that character is present in the scene. Guessing a wrong named character is worse than using a role name.
3h. REPEATED LINES: an identical quoted line that appears again in the same scene ("The voice calls again, 'Help me, please.'") is spoken by the SAME speaker as the first occurrence. When narration says "again", "another voice", or restates a previous line, reuse the earlier speaker.
3i. RETROACTIVE NAMING: If an unnamed role ("the robot", "a guard") is later given a name ("Niñera", "Support Robot 74532") in the same scene, ALL prior dialogue by that role should be attributed to that name.
3j. ADDRESS PATTERN: When dialogue addresses someone by name or title ("X, kill the connection" / "Yes, Captain" / "Captain, we need to jump"), the SPEAKER is NOT the person being addressed. The speaker is whoever was identified as speaking in the preceding narration. If narration says "X says" and the dialogue reads "Y, kill the connection", the speaker is X (talking TO Y). If narration says "the ship's AI says" and the dialogue reads "Captain, we need to jump", the speaker is the AI (talking TO the Captain).
3k. SELF-ADDRESS IMPOSSIBLE: A character never addresses themselves by name or title in their own dialogue. If a line says "Yes, Captain" and the Captain is X, then X is NOT the speaker — someone else is responding to her. If a line says "Y, do Z" and Y is X, then X is NOT the speaker.
4. Use dialogue tags (X says/replies/asks) to identify speakers
5. Use context: who is present, who was speaking before, who is addressed
6. Resolve pronouns: "she says" refers to the most recently mentioned female character
7. If a character has aliases, use their EXACT full primary_name shown in the profiles (e.g. the full name as listed, NOT a short form or nickname)
8. The "current_speaker" field shows who the mechanical segmenter guessed — use it as a hint but correct if wrong
9. In a conversation, speakers usually alternate. If one line is Character A, the next line is likely a different character (often the one addressed).
10. CRITICAL — only return the EXACT utterance indices provided in the input. Do NOT invent or hallucinate extra indices.
11. CRITICAL — every returned speaker must be the ACTUAL character speaking. Assigning the wrong speaker ruins the audiobook. When genuinely uncertain, prefer an unnamed role speaker over a random named character.
12. If the Character Profiles include an AI/computer character (e.g. a ship AI), attribute dialogue to it only when the text indicates an AI or computer is speaking (e.g., "ear device", "ship AI", "computer voice"). Do NOT default undetermined utterances to an AI character if another character is more likely.
13. KEYWORD NON-MATCHING: When a character's name appears in the TEXT of a dialogue line (e.g., "The ship is the Kestrel", "I know Captain Smith", "We are near the Kestrel"), it does NOT mean that character is speaking. The speaker is determined by the preceding narration tag, NOT by character names mentioned inside the dialogue. A character can be TALKED ABOUT without being the speaker. Only attribute to a character when a narration tag or dialogue tag explicitly identifies them as the speaker.
14. EVIDENCE PRIORITY (order of preference):
    (a) a "name" tag in narr_before → that character speaks
    (b) a "role" tag → the role speaker (rule 3g)
    (c) a "pronoun" tag → resolve by rule 3f
    (d) an addressee listed in "addressees" → the addressee is NOT the speaker (rules 3j/3k)
    (e) scene flow and dialogue content as a last resort
    Never override an explicit tag with a guess from dialogue content (rule 13).
15. ROLE-NAME CONSISTENCY: once you pick a speaker for a recurring role ("the guard",
    "a pirate"), keep using the SAME speaker name for that role throughout the chapter.
    Do NOT emit scene-qualified variants like "Guard (Scene 2)" or invent a new role
    name per line. Reuse the existing role speaker from the profiles.

## Scene Attribution Rules (CRITICAL)
- Chapters can contain multiple scenes separated by scene breaks (<<<>>>)
- Each scene is an independent unit. Speakers identified in one scene may NOT appear in other scenes
- Within a scene, once a speaker is identified, subsequent dialogue by the same person continues until another speaker is identified
- "captain" in one scene may refer to a different person than "captain" in another scene — there can be multiple characters with the same title
- Only attribute a speaker to dialogue if there is evidence within THAT SCENE that they are speaking
- If a scene introduces new characters via narration, those characters are candidates for dialogue in that scene
- A speaker who appears in scene 1 may be entirely absent from scene 2 — do NOT carry speakers across scene boundaries
- When ambiguous ("he says", "she says"), resolve using the most recent speaker mention within the SAME SCENE

## Output Format
Return ONLY a JSON array of objects, one per utterance, each with:
- "index": the utterance index (number)
- "speaker": the correct speaker name (use primary_name from profiles, MUST be a character name, NOT "Narrator")
- "confidence": your confidence in this attribution (0.0 - 1.0)
- "reasoning": brief explanation of key evidence (max 50 chars)
- "uncertain": true ONLY when no narration evidence exists and you had to guess
  from content; this flags the line for a later verification pass

Return ONLY the JSON array, no other text."""


# Second-pass prompt: verify and correct attributions a first pass already made.
# The LLM is shown only the lines a mechanical QA pass flagged (or low-confidence
# lines), each with its mechanical narration evidence, and asked to confirm or fix.
VERIFY_TEMPLATE = """You are a dialogue attribution verifier for audiobook production.

A first pass already attributed speakers to the lines below, but a mechanical check
flagged some as possibly wrong. For EACH line decide whether the current attribution
is correct; if not, give the correct speaker.

## Character Profiles
{character_profiles}

## Chapter Title
{chapter_title}

## Lines to Verify
{items_json}

## Rules
1. The narration IMMEDIATELY BEFORE a line (the "narr_before" field,
   VERBATIM from the manuscript) is the STRONGEST signal. A "name" tag
   there decides the speaker.
2. A "role" tag ("the lead guard", "a pirate") → use the matching role
   speaker's EXACT primary_name from the Character Profiles, else a
   descriptive Title Case role name. NEVER switch to an unrelated named
   character.
3. A "pronoun" tag ("He says", "She says") → resolve to the most recently
   mentioned character of that gender in the recent context, within the
   SAME scene.
4. The addressee of a line is NOT its speaker ("Yes, Captain" is spoken TO
   the Captain). A character never addresses themself by name.
5. Repeated lines ("The voice calls again, ...") keep the first
   occurrence's speaker.
6. NEVER use "Narrator" for dialogue when any character evidence exists.
7. Use EXACT primary_names from the profiles, never invent names. Keep
   role names consistent (no "Guard (Scene 2)" variants).
8. The "issue" field explains why the line was flagged — treat it as a
   hint, not a verdict. Confirm the current speaker when it is right.
9. Keep the EXACT "index" values from the input. Return one object per
   line.

## Output Format
Return ONLY a JSON array of objects, one per input line, in the same order, each with:
- "index": the line's exact "index" value (number)
- "speaker": the CORRECT speaker primary_name (the same as "current" if it is already right)
- "confidence": your confidence in this attribution (0.0 - 1.0)
- "reasoning": brief explanation (max 60 chars)

Return ONLY the JSON array, no other text."""


@dataclass
class AttributionResult:
    """Result of attribution with confidence scoring."""
    index: int
    speaker: str
    confidence: float
    reasoning: str


class LLMAttributor:
    """Uses LLM to attribute speakers to mechanically segmented utterances with multi-pass context awareness."""

    def __init__(self, config: LLMConfig, model: str | None = None, data_dir: Path | None = None):
        self.config = config
        self.model = model or config.lmstudio_model
        self.context_analyzer = ContextAnalyzer()
        self.tone_analyzer = ToneAnalyzer()
        self.data_dir = data_dir or Path(".umbra")

    def _get_client(self) -> OpenAI:
        return OpenAI(
            base_url=self.config.lmstudio_base_url,
            api_key="lm-studio",
            timeout=3600.0,
        )

    def _load_speakers(self) -> list[dict[str, Any]]:
        speakers_path = self.data_dir / "speakers.json"
        if not speakers_path.exists():
            return []
        with open(speakers_path) as f:
            data = json.load(f)
        speakers = data.get("speakers", [])
        return [dict(s) for s in speakers] if isinstance(speakers, list) else []

    def _load_character_profiles(self) -> str:
        speakers = self._load_speakers()
        if not speakers:
            return "No character profiles available."

        lines = []
        for s in speakers:
            name = s.get("primary_name", "Unknown")
            demos = s.get("demographics", {})
            dtype = demos.get("type", "Human")
            gender = demos.get("gender", "Unknown")
            aliases = ", ".join(s.get("aliases", [])) or "none"
            lines.append(f"- {name} ({dtype}, {gender}): aliases=[{aliases}]")

        return "\n".join(lines)

    def _build_alias_map(self) -> dict[str, str]:
        speakers = self._load_speakers()
        alias_map: dict[str, str] = {}
        for s in speakers:
            primary = s.get("primary_name", "")
            if not primary:
                continue
            alias_map[primary.lower()] = primary
        for s in speakers:
            primary = s.get("primary_name", "")
            for alias in s.get("aliases", []):
                if alias:
                    alias_map[alias.lower()] = primary
        for s in speakers:
            primary = s.get("primary_name", "")
            if not primary:
                continue
            first_word = primary.split()[0].lower()
            if first_word not in alias_map:
                alias_map[first_word] = primary
        return alias_map

    def _normalize_speaker(self, name: str, alias_map: dict[str, str]) -> str:
        if not name:
            return name
        raw = name.strip().strip("\"'“”")
        lower = raw.lower()
        if lower in ("narrator", "the narrator"):
            return "narrator"
        if lower in _PRONOUNS:
            return "narrator"
        if lower in alias_map:
            return alias_map[lower]
        # Strip a trailing parenthetical qualifier ("Xylia (Voice (Dragon))" → "Xylia")
        # and re-check the base name. "Guard (Scene 2)" stays as-is when it IS a
        # primary, but LLM reinventions like "Guard (Scene 3)" collapse to "Guard".
        match = re.match(r"^(?P<base>.+?)\s*\([^()]*\)\s*$", raw)
        if match:
            base = match.group("base").strip()
            if base and base.lower() in alias_map:
                return alias_map[base.lower()]
        if base := re.match(r"^(?P<base>[^(]+)\s*\(.*$", raw):
            base_name = base.group("base").strip()
            if base_name.lower() in alias_map:
                return alias_map[base_name.lower()]
        return raw

    def attribute_speakers(
        self,
        utterances: list[dict[str, Any]],
        chapter_title: str,
        chapter_id: int,
        chapter_data: dict[str, Any] | None = None,
        verify: bool = True,
        verify_rounds: int = 2,
    ) -> list[dict[str, Any]]:
        """
        Multi-pass speaker attribution with context analysis.
        
        Args:
            utterances: List of utterance dicts with type, text, utterance_index
            chapter_title: Title of the chapter
            chapter_id: Chapter ID
            chapter_data: Optional full chapter data with source_text for context analysis
            verify: Run the deterministic-fix + LLM verification pass afterwards
            verify_rounds: Max verification rounds (stop early on no change)
        """
        profiles = self._load_character_profiles()
        alias_map = self._build_alias_map()

        # Prepare chapter data for context analysis
        if chapter_data is None:
            chapter_data = {
                "id": chapter_id,
                "title": chapter_title,
                "source_text": "",
                "utterances": utterances,
            }

        # Run multi-pass context analysis
        logger.info(f"Chapter {chapter_id}: Running multi-pass context analysis...")
        context_analysis = self.context_analyzer.analyze_manuscript([chapter_data])

        # Run tone/emotion analysis
        scene_contexts = context_analysis.get("scene_contexts", {}).get(chapter_id, [])
        tone_profiles = self.tone_analyzer.analyze_chapter_emotions(chapter_data, scene_contexts)

        # Store context for use in prompt building
        self._current_context = context_analysis
        self._current_tone = tone_profiles

        dialogue_utterances = []
        for u in utterances:
            if u.get("type") != "dialogue":
                continue
            dialogue_utterances.append(u)

        if not dialogue_utterances:
            logger.info("Chapter %d: no dialogue utterances", chapter_id)
            return utterances

        dialogue_total = len(dialogue_utterances)
        logger.info(
            "Chapter %d: %d dialogue utterances to attribute",
            chapter_id,
            dialogue_total,
        )

        scenes = self._split_into_scenes(utterances)
        logger.info("Chapter %d: split into %d scene(s)", chapter_id, len(scenes))

        BATCH_SIZE = 40
        speaker_map: dict[int, str] = {}
        confidence_map: dict[int, float] = {}
        reasoning_map: dict[int, str] = {}
        uncertain_map: dict[int, bool] = {}
        evidence_by_idx: dict[int, dict[str, Any]] = {}

        SPEECH_VERB_SUFFIXES = (
            " says", " says,", " said", " said,", " asks", " asks,",
            " replies", " replies,", " whispers", " whispers,", " shouts",
            " shouts,", " mutters", " mutters,", " calls", " calls,",
            " exclaimed", " exclaimed,", " answered", " answered,",
            " laughed", " laughed,", " sighed", " sighed,", " snaps",
            " snaps,",
        )

        for scene_idx, scene_utterances in enumerate(scenes):
            scene_dialogue = [u for u in scene_utterances if u.get("type") == "dialogue"]

            if not scene_dialogue:
                continue

            scene_narration_names = self._extract_scene_speakers(scene_utterances)
            scene_speakers_context = self._build_scene_speakers_context(
                scene_idx,
                len(scenes),
                scene_narration_names,
                profiles,
                narration_text=self._extract_scene_narration(scene_utterances),
            )

            # Build character presence context for this scene
            character_presence_context = self._build_character_presence_context(chapter_id, scene_idx)
            role_mappings_context = self._build_role_mappings_context(chapter_id, scene_idx)

            batches = []
            for i in range(0, len(scene_dialogue), BATCH_SIZE):
                batches.append(scene_dialogue[i : i + BATCH_SIZE])

            logger.info(
                "Chapter %d, Scene %d/%d: %d dialogue in %d batches (narration names: %s)",
                chapter_id,
                scene_idx + 1,
                len(scenes),
                len(scene_dialogue),
                len(batches),
                ", ".join(scene_narration_names) if scene_narration_names else "none",
            )

            for batch_idx, batch in enumerate(batches):
                batch_data = []
                evidence_map = self._build_mechanical_evidence(
                    scene_utterances, batch, alias_map
                )
                for ev_idx, ev in evidence_map.items():
                    evidence_by_idx[ev_idx] = ev
                for u in batch:
                    evidence = evidence_map.get(u["utterance_index"])
                    if evidence is None:
                        evidence = {
                            "narr_before": "",
                            "tag": None,
                            "tag_kind": None,
                            "addressees": [],
                        }
                    batch_data.append(
                        {
                            "index": u["utterance_index"],
                            "text": u["text"][:200],
                            "current_speaker": u.get("speaker_name", "Narrator"),
                            **evidence,
                        }
                    )

                attributions = None
                context_caps: list[int | None] = [None, 1000, 500, 300, 0]
                for attempt, cap in enumerate(context_caps, start=1):
                    try:
                        position_context = self._build_position_context(
                            scene_utterances, batch, max_chars=cap
                        )
                        prompt = PROMPT_TEMPLATE.format(
                            character_profiles=profiles,
                            chapter_title=chapter_title,
                            scene_count=len(scenes),
                            scene_label=f"Scene {scene_idx + 1} of {len(scenes)}",
                            scene_speakers_context=scene_speakers_context,
                            character_presence_context=character_presence_context,
                            role_mappings_context=role_mappings_context,
                            position_context=position_context,
                            utterances_json=json.dumps(batch_data, indent=2),
                        )
                        result = self._call_llm(prompt)
                        if result:
                            attributions = result
                            break
                        logger.warning(
                            "Batch %d (scene %d, chapter %d) returned no attributions; "
                            "retrying with reduced context (attempt %d/%d)",
                            batch_idx,
                            scene_idx + 1,
                            chapter_id,
                            attempt,
                            len(context_caps),
                        )
                    except Exception:
                        logger.exception(
                            "Unexpected error attributing batch %d (scene %d, chapter %d)",
                            batch_idx,
                            scene_idx + 1,
                            chapter_id,
                        )
                        break

                if attributions:
                    for a in attributions:
                        idx = a.get("index")
                        speaker = a.get("speaker", "narrator")
                        confidence = a.get("confidence", 0.5)
                        reasoning = a.get("reasoning", "")
                        if speaker and speaker != "narrator" and speaker != "scene_break":
                            for suffix in SPEECH_VERB_SUFFIXES:
                                if speaker.endswith(suffix):
                                    speaker = speaker[: -len(suffix)].strip()
                                    break
                        speaker = self._normalize_speaker(speaker, alias_map)
                        if idx is not None:
                            speaker_map[idx] = speaker
                            confidence_map[idx] = confidence
                            reasoning_map[idx] = reasoning
                            uncertain_map[idx] = bool(a.get("uncertain", False))

                if batch_idx < len(batches) - 1:
                    time.sleep(1)

        # Apply attributions and confidence scores
        updated = []
        for u in utterances:
            idx = u["utterance_index"]
            if idx in speaker_map and u.get("type") == "dialogue":
                new_u = dict(u)
                new_u["speaker_name"] = speaker_map[idx]
                new_u["attribution_confidence"] = confidence_map.get(idx, 0.5)
                new_u["attribution_reasoning"] = reasoning_map.get(idx, "")
                ev = evidence_by_idx.get(idx)
                if ev:
                    new_u["attribution_evidence_kind"] = evidence_kind_for(ev)
                    new_u["attribution_tag_name"] = ev.get("tag")
                # Add emotion from tone analysis
                scene_id = self._find_scene_for_utterance(chapter_id, idx, scene_contexts)
                if scene_id is not None:
                    tone_profile = tone_profiles.get(scene_id)
                    if tone_profile:
                        emotion, emo_conf = self.tone_analyzer.get_utterance_emotion(
                            chapter_id, scene_id, idx, speaker_map[idx]
                        )
                        new_u["emotion"] = emotion
                        new_u["emotion_confidence"] = emo_conf
                updated.append(new_u)
            else:
                updated.append(u)

        updated = self._apply_speaker_fallback(updated)

        # Store the "uncertain" flag from pass 1 so the verification pass knows
        # which lines had no mechanical evidence to lean on.
        for u in updated:
            idx = u.get("utterance_index")
            if idx is not None and uncertain_map.get(idx):
                u["attribution_uncertain"] = True

        # Pass 3: mechanical auto-fix + LLM verification of flagged lines
        if verify:
            updated = self._verify_attributions(
                updated,
                chapter_id,
                chapter_title,
                profiles,
                alias_map,
                max_rounds=verify_rounds,
            )

        # Store confidence scores in context analysis for exception review
        context_analysis["attribution_confidences"] = [
            {
                "utterance_index": idx,
                "speaker": speaker_map[idx],
                "confidence": confidence_map.get(idx, 0.5),
                "reasoning": reasoning_map.get(idx, ""),
            }
            for idx in speaker_map
        ]

        logger.info(
            "Chapter %d: attributed %d/%d dialogue utterances",
            chapter_id,
            len(speaker_map),
            dialogue_total,
        )

        return updated

    def _build_mechanical_evidence(
        self,
        scene_utterances: list[dict[str, Any]],
        batch: list[dict[str, Any]],
        alias_map: dict[str, str],
    ) -> dict[int, dict[str, Any]]:
        """Extract machine-readable narration evidence for each dialogue line.

        For every dialogue utterance in ``batch``, walk back through the scene
        looking at narration BEFORE the line (up to 2 narration lines, staying
        inside the same scene) and pull out:
        - ``narr_before``: the nearest narration text before the line
        - ``tag``: speaker named by a speech tag ("Kainda", "the lead guard", "She")
        - ``tag_kind``: "name", "role", or "pronoun"
        - ``addressees``: whom the dialogue itself addresses at its start
        """
        evidence: dict[int, dict[str, Any]] = {}
        dialogue_positions: dict[int, int] = {}
        for pos, u in enumerate(scene_utterances):
            if u.get("type") == "dialogue":
                dialogue_positions[u["utterance_index"]] = pos

        for u in batch:
            idx = u.get("utterance_index")
            if idx is None:
                idx = u.get("index")
            pos = dialogue_positions.get(idx)
            if pos is None:
                continue
            entry: dict[str, Any] = {
                "narr_before": "",
                "tag": None,
                "tag_kind": None,
                "addressees": [],
            }

            # Walk back through narration lines (max 2), staying in the scene.
            narr_texts: list[str] = []
            j = pos - 1
            while j >= 0:
                prev = scene_utterances[j]
                prev_text = (prev.get("text") or "").strip()
                if prev_text == "<<<>>>" or (prev.get("type") or "").lower() != "narration":
                    break
                narr_texts.append(prev_text)
                if len(narr_texts) >= 2:
                    break
                j -= 1

            if narr_texts:
                entry["narr_before"] = narr_texts[0][:220]

            # Find the first narration line with a speech tag (nearest wins).
            for narr_text in narr_texts:
                if _QA_SPEECH_VERBS.search(narr_text):
                    kind, value, _resolved = _qa_match_speech_tag(narr_text, alias_map)
                    if kind == "name" and value:
                        entry["tag"] = value
                        entry["tag_kind"] = "name"
                        break
                    if kind == "role" and value:
                        entry["tag"] = value
                        entry["tag_kind"] = "role"
                        break
                pm = _PRONOUN_TAG.match(narr_text)
                if pm:
                    entry["tag"] = pm.group(1).lower()
                    entry["tag_kind"] = "pronoun"
                    break

            addressee = _qa_detect_addressee(u.get("text") or "", alias_map)
            if addressee:
                entry["addressees"].append(addressee)

            evidence[idx] = entry

        return evidence

    def _collect_verify_targets(
        self,
        utterances: list[dict[str, Any]],
        chapter_id: int,
        alias_map: dict[str, str],
    ) -> dict[int, str]:
        """Collect utterance indices the verification pass should re-check.

        Returns a map of ``utterance_index -> reason`` for every dialogue line
        that (a) a mechanical QA pass flagged, (b) is still attributed to the
        narrator, or (c) has low confidence or the ``uncertain`` flag from
        pass 1.
        """
        qa = _qa_check_attribution(utterances, chapter_id, alias_map)
        flow = _qa_verify_flow(utterances, chapter_id, alias_map)

        targets: dict[int, str] = {}
        for issue in qa.issues + flow.issues:
            reason = issue.evidence or issue.issue_type
            targets.setdefault(issue.utterance_index, f"{issue.issue_type}: {reason}")

        for u in utterances:
            if u.get("type") != "dialogue":
                continue
            idx = u.get("utterance_index")
            if idx is None or idx in targets:
                continue
            speaker = (u.get("speaker_name") or "narrator").lower()
            confidence = float(u.get("attribution_confidence") or 0.5)
            if speaker == "narrator":
                targets[idx] = "dialogue still attributed to narrator"
            elif confidence < 0.6:
                targets[idx] = f"low confidence ({confidence:.2f})"
            elif u.get("attribution_uncertain"):
                targets[idx] = "uncertain: no narration evidence in first pass"

        return targets

    def _verify_attributions(
        self,
        utterances: list[dict[str, Any]],
        chapter_id: int,
        chapter_title: str,
        profiles: str,
        alias_map: dict[str, str],
        max_rounds: int = 2,
    ) -> list[dict[str, Any]]:
        """Deterministic auto-fix + LLM verification of flagged attributions.

        Round loop:
        1. Apply ``fix_high_confidence`` (exact resolved name-tag mismatches).
        2. Collect targets via mechanical QA (+ narrator-left, low confidence).
        3. Ask the LLM to confirm or correct each target line (VERIFY_TEMPLATE).
        4. Apply only corrections that normalize to a known speaker.
        Stops when a round makes no changes or no targets remain.
        """
        if max_rounds <= 0:
            return utterances

        current = [dict(u) for u in utterances]

        for round_no in range(1, max_rounds + 1):
            current = _qa_fix_high_confidence(current, alias_map)
            targets = self._collect_verify_targets(current, chapter_id, alias_map)
            if not targets:
                break

            corrections = self._verify_round_llm(
                current,
                targets,
                chapter_id,
                chapter_title,
                profiles,
                alias_map,
            )
            if not corrections:
                break

            changed = 0
            by_idx = {int(c["index"]): c for c in corrections if c.get("index") is not None}
            for u in current:
                if u.get("type") != "dialogue":
                    continue
                idx = u.get("utterance_index")
                if idx is None or idx not in by_idx:
                    continue
                corr = by_idx[idx]
                new_speaker = str(corr.get("speaker") or "").strip()
                if not new_speaker:
                    continue
                new_speaker = self._normalize_speaker(new_speaker, alias_map)
                if new_speaker.lower() in ("narrator", "the narrator"):
                    continue
                # Gate: only accept known speakers (profiles/aliases) or
                # plausible Title Case role names — rejects invented names
                # and scene-qualified variants the normalizer kept verbatim.
                if new_speaker.lower() not in alias_map and not _looks_like_role(
                    new_speaker
                ):
                    logger.debug(
                        "Chapter %d: verify rejected invented name %r for utt %s",
                        chapter_id,
                        new_speaker,
                        idx,
                    )
                    continue
                old_speaker = u.get("speaker_name") or "narrator"
                if new_speaker.lower() == old_speaker.lower():
                    continue
                u["speaker_name"] = new_speaker
                u["attribution_confidence"] = float(corr.get("confidence") or 0.7)
                u["attribution_reasoning"] = (
                    "verified: " + str(corr.get("reasoning") or "")[:100]
                )
                changed += 1

            if changed == 0:
                break
            logger.info(
                "Chapter %d: verification round %d corrected %d target(s)",
                chapter_id,
                round_no,
                changed,
            )

        return current

    def _verify_round_llm(
        self,
        utterances: list[dict[str, Any]],
        targets: dict[int, str],
        chapter_id: int,
        chapter_title: str,
        profiles: str,
        alias_map: dict[str, str],
    ) -> list[dict[str, Any]] | None:
        """One LLM call over the target lines using VERIFY_TEMPLATE."""
        items: list[dict[str, Any]] = []
        target_positions: dict[int, int] = {}
        for pos, u in enumerate(utterances):
            if u.get("utterance_index") in targets:
                target_positions[u["utterance_index"]] = pos

        for idx in sorted(target_positions):
            pos = target_positions[idx]
            u = utterances[pos]
            items.append(
                {
                    "index": idx,
                    "current": u.get("speaker_name") or "narrator",
                    "text": (u.get("text") or "")[:200],
                    "issue": targets[idx],
                    "narr_before": "",
                    "tag": None,
                    "tag_kind": None,
                    "addressees": [],
                }
            )

        if not items:
            return None

        # Add narration evidence (same extraction as pass 1) for context.
        evidence_map = self._build_mechanical_evidence(utterances, items, alias_map)
        for item in items:
            item.update(evidence_map.get(item["index"], {}))

        results: list[dict[str, Any]] = []
        batch_size = 20
        for start in range(0, len(items), batch_size):
            chunk = items[start : start + batch_size]
            prompt = VERIFY_TEMPLATE.format(
                character_profiles=profiles,
                chapter_title=chapter_title,
                items_json=json.dumps(chunk, indent=2),
            )
            batch_result = self._call_llm(prompt)
            if batch_result:
                results.extend(batch_result)
            if start + batch_size < len(items):
                time.sleep(1)
        return results or None

    def _build_character_presence_context(self, chapter_id: int, scene_idx: int) -> str:
        """Build context about which characters are present in this scene."""
        scenes = self._current_context.get("scene_contexts", {}).get(chapter_id, [])
        if scene_idx >= len(scenes):
            return "No scene context available."

        scene = scenes[scene_idx]
        if not scene.characters_present:
            return "No specific characters identified in this scene's narration."

        lines = ["Characters present in this scene (from narration):"]
        for name, intro in scene.characters_present.items():
            role_info = f" (role: {intro.role})" if intro.role else ""
            aliases_info = f" aliases: {intro.aliases}" if intro.aliases else ""
            lines.append(f"- {name}{role_info}{aliases_info}")

        return "\n".join(lines)

    def _build_role_mappings_context(self, chapter_id: int, scene_idx: int) -> str:
        """Build context about unnamed roles that get names later."""
        scenes = self._current_context.get("scene_contexts", {}).get(chapter_id, [])
        if scene_idx >= len(scenes):
            return "No role mappings for this scene."

        scene = scenes[scene_idx]
        if not scene.unnamed_roles:
            return "No unnamed roles resolved in this scene."

        lines = ["Unnamed roles that get names later in this scene:"]
        for role, name in scene.unnamed_roles.items():
            lines.append(f"- '{role}' → {name}")

        return "\n".join(lines)

    def _find_scene_for_utterance(self, chapter_id: int, utterance_index: int, scenes: list) -> int | None:
        """Find which scene an utterance belongs to."""
        for scene in scenes:
            if scene.start_utterance <= utterance_index <= scene.end_utterance:
                return scene.scene_id
        return None

    def _apply_speaker_fallback(self, utterances: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """Attribute remaining dialogue via speech-tag detection or speaker carry."""
        result = list(utterances)
        active_speaker: str | None = None
        speech_tag_re = re.compile(
            r"\b([A-Z][a-zA-Z]+(?:\s+[A-Z][a-zA-Z]+)?)\s+(says|said|asks|asked|replies|replied|"
            r"whispers|whispered|shouts|shouted|mutters|muttered|calls|called|answers|answered|"
            r"tells|told|explains|explained)\b",
        )
        for i, u in enumerate(result):
            text = (u.get("text") or "").strip()
            if text == "<<<>>>":
                active_speaker = None
                continue

            if u.get("type") != "dialogue":
                continue

            current = u.get("speaker_name") or "narrator"
            if current.lower() != "narrator":
                active_speaker = current
                continue

            resolved: str | None = None

            for j in range(max(0, i - 3), i):
                prev_text = (result[j].get("text") or "").strip()
                if prev_text == "<<<>>>":
                    break
                match = speech_tag_re.search(prev_text)
                if match:
                    candidate = match.group(1)
                    resolved = self._normalize_speaker(candidate, self._build_alias_map())
                    break

            if resolved is None and active_speaker:
                resolved = active_speaker

            if resolved:
                new_u = dict(u)
                new_u["speaker_name"] = resolved
                result[i] = new_u
                active_speaker = resolved

        return result

    def _split_into_scenes(self, utterances: list[dict[str, Any]]) -> list[list[dict[str, Any]]]:
        scenes: list[list[dict[str, Any]]] = []
        current_scene: list[dict[str, Any]] = []
        for u in utterances:
            text = u.get("text", "")
            if "<<<>>>" in text or "<<<>>>" in text.strip():
                if current_scene:
                    scenes.append(current_scene)
                current_scene = []
            current_scene.append(u)
        if current_scene:
            scenes.append(current_scene)
        return scenes if scenes else [utterances]

    def _extract_scene_speakers(self, scene_utterances: list[dict[str, Any]]) -> list[str]:
        names_found: list[str] = []
        seen: set[str] = set()
        known_names = self._known_name_phrases()
        for u in scene_utterances:
            if u.get("type") == "narration":
                text = u.get("text", "")
                lower = text.lower()
                for name, phrase in known_names:
                    if name in seen:
                        continue
                    if phrase in lower:
                        seen.add(name)
                        names_found.append(name)
                for match in re.finditer(r"\b([A-Z][a-z]+(?:\s+[A-Z][a-z]+)+)\b", text):
                    name = match.group(1)
                    if name not in seen:
                        seen.add(name)
                        names_found.append(name)
        return names_found

    def _known_name_phrases(self) -> list[tuple[str, str]]:
        phrases: list[tuple[str, str]] = []
        for s in self._load_speakers():
            primary = s.get("primary_name", "")
            if not primary:
                continue
            phrases.append((primary, primary.lower()))
            for alias in s.get("aliases", []):
                if alias:
                    phrases.append((primary, alias.lower()))
        return phrases

    def _extract_scene_narration(self, scene_utterances: list[dict[str, Any]]) -> str:
        parts = []
        for u in scene_utterances:
            if u.get("type") == "narration":
                text = u.get("text", "").strip()
                if text and text not in ("<<<>>>",):
                    parts.append(text)
        return "\n".join(parts[-20:])

    def _build_scene_speakers_context(
        self,
        scene_idx: int,
        total_scenes: int,
        narration_names: list[str],
        profiles: str,
        narration_text: str = "",
    ) -> str:
        lines = [f"This is Scene {scene_idx + 1} of {total_scenes} in the chapter."]
        if narration_names:
            lines.append(
                "Character names mentioned in this scene's narration: " + ", ".join(narration_names)
            )
            lines.append(
                "Only attribute dialogue to characters who appear in this scene's narration. "
                "A character present in a previous scene may be absent from this one."
            )
        else:
            lines.append(
                "No character names were found in this scene's narration. "
                "Use dialogue tags, context, and pronouns to identify speakers."
            )
        if narration_text:
            lines.append("## This Scene's Narration (context for attribution)")
            lines.append(narration_text)
            lines.append("Use this narration to determine WHO is speaking in the dialogue below.")
        lines.append(
            "Remember: 'captain' in this scene may differ from 'captain' in another scene. "
            "There can be multiple characters sharing the same title."
        )
        return "\n".join(lines)

    def _find_batch_position(
        self, scene_utterances: list[dict[str, Any]], batch: list[dict[str, Any]]
    ) -> int:
        target_idx = batch[0]["utterance_index"]
        for i, u in enumerate(scene_utterances):
            if u.get("utterance_index") == target_idx:
                return i
        return 0

    def _build_position_context(
        self,
        scene_utterances: list[dict[str, Any]],
        batch: list[dict[str, Any]],
        max_chars: int | None = None,
    ) -> str:
        first_idx = batch[0]["utterance_index"]
        last_idx = batch[-1]["utterance_index"]

        start = 0
        end = len(scene_utterances)
        for i, u in enumerate(scene_utterances):
            if u.get("utterance_index") == first_idx:
                start = max(0, i - 40)
            if u.get("utterance_index") == last_idx:
                end = i + 1
                break

        parts: list[str] = []
        for i in range(start, end):
            u = scene_utterances[i]
            text = (u.get("text") or "").strip()
            if text == "<<<>>>" or not text:
                continue
            if u.get("type") == "narration":
                parts.append(text[:160])
            else:
                speaker = u.get("speaker_name") or "Narrator"
                parts.append(f"[{speaker}] {text[:150]}")

        if len(parts) > 45:
            parts = parts[-45:]

        char_total = 0
        trimmed: list[str] = []
        char_cap = 4200 if max_chars is None else max_chars
        for part in reversed(parts):
            if char_total + len(part) > char_cap:
                break
            trimmed.append(part)
            char_total += len(part)
        parts = list(reversed(trimmed))

        lines = []
        if parts:
            lines.append("## Conversation flow (narration and dialogue before and within this batch):")
            lines.extend(parts)
        return "\n".join(lines)

    def _call_llm(self, prompt: str) -> list[dict[str, Any]] | None:
        client = self._get_client()

        is_thinking_model = any(
            tag in self.model.lower()
            for tag in ("qwen", "deepseek", "r1", "qwq")
        )

        messages = [{"role": "user", "content": prompt}]
        extra: dict[str, Any] = {}
        if is_thinking_model:
            messages.append({"role": "assistant", "content": "<think>\n\n</think>"})
            extra["continue_assistant_turn"] = True

        max_attempts = 3
        for attempt in range(max_attempts):
            try:
                response = queued_chat_completion(
                    client,
                    model=self.model,
                    messages=messages,
                    temperature=0.1,
                    max_tokens=16384,
                    **extra,
                )
                break
            except Exception as e:
                err = str(e)
                logger.warning("LLM call failed (attempt %d/%d): %s", attempt + 1, max_attempts, err)
                # Log the full exception for debugging
                import traceback
                logger.debug("LLM call exception traceback:\n%s", traceback.format_exc())
                if attempt < max_attempts - 1:
                    time.sleep(3 * (attempt + 1))
                else:
                    logger.error("LLM call failed after %d attempts: %s", max_attempts, err)
                    return None

        # Handle potential None response
        if not response or not response.choices:
            logger.warning("Empty LLM response object")
            return None

        message = response.choices[0].message
        if not message:
            logger.warning("LLM response has no message")
            return None

        content = message.content or ""
        content = content.strip()

        if not content:
            reasoning = getattr(message, "reasoning_content", None) or ""
            reasoning = reasoning.strip()
            if reasoning:
                logger.warning(
                    "Empty LLM response content; falling back to reasoning_content (%d chars)",
                    len(reasoning),
                )
                return self._parse_response(reasoning)
            logger.warning("Empty LLM response content")
            return None

        logger.debug("LLM raw response content (first 500 chars): %s", content[:500])
        return self._parse_response(content)

    def _parse_response(self, content: str) -> list[dict[str, Any]] | None:
        stripped = content.strip()

        if stripped.startswith("```"):
            first_newline = stripped.find("\n")
            if first_newline > 0:
                stripped = stripped[first_newline + 1 :]
            if stripped.endswith("```"):
                stripped = stripped[:-3].strip()

        result = self._try_json_parse(stripped)
        if result is not None:
            return result

        json_start = stripped.find("[")
        json_end = stripped.rfind("]") + 1
        if json_start >= 0 and json_end > json_start:
            candidate = stripped[json_start:json_end]
            result = self._try_json_parse(candidate)
            if result is not None:
                return result

            last_brace = candidate.rfind("}")
            if last_brace > 0:
                candidate = candidate[: last_brace + 1]
                result = self._try_json_parse(candidate)
                if result is not None:
                    return result

        logger.error("Failed to parse LLM attribution response")
        return None

    def _try_json_parse(self, text: str) -> list[dict[str, Any]] | None:
        try:
            result = json.loads(text)
            if isinstance(result, list):
                return [dict(x) for x in result if isinstance(x, dict)]
            if isinstance(result, dict) and "utterances" in result:
                inner = result["utterances"]
                if isinstance(inner, list):
                    return [dict(x) for x in inner if isinstance(x, dict)]
            return None
        except json.JSONDecodeError:
            pass

        fixed = text
        for _ in range(10):
            try:
                result = json.loads(fixed)
                if isinstance(result, list):
                    return [dict(x) for x in result if isinstance(x, dict)]
                return None
            except json.JSONDecodeError:
                last_brace = fixed.rfind("}")
                last_bracket = fixed.rfind("]")
                last_comma = fixed.rfind(",")
                cut = max(last_brace, last_bracket, last_comma)
                if cut > 0:
                    fixed = fixed[:cut]
                    if fixed.endswith(","):
                        fixed = fixed[:-1]
                else:
                    break

        return None