InkFlow/backend/inkflow/casting/dedup.py

"""Reconciliation du casting : deduplication des variantes de noms.

Probleme : un meme personnage apparait sous plusieurs formes ("Holden",
"James Holden", "James", "Jim"). Sans reconciliation, chaque forme devient un
personnage distinct avec sa propre voix -> incoherence a l'ecoute.

Strategie hybride :
1. Heuristique (sans LLM) : match exact sur nom/alias, puis sous-ensemble de
   tokens ("Holden" contenu dans "James Holden").
2. Gemma tranche les cas ambigus (plusieurs candidats compatibles, ou variante
   non evidente type "Jim" <-> "James") a l'aide des descriptions.

Chaque variante rencontree est conservee comme `alias` du personnage canonique ;
le nom canonique est la forme la plus complete vue ("James Holden"). Les
artefacts d'analyse (segments) ne sont PAS modifies : la resolution de voix au
rendu s'appuie sur les aliases (`casting/assign.py`).
"""
from __future__ import annotations

import re
from typing import Optional

from ..models import Character
from ..settings import get_settings

# Sentinelles internes.
_AMBIGUOUS = object()   # heuristique : plusieurs candidats -> on delegue a Gemma
_NEW = object()         # decision Gemma : nouveau personnage

# Mots vides / titres a ignorer pour le rapprochement par tokens.
_STOPWORDS = {
    "le", "la", "les", "un", "une", "de", "du", "des", "monsieur", "madame",
    "mademoiselle", "m", "mme", "mlle", "mr", "dr", "docteur", "capitaine",
    "lieutenant", "sergent", "general", "amiral", "the", "of",
}
_SPLIT_RE = re.compile(r"[^\wÀ-ÿ]+")

# Garde-fou de contexte (caracteres) pour le prompt Gemma.
_MAX_PROMPT_CHARS = 24000


def _norm(name: str) -> str:
    return name.strip().lower()


def _tokens(name: str) -> set[str]:
    """Tokens significatifs d'un nom (minuscules, sans titres ni mots vides)."""
    parts = [p for p in _SPLIT_RE.split(name.strip()) if p]
    return {p.lower() for p in parts
            if len(p) >= 2 and p.lower() not in _STOPWORDS}


def _completeness(name: str) -> tuple[int, int]:
    """Cle de tri du nom le plus "complet" : plus de tokens, puis plus long."""
    return (len(_tokens(name)), len(name.strip()))


def _forms(c: Character) -> list[str]:
    return [c.name, *c.aliases]


def _token_freq(characters: list[Character], extra: Optional[list[str]] = None):
    """Compte, pour chaque token, le nb de surfaces distinctes le contenant.

    Sert a juger la distinctivite d'un token : "holden" present dans une seule
    famille est sur a fusionner ; "alex" present dans plusieurs ne l'est pas.
    """
    from collections import Counter
    freq: Counter[str] = Counter()
    surfaces = {_norm(f) for c in characters for f in _forms(c)}
    surfaces |= {_norm(s) for s in (extra or [])}
    for s in surfaces:
        for t in _tokens(s):
            freq[t] += 1
    return freq


def heuristic_match(surface: str, characters: list[Character], tokfreq=None):
    """Rapproche `surface` d'un personnage connu sans LLM (conservateur).

    Renvoie le `Character` correspondant, `None` si aucun, ou `_AMBIGUOUS` si le
    rapprochement est plausible mais incertain (decision laissee a Gemma).

    Un lien par sous-ensemble de tokens n'est considere SUR que si le plus petit
    cote a >=2 tokens, ou si les tokens partages sont globalement distinctifs
    (presents dans <=2 surfaces). Sinon le lien est ambigu (ex: un prenom
    courant "Alex" partage par plusieurs personnages).
    """
    s_norm = _norm(surface)
    for c in characters:
        if _norm(c.name) == s_norm or any(_norm(a) == s_norm for a in c.aliases):
            return c
    s_tok = _tokens(surface)
    if not s_tok:
        return None
    if tokfreq is None:
        tokfreq = _token_freq(characters, [surface])

    safe: list[Character] = []
    ambiguous = False
    for c in characters:
        linked = is_safe = False
        for form in _forms(c):
            f_tok = _tokens(form)
            if not f_tok or not (s_tok <= f_tok or f_tok <= s_tok):
                continue
            linked = True
            shared = s_tok & f_tok
            if min(len(s_tok), len(f_tok)) >= 2 or all(tokfreq[t] <= 2 for t in shared):
                is_safe = True
        if is_safe:
            safe.append(c)
        elif linked:
            ambiguous = True
    if len(safe) == 1 and not ambiguous:
        return safe[0]
    if safe or ambiguous:
        return _AMBIGUOUS
    return None


def canonical_of(a: str, b: str) -> str:
    """Forme canonique entre deux variantes : la plus complete."""
    return a if _completeness(a) >= _completeness(b) else b


def _absorb(
    target: Character,
    name: str,
    *,
    gender: Optional[str] = None,
    age: Optional[str] = None,
    description: Optional[str] = None,
    voice_id: Optional[str] = None,
    keep_canonical: bool = False,
) -> None:
    """Fusionne la variante `name` dans `target` (mutation en place).

    Enrichit les attributs manquants, recalcule le nom canonique et range les
    autres formes en aliases. `keep_canonical=True` GARDE le nom actuel de
    `target` comme canonique (les autres formes deviennent aliases) : sert a
    rendre stable un nom deja etabli dans le cast (un chapitre ne doit pas
    renommer "Sagale" en "Amiral Mehmet Sagale")."""
    target.gender = target.gender or gender
    target.age = target.age or age
    target.description = target.description or description
    target.voice_id = target.voice_id or voice_id

    forms: dict[str, str] = {}  # norm -> graphie d'origine (1re vue conservee)
    for f in [target.name, *target.aliases, name]:
        f = (f or "").strip()
        if f:
            forms.setdefault(_norm(f), f)
    canon = (_norm(target.name) if keep_canonical
             else max(forms, key=lambda n: _completeness(forms[n])))
    target.name = forms[canon]
    target.aliases = sorted(v for k, v in forms.items() if k != canon)


# Genre/age d'un locuteur anonyme "anonyme (homme, adulte)" (inverse de
# segmenter._anon_identity) -> pour qu'il herite d'une voix du bon genre.
_ANON_GENDER = {"homme": "male", "femme": "female"}
_ANON_AGE = {"enfant": "child", "jeune": "young", "adulte": "adult", "vieux": "old"}


def _anon_attrs(name: str) -> tuple[Optional[str], Optional[str]]:
    low = name.strip().lower()
    if not low.startswith("anonyme"):
        return None, None
    inside = low[low.find("(") + 1: low.find(")")] if "(" in low else ""
    toks = [t.strip() for t in inside.split(",")]
    gender = next((_ANON_GENDER[t] for t in toks if t in _ANON_GENDER), None)
    age = next((_ANON_AGE[t] for t in toks if t in _ANON_AGE), None)
    return gender, age


def _item(c) -> dict:
    """Normalise un personnage ou un nom brut en entree de reconciliation."""
    if isinstance(c, Character):
        return {"name": c.name, "gender": c.gender, "age": c.age,
                "description": c.description, "voice_id": c.voice_id}
    gender, age = _anon_attrs(str(c))  # figurant anonyme -> genre/age depuis le nom
    return {"name": str(c), "gender": gender, "age": age,
            "description": None, "voice_id": None}


def _find(chars: list[Character], name: str) -> Optional[Character]:
    n = _norm(name)
    return next((c for c in chars
                 if _norm(c.name) == n or any(_norm(a) == n for a in c.aliases)),
                None)


def _create(chars: list[Character], it: dict, name_map: dict[str, str]) -> None:
    new = Character(name=it["name"].strip(), gender=it["gender"], age=it["age"],
                    description=it["description"], voice_id=it["voice_id"])
    chars.append(new)
    name_map[_norm(it["name"])] = new.name


def reconcile_characters(
    book_chars: list[Character],
    new_chars,
    gemma=None,
    *,
    speaker_names: Optional[list[str]] = None,
) -> tuple[list[Character], dict[str, str]]:
    """Reconcilie de nouvelles detections dans le casting du livre.

    `new_chars` : personnages extraits (objets `Character`) du/des chapitre(s).
    `speaker_names` : formes de locuteur brutes vues dans les segments (absorbees
    comme aliases pour que la resolution de voix matche au rendu).
    `gemma` : si fourni, tranche les cas ambigus ; sinon heuristique seule.

    Renvoie (liste canonique mise a jour, map nom_surface_normalise -> canonique).
    """
    chars = [c.model_copy(deep=True) for c in book_chars]
    name_map: dict[str, str] = {}
    # Noms deja etablis dans le cast : on les garde canoniques (un chapitre ne
    # doit pas renommer un personnage existant en une forme plus longue/titree).
    established = {_norm(c.name) for c in book_chars}

    items = [_item(c) for c in new_chars]
    seen = {_norm(it["name"]) for it in items}
    for sp in (speaker_names or []):
        n = _norm(sp or "")
        if n and n not in seen and n not in {"narrateur", "inconnu", "?"}:
            items.append(_item(sp))
            seen.add(n)

    # Fréquence globale des tokens (base + entrants) -> distinctivite stable,
    # independante de l'ordre de traitement.
    tokfreq = _token_freq(chars, [it["name"] for it in items])

    pending: list[dict] = []
    for it in items:
        m = heuristic_match(it["name"], chars, tokfreq)
        if m is _AMBIGUOUS:
            pending.append(it)
        elif m is not None:
            _absorb(m, it["name"], gender=it["gender"], age=it["age"],
                    description=it["description"], voice_id=it["voice_id"],
                    keep_canonical=_norm(m.name) in established)
            name_map[_norm(it["name"])] = m.name
        elif gemma is not None:
            pending.append(it)  # peut etre une variante non evidente ("Jim")
        else:
            _create(chars, it, name_map)

    if pending and gemma is not None:
        decisions = _gemma_reconcile(chars, pending, gemma)
        for it in pending:
            canon = decisions.get(_norm(it["name"]))
            target = _find(chars, canon) if isinstance(canon, str) else None
            if target is None:  # Gemma dit NOUVEAU/inconnu : ultime essai heuristique
                hm = heuristic_match(it["name"], chars, tokfreq)
                target = hm if isinstance(hm, Character) else None
            if target is not None:
                _absorb(target, it["name"], gender=it["gender"], age=it["age"],
                        description=it["description"], voice_id=it["voice_id"],
                        keep_canonical=_norm(target.name) in established)
                name_map[_norm(it["name"])] = target.name
            else:
                _create(chars, it, name_map)
    elif pending:
        # Sans Gemma : on ne devine pas les cas ambigus, on les garde distincts.
        for it in pending:
            _create(chars, it, name_map)

    return chars, name_map


def dedup_cast(characters: list[Character], gemma=None) -> list[Character]:
    """Replie les doublons d'un casting existant (conserve les voix attribuees).

    Deux phases : (1) regroupement heuristique sur (gemma=None) -> liste reduite
    et sure ; (2) si `gemma` fourni, passe de regroupement Gemma sur les seuls
    noms candidats (partageant un token avec un autre), pour fusionner les
    variantes que l'heuristique laisse de cote (ex: "Okoye" -> "Elvi Okoye").
    """
    base, _ = reconcile_characters([], characters, gemma=None)
    if gemma is None:
        return base
    return _gemma_merge_pass(base, gemma)


def _gemma_merge_pass(base: list[Character], gemma) -> list[Character]:
    """Rattache via Gemma les formes courtes a un nom complet (ancre).

    Tache volontairement contrainte (et plus fiable qu'un regroupement libre) :
    une "forme courte" est un nom dont les tokens sont strictement inclus dans
    ceux d'un autre (ex: "Okoye" vs "Elvi Okoye"). Gemma mappe chaque forme
    courte vers le nom canonique EXACT d'une ancre, ou "NOUVEAU". Traite par
    petits lots pour rester dans la zone de fiabilite du modele.
    """
    shorts: list[Character] = []
    anchors: list[Character] = []
    for i, c in enumerate(base):
        ts = _tokens(c.name)
        if ts and any(j != i and ts < _tokens(d.name) for j, d in enumerate(base)):
            shorts.append(c)
        else:
            anchors.append(c)
    if not shorts:
        return base

    result = [a.model_copy(deep=True) for a in anchors]
    leftovers: list[Character] = []
    for start in range(0, len(shorts), 12):
        chunk = shorts[start:start + 12]
        decisions = _gemma_reconcile(result, [_item(s) for s in chunk], gemma)
        for s in chunk:
            canon = decisions.get(_norm(s.name))
            tgt = _find(result, canon) if isinstance(canon, str) else None
            if tgt is None:
                hm = heuristic_match(s.name, result)
                tgt = hm if isinstance(hm, Character) else None
            # Garde-fou : ne pas fusionner deux genres connus opposes.
            if tgt is not None and s.gender and tgt.gender and s.gender != tgt.gender:
                tgt = None
            if tgt is not None:
                _absorb(tgt, s.name, gender=s.gender, age=s.age,
                        description=s.description, voice_id=s.voice_id)
                for a in s.aliases:
                    _absorb(tgt, a)
            else:
                leftovers.append(s)
    return result + leftovers


def _gemma_reconcile(
    chars: list[Character], pending: list[dict], gemma
) -> dict[str, object]:
    """Un appel groupe : pour chaque nom en attente, son canonique ou _NEW."""
    known = []
    for c in chars:
        al = f" (alias: {', '.join(c.aliases)})" if c.aliases else ""
        desc = f" — {c.description}" if c.description else ""
        known.append(f"- {c.name}{al}{desc}")
    new_lines = []
    for n, it in enumerate(pending):
        desc = f" — {it['description']}" if it.get("description") else ""
        new_lines.append(f"[{n}] {it['name']}{desc}")

    prompt = (
        "Personnages DEJA connus du livre :\n"
        + ("\n".join(known) if known else "(aucun)")
        + "\n\nNoms DETECTES a classer :\n" + "\n".join(new_lines)
        + "\n\nPour chaque nom detecte, indique s'il designe un personnage deja "
        "connu (donne alors son nom canonique EXACT tel qu'ecrit ci-dessus) ou "
        "s'il s'agit d'un nouveau personnage (\"NOUVEAU\"). Ne fusionne que si "
        "c'est, avec certitude, la meme personne. EN CAS DE DOUTE, ou si "
        "plusieurs personnages connus pourraient correspondre, reponds "
        "\"NOUVEAU\". Ne rapproche jamais deux personnes differentes qui "
        "partagent seulement un prenom ou un nom de famille.\n\n"
        'Reponds par un tableau JSON: '
        '[{"i":0,"canonical":"James Holden"},{"i":1,"canonical":"NOUVEAU"}]'
    )
    if len(prompt) > _MAX_PROMPT_CHARS:
        prompt = prompt[:_MAX_PROMPT_CHARS]
    result = gemma.generate_json(prompt, system=get_settings().prompt_dedup)

    decisions: dict[str, object] = {}
    for item in result:
        if not isinstance(item, dict) or "i" not in item:
            continue
        n = item["i"]
        canon = str(item.get("canonical") or "").strip()
        if isinstance(n, int) and 0 <= n < len(pending) and canon:
            decisions[_norm(pending[n]["name"])] = (
                _NEW if canon.upper() == "NOUVEAU" else canon)
    return decisions