UmbraAudio/umbraaudio/analysis/character_discovery.py

from __future__ import annotations

import json
import logging
import re
import time
from dataclasses import dataclass, field
from pathlib import Path
from collections.abc import Callable
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from umbraaudio.config import LLMConfig

logger = logging.getLogger(__name__)

from umbraaudio.voices.lmstudio_gate import queued_chat_completion

PROMPT_PATH = Path(__file__).parent / "CHARACTER_DISCOVERY_PROMPT.md"
SCENE_BREAK = "<<<>>>"


@dataclass
class Demographics:
    type: str = "Human"
    gender: str = "Unknown"

    def to_dict(self) -> dict:
        return {"type": self.type, "gender": self.gender}

    @classmethod
    def from_dict(cls, data: dict) -> Demographics:
        return cls(type=data.get("type", "Human"), gender=data.get("gender", "Unknown"))


@dataclass
class CommunicationStyle:
    primary_register: str = "Casual"
    style_shifts_or_notes: str = ""

    def to_dict(self) -> dict:
        return {
            "primary_register": self.primary_register,
            "style_shifts_or_notes": self.style_shifts_or_notes,
        }

    @classmethod
    def from_dict(cls, data: dict) -> CommunicationStyle:
        return cls(
            primary_register=data.get("primary_register", "Casual"),
            style_shifts_or_notes=data.get("style_shifts_or_notes", ""),
        )


@dataclass
class Speaker:
    primary_name: str
    aliases: list[str] = field(default_factory=list)
    demographics: Demographics = field(default_factory=Demographics)
    vocal_traits: list[str] = field(default_factory=list)
    communication_style: CommunicationStyle = field(default_factory=CommunicationStyle)
    is_narrator: bool = False
    is_group: bool = False
    description: str = ""
    sample_dialogue: str = ""
    chapters_appeared: list[int] = field(default_factory=list)
    unidentified: bool = False
    unidentified_reason: str = ""
    alias_evidence: list[str] = field(default_factory=list)

    @property
    def all_names(self) -> list[str]:
        return [self.primary_name] + self.aliases

    def to_dict(self) -> dict:
        return {
            "primary_name": self.primary_name,
            "aliases": self.aliases,
            "demographics": self.demographics.to_dict(),
            "vocal_traits": self.vocal_traits,
            "communication_style": self.communication_style.to_dict(),
            "is_narrator": self.is_narrator,
            "is_group": self.is_group,
            "description": self.description,
            "sample_dialogue": self.sample_dialogue,
            "chapters_appeared": self.chapters_appeared,
            "unidentified": self.unidentified,
            "unidentified_reason": self.unidentified_reason,
            "alias_evidence": self.alias_evidence,
        }

    @classmethod
    def from_dict(cls, data: dict) -> Speaker:
        demos = data.get("demographics", {})
        comm = data.get("communication_style", {})
        return cls(
            primary_name=data.get("primary_name", ""),
            aliases=data.get("aliases", []),
            demographics=Demographics.from_dict(demos),
            vocal_traits=data.get("vocal_traits", []),
            communication_style=CommunicationStyle.from_dict(comm),
            is_narrator=data.get("is_narrator", False),
            is_group=data.get("is_group", False),
            description=data.get("description", ""),
            sample_dialogue=data.get("sample_dialogue", ""),
            chapters_appeared=data.get("chapters_appeared", []),
            unidentified=data.get("unidentified", False),
            unidentified_reason=data.get("unidentified_reason", ""),
            alias_evidence=data.get("alias_evidence", []),
        )


@dataclass
class ChapterSpeakers:
    chapter_index: int
    chapter_title: str
    speaking_characters: list[dict]
    non_speaking_characters: list[dict]
    validation_notes: list[str] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "chapter_index": self.chapter_index,
            "chapter_title": self.chapter_title,
            "speaking_characters": self.speaking_characters,
            "non_speaking_characters": self.non_speaking_characters,
            "validation_notes": self.validation_notes,
        }


@dataclass
class DiscoveryResult:
    speakers: list[Speaker]
    chapter_results: list[ChapterSpeakers] = field(default_factory=list)
    raw_responses: list[str] = field(default_factory=list)
    book_title: str = ""
    total_chapters: int = 0
    total_words: int = 0
    validation_log: list[str] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "speakers": [s.to_dict() for s in self.speakers],
            "chapter_results": [cr.to_dict() for cr in self.chapter_results],
            "book_title": self.book_title,
            "total_chapters": self.total_chapters,
            "total_words": self.total_words,
            "validation_log": self.validation_log,
        }

    def save(self, output_path: Path) -> None:
        has_narrator = any(s.is_narrator for s in self.speakers)
        if not has_narrator:
            narrator = Speaker(
                primary_name="Narrator",
                aliases=["narrator", "the narrator"],
                demographics=Demographics(type="Human", gender="Male"),
                vocal_traits=["deep", "resonant", "calm", "authoritative"],
                communication_style=CommunicationStyle(
                    primary_register="Formal",
                    style_shifts_or_notes="Steady, even-toned throughout. Never rushes.",
                ),
                description=(
                    "Deep, resonant male voice. Calm, authoritative tone. "
                    "Rich baritone timbre, slow and deliberate pacing, "
                    "like a late-night radio host."
                ),
                is_narrator=True,
            )
            self.speakers.insert(0, narrator)
            logger.info("Added default narrator speaker")

        output_path.parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "w") as f:
            json.dump(self.to_dict(), f, indent=2)
        logger.info("Saved discovery results to %s", output_path)

    @classmethod
    def load(cls, path: Path) -> DiscoveryResult:
        with open(path) as f:
            data = json.load(f)
        return cls(
            speakers=[Speaker.from_dict(s) for s in data.get("speakers", [])],
            chapter_results=[
                ChapterSpeakers(
                    chapter_index=cr.get("chapter_index", 0),
                    chapter_title=cr.get("chapter_title", ""),
                    speaking_characters=cr.get("speaking_characters", []),
                    non_speaking_characters=cr.get("non_speaking_characters", []),
                    validation_notes=cr.get("validation_notes", []),
                )
                for cr in data.get("chapter_results", [])
            ],
            book_title=data.get("book_title", ""),
            total_chapters=data.get("total_chapters", 0),
            total_words=data.get("total_words", 0),
            validation_log=data.get("validation_log", []),
        )


class CharacterDiscoverer:
    def __init__(self, llm_config: LLMConfig, data_dir: Path | None = None) -> None:
        self.config = llm_config
        self._client = None
        self._prompt_template: str | None = None
        self._active_model: str | None = None
        self.data_dir = data_dir or Path(".umbra")

    def _load_prompt(self) -> str:
        if self._prompt_template is None:
            self._prompt_template = PROMPT_PATH.read_text()
        return self._prompt_template

    def _get_client(self) -> Any:
        if self._client is not None:
            return self._client

        if self.config.provider in ("lmstudio", "auto"):
            try:
                import httpx

                resp = httpx.get(f"{self.config.lmstudio_base_url}/models", timeout=5)
                if resp.status_code == 200:
                    data = resp.json()
                    models = data.get("data", [])
                    embedding_models = {
                        "text-embedding-nomic-embed-text-v1.5",
                        "text-embedding-ada-002",
                    }
                    chat_models = [
                        m["id"] for m in models if m.get("id") and m["id"] not in embedding_models
                    ]
                    if chat_models:
                        preferred = [
                            self.config.discovery_model or "google/gemma-4-e4b",
                            "google/gemma-4-e4b",
                            "qwen3.5-9b",
                        ]
                        self._active_model = next(
                            (m for m in preferred if m in chat_models),
                            chat_models[0],
                        )
                        logger.info("Discovery using LM Studio model: %s", self._active_model)
                    from openai import OpenAI

                    self._client = OpenAI(
                        base_url=self.config.lmstudio_base_url,
                        api_key="lm-studio",
                        timeout=600,
                    )
                    return self._client
            except Exception:
                pass

        if self.config.provider in ("openai", "auto"):
            from openai import OpenAI

            self._client = OpenAI()
            self._active_model = self.config.openai_model
            return self._client

        raise RuntimeError("No LLM provider available")

    def _call_llm(
        self,
        text: str,
        context: str = "",
        max_chars: int = 40000,
        known_speakers: str = "",
        pass_context: str = "",
        on_retry: Callable[[int], None] | None = None,
    ) -> dict:
        if len(text) > max_chars:
            logger.info("Truncating chapter text from %d to %d chars", len(text), max_chars)
            text = text[:max_chars]

        prompt_template = self._load_prompt()
        prompt = prompt_template.replace("{book_text}", text)

        if known_speakers:
            prompt = prompt.replace("{known_speakers_section}", known_speakers)
        else:
            prompt = prompt.replace("{known_speakers_section}", "")

        if pass_context:
            prompt = prompt.replace("{pass_context}", pass_context)
        else:
            prompt = prompt.replace("{pass_context}", "")

        if context:
            prompt = f"{context}\n\n{prompt}"

        client = self._get_client()
        model = self._active_model or self.config.openai_model

        max_attempts = 6
        response = None
        for attempt in range(max_attempts):
            try:
                response = queued_chat_completion(
                    client,
                    model=model,
                    messages=[{"role": "user", "content": prompt}],
                    temperature=0.2,
                    max_tokens=8192,
                    timeout=600,
                )
                break
            except Exception as e:
                err = str(e)
                logger.warning(
                    "LLM call failed (attempt %d/%d): %s", attempt + 1, max_attempts, err
                )
                if on_retry is not None:
                    try:
                        on_retry(attempt + 1)
                    except Exception:
                        pass
                if attempt < max_attempts - 1:
                    time.sleep(min(60, 5 * (attempt + 1)))
                else:
                    logger.error("LLM call failed after %d attempts: %s", max_attempts, err)
                    raise
        assert response is not None

        content = response.choices[0].message.content or ""
        content = content.strip()

        if not content:
            logger.warning("Empty LLM response")
            return {"speaking_characters": [], "non_speaking_characters": []}

        code_block = re.search(r"```(?:json)?\s*\n?(.*?)\n?\s*```", content, re.DOTALL)
        if code_block:
            content = code_block.group(1).strip()

        try:
            return json.loads(content)
        except json.JSONDecodeError as e:
            logger.warning("JSON parse error: %s", e)

        json_start = content.find("{")
        json_end = content.rfind("}") + 1
        if json_start >= 0 and json_end > json_start:
            candidate = content[json_start:json_end]
            try:
                return json.loads(candidate)
            except json.JSONDecodeError:
                pass

            for _ in range(5):
                open_braces = candidate.count("{") - candidate.count("}")
                open_brackets = candidate.count("[") - candidate.count("]")
                close = "}" * open_braces + "]" * open_brackets
                fixed = candidate + close
                try:
                    return json.loads(fixed)
                except json.JSONDecodeError:
                    last_brace = candidate.rfind("}")
                    last_bracket = candidate.rfind("]")
                    last_comma = candidate.rfind(",")
                    cut = max(last_brace, last_bracket, last_comma)
                    if cut > 0:
                        candidate = candidate[:cut]
                        if candidate.endswith(","):
                            candidate = candidate[:-1]
                    else:
                        break

        logger.error("Failed to parse LLM response")
        return {"speaking_characters": [], "non_speaking_characters": []}

    def _parse_scenes(self, chapter_text: str) -> list[str]:
        if SCENE_BREAK in chapter_text:
            return [s.strip() for s in chapter_text.split(SCENE_BREAK) if s.strip()]
        paragraphs = re.split(r"\n\s*\n", chapter_text)
        scenes: list[str] = []
        current = ""
        for para in paragraphs:
            if len(current) + len(para) > 3000:
                if current.strip():
                    scenes.append(current.strip())
                current = para
            else:
                current = current + "\n\n" + para if current else para
        if current.strip():
            scenes.append(current.strip())
        return scenes if scenes else [chapter_text]

    def _find_speakers_in_text(self, text: str, known_names: list[str]) -> list[str]:
        found: list[str] = []
        sorted_names = sorted(known_names, key=len, reverse=True)
        pattern = re.compile(
            r"\b(" + "|".join(re.escape(n) for n in sorted_names) + r")\b",
            re.IGNORECASE,
        )
        for match in pattern.finditer(text):
            matched = match.group(0)
            for name in sorted_names:
                if name.lower() == matched.lower():
                    if name not in found:
                        found.append(name)
                    break
        return found

    TITLE_WORDS = {
        "captain", "princess", "prince", "king", "queen", "doctor", "dr",
        "professor", "prof", "aunt", "auntie", "uncle", "sir", "lady",
        "lord", "mr", "mrs", "ms", "miss", "master", "sister", "brother",
        "father", "mother", "grandfather", "grandmother", "grandma", "grandpa",
        "chief", "commander", "major", "colonel", "general", "admiral",
        "sergeant", "lieutenant", "corporal", "private", "ensign", "mister",
        "madam", "madame", "dame", "saint", "emperor", "empress",
        "duke", "duchess", "count", "countess", "baron", "baroness",
        "ambassador", "senator", "governor", "judge", "justice",
        "reverend", "pastor", "bishop", "cardinal", "rabbi",
        "imam", "sheikh", "sultan",
    }

    def _has_two_proper_names(self, name_a: str, name_b: str) -> bool:
        a_lower = name_a.lower()
        b_lower = name_b.lower()

        if a_lower in b_lower or b_lower in a_lower:
            return False

        def strip_titles(name: str) -> list[str]:
            return [w for w in name.split() if w.lower() not in self.TITLE_WORDS]

        words_a = strip_titles(name_a)
        words_b = strip_titles(name_b)

        if not words_a or not words_b:
            return False

        all_proper_a = all(w[0].isupper() and len(w) > 1 for w in words_a)
        all_proper_b = all(w[0].isupper() and len(w) > 1 for w in words_b)

        if not (all_proper_a and all_proper_b):
            return False

        a_word_set = set(w.lower() for w in words_a)
        b_word_set = set(w.lower() for w in words_b)

        if a_word_set & b_word_set:
            return False

        for wa in words_a:
            for wb in words_b:
                wa_l = wa.lower()
                wb_l = wb.lower()
                min_len = min(len(wa_l), len(wb_l))
                if min_len >= 4 and wa_l[:min_len] == wb_l[:min_len]:
                    return False

        return True

    @staticmethod
    def _alias_has_evidence(char_data: dict[str, Any], alias: str) -> bool:
        alias_lower = alias.lower()
        return any(
            alias_lower in ev.lower()
            for ev in char_data.get("alias_evidence", [])
        )

    def validate_aliases(
        self, chapter_text: str, chapter_title: str, characters: list[dict[str, Any]]
    ) -> tuple[list[dict[str, Any]], list[str]]:
        scenes = self._parse_scenes(chapter_text)
        all_names: list[str] = []
        for char in characters:
            name = char.get("primary_name", "")
            if name and name not in all_names:
                all_names.append(name)
            for alias in char.get("aliases", []):
                if alias not in all_names:
                    all_names.append(alias)

        invalid_aliases: list[tuple[str, str, str]] = []
        scene_speakers: list[list[str]] = []
        for _scene_idx, scene_text in enumerate(scenes):
            speakers_in_scene = self._find_speakers_in_text(scene_text, all_names)
            scene_speakers.append(speakers_in_scene)

        for char in characters:
            primary = char.get("primary_name", "")
            aliases = list(char.get("aliases", []))
            valid_aliases: list[str] = []

            for alias in aliases:
                is_invalid = False

                if not is_invalid:
                    for scene_idx, speakers in enumerate(scene_speakers):
                        if (
                            primary in speakers
                            and alias in speakers
                            and primary.lower() != alias.lower()
                        ):
                            invalid_aliases.append(
                                (
                                    primary,
                                    alias,
                                    f"Both appear in scene {scene_idx + 1}",
                                )
                            )
                            is_invalid = True
                            break

                if not is_invalid:
                    valid_aliases.append(alias)

            char["aliases"] = valid_aliases

        validation_notes: list[str] = []
        for primary, alias, reason in invalid_aliases:
            note = f"Removed alias '{alias}' from '{primary}': {reason}"
            validation_notes.append(note)
            logger.info("Alias validation: %s", note)

        for char in characters:
            if not char.get("alias_evidence"):
                char["alias_evidence"] = []
            primary = char.get("primary_name", "")
            for alias in char.get("aliases", []):
                has_evidence = any(
                    primary.lower() in e.lower() and alias.lower() in e.lower()
                    for e in char["alias_evidence"]
                )
                if not has_evidence:
                    char["alias_evidence"].append(
                        f"'{alias}' validated by text co-occurrence check"
                    )

        return characters, validation_notes

    def _merge_chapter_results(
        self, chapter_results: list[ChapterSpeakers], book_title: str, total_words: int
    ) -> list[Speaker]:
        merged: dict[str, Speaker] = {}
        alias_frequency: dict[str, dict[str, set[int]]] = {}

        for cr in chapter_results:
            chapter_num = cr.chapter_index + 1

            for char_data in cr.speaking_characters:
                primary = char_data.get("primary_name", "").strip()
                if not primary:
                    continue

                if primary in merged:
                    existing = merged[primary]
                    existing.chapters_appeared.append(chapter_num)

                    for alias in char_data.get("aliases", []):
                        if alias.lower() == primary.lower():
                            continue
                        alias_frequency.setdefault(primary, {}).setdefault(
                            alias, set()
                        ).add(chapter_num)
                        if (
                            alias not in existing.aliases
                            and (
                                not self._has_two_proper_names(primary, alias)
                                or self._alias_has_evidence(char_data, alias)
                            )
                        ):
                            existing.aliases.append(alias)

                    for trait in char_data.get("vocal_traits", []):
                        if trait not in existing.vocal_traits:
                            existing.vocal_traits.append(trait)

                    demos = char_data.get("demographics", {})
                    if demos.get("type") and demos["type"] != "Human":
                        existing.demographics.type = demos["type"]
                    if demos.get("gender") and demos["gender"] != "Unknown":
                        existing.demographics.gender = demos["gender"]

                    comm = char_data.get("communication_style", {})
                    if (
                        comm.get("style_shifts_or_notes")
                        and comm["style_shifts_or_notes"]
                        not in existing.communication_style.style_shifts_or_notes
                    ):
                        existing.communication_style.style_shifts_or_notes = (
                            f"{existing.communication_style.style_shifts_or_notes} | "
                            f"Ch.{chapter_num}: {comm['style_shifts_or_notes']}"
                            if existing.communication_style.style_shifts_or_notes
                            else f"Ch.{chapter_num}: {comm['style_shifts_or_notes']}"
                        )

                    for ev in char_data.get("alias_evidence", []):
                        if ev not in existing.alias_evidence:
                            existing.alias_evidence.append(ev)

                    if char_data.get("unidentified") and not existing.unidentified:
                        existing.unidentified = True
                        existing.unidentified_reason = char_data.get("unidentified_reason", "")
                else:
                    demos = char_data.get("demographics", {})
                    comm = char_data.get("communication_style", {})
                    raw_aliases = char_data.get("aliases", [])
                    filtered_aliases = [
                        a for a in raw_aliases
                        if a.lower() == primary.lower()
                        or not self._has_two_proper_names(primary, a)
                        or self._alias_has_evidence(char_data, a)
                    ]
                    merged[primary] = Speaker(
                        primary_name=primary,
                        aliases=filtered_aliases,
                        demographics=Demographics.from_dict(demos),
                        vocal_traits=char_data.get("vocal_traits", []),
                        communication_style=CommunicationStyle.from_dict(comm),
                        description=char_data.get("description", ""),
                        chapters_appeared=[chapter_num],
                        unidentified=char_data.get("unidentified", False),
                        unidentified_reason=char_data.get("unidentified_reason", ""),
                        alias_evidence=char_data.get("alias_evidence", []),
                    )

            for char_data in cr.non_speaking_characters:
                primary = char_data.get("primary_name", "").strip()
                if not primary:
                    continue

                if primary in merged:
                    existing = merged[primary]
                    if chapter_num not in existing.chapters_appeared:
                        existing.chapters_appeared.append(chapter_num)

        for primary_name, speaker in list(merged.items()):
            chapters_count = len(speaker.chapters_appeared)
            if chapters_count >= 3 and primary_name in alias_frequency:
                freq = alias_frequency[primary_name]
                speaker.aliases = [
                    a for a in speaker.aliases
                    if a.lower() == primary_name.lower()
                    or len(freq.get(a, set())) > 1
                    or a.lower() not in {n.lower() for n in freq}
                ]

        merged = self._deduplicate_by_aliases(merged)
        return sorted(merged.values(), key=lambda s: len(s.chapters_appeared), reverse=True)

    def _deduplicate_by_aliases(self, merged: dict[str, Speaker]) -> dict[str, Speaker]:
        changed = True
        while changed:
            changed = False
            names = list(merged.keys())
            for i, name_a in enumerate(names):
                if name_a not in merged:
                    continue
                speaker_a = merged[name_a]

                for name_b in names[i + 1 :]:
                    if name_b not in merged:
                        continue
                    speaker_b = merged[name_b]

                    name_a_lower = name_a.lower()
                    name_b_lower = name_b.lower()

                    if name_a_lower == name_b_lower:
                        self._absorb_into_richer(merged, name_a, name_b)
                        changed = True
                        if name_a not in merged:
                            break
                        continue

                    b_aliases_lower = [a.lower() for a in speaker_b.aliases]
                    a_aliases_lower = [a.lower() for a in speaker_a.aliases]

                    b_has_a = name_a_lower in b_aliases_lower
                    a_has_b = name_b_lower in a_aliases_lower
                    shared_specific = self._shared_specific_alias(speaker_a, speaker_b)

                    if (b_has_a or a_has_b) and self._should_merge_by_context(
                        speaker_a, speaker_b
                    ):
                        self._absorb_into_richer(merged, name_a, name_b)
                        changed = True
                        if name_a not in merged:
                            break
                    elif shared_specific:
                        logger.info(
                            "Dedup: merging '%s' and '%s' (shared specific alias)",
                            name_a, name_b,
                        )
                        self._absorb_into_richer(merged, name_a, name_b)
                        changed = True
                        if name_a not in merged:
                            break

        return merged

    def _shared_specific_alias(self, speaker_a: Speaker, speaker_b: Speaker) -> bool:
        def is_specific(name: str) -> bool:
            words = [w for w in name.split() if w.lower() not in self.TITLE_WORDS]
            return bool(words)

        a_names = {
            n.lower()
            for n in [speaker_a.primary_name] + speaker_a.aliases
            if is_specific(n)
        }
        b_names = {
            n.lower()
            for n in [speaker_b.primary_name] + speaker_b.aliases
            if is_specific(n)
        }
        return bool(a_names & b_names)

    def _absorb_into_richer(
        self, merged: dict[str, Speaker], name_a: str, name_b: str
    ) -> None:
        if name_a not in merged or name_b not in merged:
            return
        speaker_a = merged[name_a]
        speaker_b = merged[name_b]

        if len(speaker_b.chapters_appeared) > len(speaker_a.chapters_appeared):
            self._absorb_speaker(speaker_b, speaker_a)
            merged.pop(name_a, None)
        else:
            self._absorb_speaker(speaker_a, speaker_b)
            merged.pop(name_b, None)

    def _should_merge_by_context(self, speaker_a: Speaker, speaker_b: Speaker) -> bool:
        all_a = [speaker_a.primary_name] + speaker_a.aliases
        all_b = [speaker_b.primary_name] + speaker_b.aliases
        for name_a in all_a:
            for name_b in all_b:
                if self._has_two_proper_names(name_a, name_b):
                    logger.info(
                        "Skipping merge: '%s' and '%s' have distinct proper names",
                        name_a,
                        name_b,
                    )
                    return False
        return True

    def _absorb_speaker(self, target: Speaker, source: Speaker) -> None:
        for ch in source.chapters_appeared:
            if ch not in target.chapters_appeared:
                target.chapters_appeared.append(ch)
        for alias in source.aliases:
            if alias not in target.aliases:
                target.aliases.append(alias)
        for trait in source.vocal_traits:
            if trait not in target.vocal_traits:
                target.vocal_traits.append(trait)
        for ev in source.alias_evidence:
            if ev not in target.alias_evidence:
                target.alias_evidence.append(ev)

        if not target.description and source.description:
            target.description = source.description
        if not target.sample_dialogue and source.sample_dialogue:
            target.sample_dialogue = source.sample_dialogue
        if source.demographics.gender != "Unknown" and target.demographics.gender == "Unknown":
            target.demographics.gender = source.demographics.gender
        if source.demographics.type != "Human" and target.demographics.type == "Human":
            target.demographics.type = source.demographics.type
        if not target.communication_style.style_shifts_or_notes:
            target.communication_style = source.communication_style
        target.is_group = target.is_group or source.is_group
        target.unidentified = target.unidentified or source.unidentified
        if not target.unidentified_reason and source.unidentified_reason:
            target.unidentified_reason = source.unidentified_reason

    def _save_temp(self, data: dict[str, Any], path: Path) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            json.dump(data, f, indent=2)
        logger.info("Saved temporary data to %s", path)

    def _load_temp(self, path: Path) -> dict[str, Any] | None:
        if not path.exists():
            return None
        try:
            with open(path) as f:
                return json.load(f)
        except (json.JSONDecodeError, OSError) as e:
            logger.warning("Failed to load temp file %s: %s", path, e)
            return None

    async def _run_pass1(
        self,
        chapters: list[dict[str, Any]],
        book_title: str,
        max_chars: int,
        temp_dir: Path,
    ) -> DiscoveryResult:
        chapter_results: list[ChapterSpeakers] = []
        raw_responses: list[str] = []
        known_speakers: list[dict[str, Any]] = []

        logger.info("Pass 1: Chapter-by-chapter discovery in %d chapters", len(chapters))

        ckpt_path = temp_dir / "discovery_pass1_chapters.json"
        start_idx = 0
        if ckpt_path.exists():
            ckpt = self._load_temp(ckpt_path)
            if ckpt:
                chapter_results = [
                    ChapterSpeakers(
                        chapter_index=cr.get("chapter_index", 0),
                        chapter_title=cr.get("chapter_title", ""),
                        speaking_characters=cr.get("speaking_characters", []),
                        non_speaking_characters=cr.get("non_speaking_characters", []),
                        validation_notes=cr.get("validation_notes", []),
                    )
                    for cr in ckpt.get("chapter_results", [])
                ]
                raw_responses = ckpt.get("raw_responses", [])
                known_speakers = ckpt.get("known_speakers", [])
                start_idx = len(chapter_results)
                logger.info(
                    "Pass 1: resuming from checkpoint with %d/%d chapters done",
                    start_idx,
                    len(chapters),
                )

        for i, ch in enumerate(chapters):
            if i < start_idx:
                continue
            ch_text = ch.get("source_text", "")
            ch_title = ch.get("title", "") or f"Chapter {i + 1}"

            logger.info(
                "Pass 1 — Chapter %d/%d: %s (%d chars)",
                i + 1,
                len(chapters),
                ch_title,
                len(ch_text),
            )

            known_section = ""
            if known_speakers:
                lines = [
                    "## Previously Discovered Speakers",
                    (
                        "The following speakers have been identified in earlier "
                        "chapters. Use this list to avoid creating duplicates."
                    ),
                    "",
                ]
                for ks in known_speakers:
                    name = ks.get("primary_name", "")
                    aliases = ks.get("aliases", [])
                    demos = ks.get("demographics", {})
                    gender = demos.get("gender", "Unknown")
                    stype = demos.get("type", "Human")
                    alias_str = f" (aliases: {', '.join(aliases)})" if aliases else ""
                    lines.append(f"- **{name}**{alias_str} — {gender}, {stype}")
                lines.append("")
                known_section = "\n".join(lines)

            chapter_prompt = f"--- Chapter {i + 1}: {ch_title} ---\n\n"
            data = self._call_llm(
                ch_text,
                context=chapter_prompt,
                max_chars=max_chars,
                known_speakers=known_section,
            )

            raw_responses.append(json.dumps(data))

            chapter_results.append(
                ChapterSpeakers(
                    chapter_index=i,
                    chapter_title=ch_title,
                    speaking_characters=data.get("speaking_characters", []),
                    non_speaking_characters=data.get("non_speaking_characters", []),
                    validation_notes=data.get("validation_notes", []),
                )
            )

            for char_data in data.get("speaking_characters", []):
                primary = char_data.get("primary_name", "").strip()
                if not primary:
                    continue
                exists = any(k["primary_name"] == primary for k in known_speakers)
                if not exists:
                    known_speakers.append(
                        {
                            "primary_name": primary,
                            "aliases": char_data.get("aliases", []),
                            "demographics": char_data.get("demographics", {}),
                        }
                    )

            import asyncio

            await asyncio.sleep(1)

            self._save_temp(
                {
                    "chapter_results": [cr.to_dict() for cr in chapter_results],
                    "raw_responses": raw_responses,
                    "known_speakers": known_speakers,
                },
                ckpt_path,
            )

        logger.info("Pass 1: Merging results from %d chapters", len(chapter_results))
        merged_speakers = self._merge_chapter_results(chapter_results, book_title, 0)

        total_words = sum(len(ch.get("source_text", "").split()) for ch in chapters)

        result = DiscoveryResult(
            speakers=merged_speakers,
            chapter_results=chapter_results,
            raw_responses=raw_responses,
            book_title=book_title,
            total_chapters=len(chapters),
            total_words=total_words,
        )

        pass1_data = result.to_dict()
        self._save_temp(pass1_data, temp_dir / "discovery_pass1.json")

        if ckpt_path.exists():
            ckpt_path.unlink()
            logger.info("Pass 1: checkpoint cleared after full run")

        logger.info(
            "Pass 1 complete: %d speakers across %d chapters",
            len(merged_speakers),
            len(chapters),
        )

        return result

    async def _run_reconciliation(
        self,
        pass1_result: DiscoveryResult,
        chapters: list[dict[str, Any]],
        max_chars: int,
        temp_dir: Path,
    ) -> DiscoveryResult:
        if len(pass1_result.speakers) < 3:
            logger.info("Reconciliation skipped: only %d speakers", len(pass1_result.speakers))
            return pass1_result

        logger.info(
            "Reconciliation: checking %d speakers for duplicates", len(pass1_result.speakers)
        )

        speaker_list_lines: list[str] = []
        for s in pass1_result.speakers:
            aliases_str = f" (aliases: {', '.join(s.aliases)})" if s.aliases else ""
            ch_str = f" [chapters: {s.chapters_appeared}]" if s.chapters_appeared else ""
            speaker_list_lines.append(
                f"- {s.primary_name}{aliases_str}{ch_str}: "
                f"{s.demographics.gender}, {s.description[:80]}"
            )
        speaker_list = "\n".join(speaker_list_lines)

        excerpt_parts: list[str] = []
        total_chars = 0
        for ch in chapters:
            text = ch.get("source_text", "")
            title = ch.get("title", "")
            if not text:
                continue
            sample = text[:2000] if len(text) > 2000 else text
            excerpt_parts.append(f"--- {title} (excerpt) ---\n{sample}")
            total_chars += len(sample)
            if total_chars > max_chars:
                break
        excerpts = "\n\n".join(excerpt_parts)

        prompt = f"""You are analyzing a list of characters discovered in a book to find DUPLICATE entries.

Two speaker entries should be MERGED if they refer to the SAME person, even if called by different names.
For example: "Robert" and "Bob" in different chapters are the same person if context confirms it.

Two entries should NOT be merged if:
- They have different proper names AND both appear as speakers in the same scene
- One is a title shared by multiple people ("Captain" in different contexts)
- There is no textual evidence they are the same person

## Discovered Speakers
{speaker_list}

## Chapter Excerpts (for context)
{excerpts}

## Instructions
Return a JSON object with a single key "merges" containing a list of merge operations.
Each merge: {{"keep": "<primary_name to keep>", "absorb": "<primary_name to remove>", "reason": "<brief explanation>"}}
If no merges are needed, return {{"merges": []}}.

Return ONLY valid JSON."""

        client = self._get_client()
        model = self._active_model or self.config.openai_model

        try:
            response = queued_chat_completion(
                client,
                model=model,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.1,
                max_tokens=4096,
            )
            content = response.choices[0].message.content or ""
            content = content.strip()

            code_block = re.search(r"```(?:json)?\s*\n?(.*?)\n?\s*```", content, re.DOTALL)
            if code_block:
                content = code_block.group(1).strip()

            data = json.loads(content)
            merges = data.get("merges", [])
        except Exception as e:
            logger.warning("Reconciliation LLM call failed: %s", e)
            return pass1_result

        if not merges:
            logger.info("Reconciliation: no merges recommended")
            return pass1_result

        speaker_map: dict[str, Speaker] = {
            s.primary_name: s for s in pass1_result.speakers
        }
        merged_count = 0
        for merge in merges:
            keep_name = merge.get("keep", "")
            absorb_name = merge.get("absorb", "")
            reason = merge.get("reason", "")

            if keep_name not in speaker_map or absorb_name not in speaker_map:
                continue
            if keep_name == absorb_name:
                continue

            keep_speaker = speaker_map[keep_name]
            absorb_speaker = speaker_map[absorb_name]

            if not self._should_merge_by_context(keep_speaker, absorb_speaker):
                logger.info(
                    "Reconciliation: skipping merge %s -> %s (context check failed)",
                    absorb_name, keep_name,
                )
                continue

            self._absorb_speaker(keep_speaker, absorb_speaker)
            if absorb_name not in keep_speaker.aliases:
                keep_speaker.aliases.append(absorb_name)
            del speaker_map[absorb_name]
            merged_count += 1
            logger.info(
                "Reconciliation: merged '%s' into '%s' — %s",
                absorb_name, keep_name, reason,
            )

        pass1_result.speakers = sorted(
            speaker_map.values(), key=lambda s: len(s.chapters_appeared), reverse=True
        )

        pass1b_data = pass1_result.to_dict()
        self._save_temp(pass1b_data, temp_dir / "discovery_pass1b_reconciliation.json")

        logger.info(
            "Reconciliation complete: %d merges applied, %d speakers remaining",
            merged_count, len(pass1_result.speakers),
        )
        return pass1_result

    async def _run_pass2(
        self,
        pass1_result: DiscoveryResult,
        chapters: list[dict[str, Any]],
        max_chars: int,
        temp_dir: Path,
    ) -> DiscoveryResult:
        logger.info("Pass 2: Alias validation across %d chapters", len(chapters))

        validation_log: list[str] = []

        for cr in pass1_result.chapter_results:
            chapter_idx = cr.chapter_index
            if chapter_idx >= len(chapters):
                continue

            ch_text = chapters[chapter_idx].get("source_text", "")
            ch_title = cr.chapter_title

            validated_chars, notes = self.validate_aliases(
                ch_text, ch_title, cr.speaking_characters
            )
            cr.speaking_characters = validated_chars
            cr.validation_notes = notes
            validation_log.extend(notes)

            if notes:
                logger.info(
                    "Pass 2 — Chapter %d (%s): %d alias corrections",
                    chapter_idx + 1,
                    ch_title,
                    len(notes),
                )

            import asyncio

            await asyncio.sleep(0.1)

        pass1_result.validation_log = validation_log

        logger.info("Pass 2: Re-merging with validated aliases")
        merged_speakers = self._merge_chapter_results(
            pass1_result.chapter_results,
            pass1_result.book_title,
            pass1_result.total_words,
        )
        pass1_result.speakers = merged_speakers

        pass2_data = pass1_result.to_dict()
        self._save_temp(pass2_data, temp_dir / "discovery_pass2.json")

        removed_count = sum(1 for n in validation_log if "Removed alias" in n)
        logger.info(
            "Pass 2 complete: %d alias corrections applied, %d speakers",
            removed_count,
            len(merged_speakers),
        )

        return pass1_result

    async def _run_pass3(
        self,
        pass2_result: DiscoveryResult,
        chapters: list[dict[str, Any]],
        max_chars: int,
        temp_dir: Path,
    ) -> DiscoveryResult:
        logger.info("Pass 3: Cross-chapter reconciliation")

        unidentified = [s for s in pass2_result.speakers if s.unidentified]
        if unidentified:
            logger.info("Pass 3: %d unidentified speakers flagged for review", len(unidentified))
            for s in unidentified:
                logger.info(
                    "  Unidentified: '%s' — %s",
                    s.primary_name,
                    s.unidentified_reason,
                )

        known_by_name: dict[str, Speaker] = {}
        for speaker in pass2_result.speakers:
            known_by_name[speaker.primary_name.lower()] = speaker
            for alias in speaker.aliases:
                if alias.lower() not in known_by_name:
                    known_by_name[alias.lower()] = speaker

        suspicious_merges: list[str] = []
        for speaker in pass2_result.speakers:
            all_names = [speaker.primary_name] + speaker.aliases
            bad_aliases: set[str] = set()
            for i, name_a in enumerate(all_names):
                for name_b in all_names[i + 1 :]:
                    if not self._has_two_proper_names(name_a, name_b):
                        continue
                    if self._alias_has_evidence({"alias_evidence": speaker.alias_evidence}, name_b):
                        continue
                    msg = (
                        f"'{speaker.primary_name}': removed alias "
                        f"'{name_b}' (distinct proper names: '{name_a}' vs '{name_b}')"
                    )
                    suspicious_merges.append(msg)
                    if name_b != speaker.primary_name:
                        bad_aliases.add(name_b)
                    if name_a != speaker.primary_name:
                        bad_aliases.add(name_a)
            if bad_aliases:
                speaker.aliases = [a for a in speaker.aliases if a not in bad_aliases]
                logger.warning(
                    "Pass 3: split %d bad aliases from '%s': %s",
                    len(bad_aliases), speaker.primary_name, list(bad_aliases),
                )

        for msg in suspicious_merges:
            logger.warning("Pass 3: %s", msg)
            pass2_result.validation_log.extend(suspicious_merges)

        pass3_data = pass2_result.to_dict()
        self._save_temp(pass3_data, temp_dir / "discovery_pass3.json")

        logger.info(
            "Pass 3 complete: %d speakers, %d unidentified, %d suspicious merges",
            len(pass2_result.speakers),
            len(unidentified),
            len(suspicious_merges),
        )

        return pass2_result

    async def discover_speakers(
        self,
        chapters: list[dict[str, Any]],
        book_title: str = "",
        max_chars: int = 40000,
    ) -> DiscoveryResult:
        data_dir = self.data_dir
        temp_dir = data_dir / "discovery"

        logger.info("Starting 4-pass character discovery for %d chapters", len(chapters))

        pass1_result = await self._run_pass1(chapters, book_title, max_chars, temp_dir)

        pass1b_result = await self._run_reconciliation(pass1_result, chapters, max_chars, temp_dir)

        pass2_result = await self._run_pass2(pass1b_result, chapters, max_chars, temp_dir)

        final_result = await self._run_pass3(pass2_result, chapters, max_chars, temp_dir)

        logger.info(
            "Discovery complete: %d speakers across %d chapters",
            len(final_result.speakers),
            len(chapters),
        )

        return final_result

    def get_speaker_names(self, result: DiscoveryResult) -> list[str]:
        names = []
        for speaker in result.speakers:
            names.append(speaker.primary_name)
            names.extend(speaker.aliases)
        return list(set(names))

    def get_canonical_for_name(self, name: str, result: DiscoveryResult) -> str | None:
        name_lower = name.lower()
        for speaker in result.speakers:
            if speaker.primary_name.lower() == name_lower:
                return speaker.primary_name
            for alias in speaker.aliases:
                if alias.lower() == name_lower:
                    return speaker.primary_name
        return None

    def get_speaker_by_canonical(
        self, canonical_name: str, result: DiscoveryResult
    ) -> Speaker | None:
        for speaker in result.speakers:
            if speaker.primary_name == canonical_name:
                return speaker
        return None