Kansho/backend/identity_store.py

"""Local identity mappings. Class A; never part of an external prompt."""
from __future__ import annotations

import json
import re
import uuid
from functools import lru_cache
from pathlib import Path

from db import get_db, row_to_dict

PROMPTS_SEED_PATH = Path(__file__).resolve().parent / "config" / "prompts.seed.json"

TOKEN_RE = re.compile(r"^(SELF|PERSON:[A-Z][A-Z0-9_]{0,24}|PLACE:[A-Z][A-Z0-9_]{0,24}|ORG:[A-Z][A-Z0-9_]{0,24})$")
MIN_LABEL_LEN = 4
TIME_RE = re.compile(r"\d|:\d|\buhr\b|\bminuten\b|\bstunde\b", re.IGNORECASE)
DETERMINERS = {
    "der", "die", "das", "den", "dem", "des",
    "ein", "eine", "einen", "einem", "eines",
    "mein", "meine", "meinen", "meinem", "meiner", "meines",
    "dein", "deine", "sein", "seine", "ihr", "ihre",
    "unser", "unsere", "euer", "eure",
}
KINSHIP = {"frau", "herr", "mann", "kind", "tochter", "sohn", "partner", "partnerin"}
UNMASKABLE = {
    "ich", "du", "er", "sie", "es", "wir", "ihr", "mich", "mir", "mein", "meine",
    "dir", "dich", "uns", "euch", "sein", "seine", "ihre",
    "heute", "gestern", "morgen", "user", "assistant", "self", "kontext",
    "himmel", "luft", "regen", "sonne", "wind", "meer", "tee", "kaffee",
    "balkon", "laden", "brot", "brote", "bad", "gasse", "gassen", "weg",
    "ort", "ortes", "stadt", "zimmer", "haus",
    "frau", "herr", "mann", "kind", "lektion", "bootstour", "frühstück",
    "tagebuch", "omlett", "haferflocken", "speiseplan",
    "impuls", "tatsache", "motivation", "erzählung", "operation",
    "uhr", "zeit", "stunde", "minute", "tag", "tages", "abend", "nacht",
    "warm", "klar", "draußen", "menge", "teil", "stück",
    "relativ", "allerdings", "zunächst", "danach", "deshalb", "trotzdem",
    "außerdem", "schließlich", "während", "bevor", "seitdem", "deswegen",
    "überlegen", "feststellen", "frage", "antwort",
    "markt", "kirschen", "hafen", "wetter", "vormittag", "nachmittag",
    "anreise", "details", "absatz", "essen", "licht",
    "zwiebeln", "zwiebel", "tomaten", "tomate", "paprika",
    "omelett", "trockenobst", "trockenfrüchte", "glutenfrei",
    "schiff", "boot", "zweimaster", "fähre",
}
FUNCTION_STEMS = ("letzt", "nächst")
NOUN_SUFFIXES = ("heit", "keit", "schaft", "tum")
COLLECTIVE_NOUN = re.compile(r"^ge[a-zäöüß]+öse$", re.IGNORECASE)
RESERVED_TOKEN_SUFFIXES = {"KURZ", "NAME", "TEXT", "LABEL", "EXAMPLE", "FOO", "BAR"}


@lru_cache(maxsize=1)
def instruction_vocab() -> frozenset[str]:
    """Words from Kanshō prompts must never become identity labels."""
    items = json.loads(PROMPTS_SEED_PATH.read_text(encoding="utf-8"))
    words: set[str] = set()
    for item in items:
        if (item.get("slug") or "") == "mvp.entity_detect":
            continue
        template = item.get("template") or ""
        template = re.sub(r"\{\{[^}]+\}\}|\[\[[^\]]+\]\]", " ", template)
        for word in re.findall(r"[A-Za-zÄÖÜäöüß]{4,}", template):
            words.add(word.casefold())
    return frozenset(words)


def is_function_label(label: str) -> bool:
    word = (label or "").strip().casefold()
    if not word:
        return False
    for stem in FUNCTION_STEMS:
        if word == stem or word.startswith(stem):
            return True
    return False


def is_common_noun_shape(label: str) -> bool:
    word = (label or "").strip()
    if COLLECTIVE_NOUN.match(word):
        return True
    lower = word.casefold()
    if lower.endswith("ung") and len(lower) > 6:
        return True
    return any(lower.endswith(suffix) for suffix in NOUN_SUFFIXES)


def is_given_name_candidate(label: str) -> bool:
    core = normalize_label(label)
    if not core or " " in core:
        return False
    if is_function_label(core) or is_common_noun_shape(core):
        return False
    return is_maskable_label(core)


def normalize_label(label: str) -> str:
    words = re.findall(r"[A-Za-zÄÖÜäöüß0-9.:]+", label or "")
    keep = [
        word
        for word in words
        if word.lower() not in DETERMINERS
        and word.lower() not in KINSHIP
        and word.lower() not in UNMASKABLE
    ]
    return " ".join(keep).strip()


def is_maskable_label(label: str) -> bool:
    core = normalize_label(label)
    if len(core) < MIN_LABEL_LEN:
        return False
    if core.lower() in UNMASKABLE:
        return False
    if is_function_label(core) or is_common_noun_shape(core):
        return False
    if core.casefold() in instruction_vocab():
        return False
    if TIME_RE.search(core):
        return False
    if re.fullmatch(r"user|assistant|self", core, re.IGNORECASE):
        return False
    if not re.search(r"[A-Za-zÄÖÜäöüß]", core):
        return False
    return True


def purge_unmaskable(profile_id: str) -> int:
    removed = 0
    with get_db() as conn:
        rows = conn.execute(
            "SELECT id, local_label FROM identity_mappings WHERE profile_id = ?",
            (profile_id,),
        ).fetchall()
        for row in rows:
            if is_maskable_label(row["local_label"]):
                continue
            conn.execute("DELETE FROM identity_mappings WHERE id = ?", (row["id"],))
            removed += 1
    return removed


def list_mappings(profile_id: str) -> list[dict]:
    with get_db() as conn:
        rows = conn.execute(
            "SELECT token, local_label FROM identity_mappings WHERE profile_id = ? ORDER BY created",
            (profile_id,),
        ).fetchall()
    return [row_to_dict(row) for row in rows]


def _token_suffix(token: str) -> str:
    raw = (token or "").upper()
    return raw.split(":")[-1] if ":" in raw else raw


def normalize_token(token: str | None) -> str | None:
    raw = (token or "").strip().upper().replace(" ", "_")
    raw = raw[2:-2] if raw.startswith("[[") and raw.endswith("]]") else raw
    if not TOKEN_RE.match(raw):
        return None
    if raw.count("_") > 1:
        return None
    if _token_suffix(raw) in RESERVED_TOKEN_SUFFIXES:
        return None
    return raw


def _same_core(left: str, right: str) -> bool:
    a = normalize_label(left).casefold()
    b = normalize_label(right).casefold()
    return bool(a and b and a == b)


def find_existing_mapping(existing: list[dict], label: str) -> dict | None:
    core = normalize_label(label).casefold()
    if not core:
        return None
    for item in existing:
        stored = (item.get("local_label") or "").strip()
        if stored.casefold() == label.casefold() or _same_core(stored, label):
            return item
    return None


def _pick_keeper(rows: list) -> dict:
    def sort_key(row: dict) -> tuple:
        suffix = _token_suffix(row["token"] or "")
        reserved = 1 if suffix in RESERVED_TOKEN_SUFFIXES else 0
        numeric = 0 if suffix.isdigit() else 1
        return (reserved, numeric, row["created"] or "", row["id"] or "")

    return min(rows, key=sort_key)


def coalesce_mappings(profile_id: str) -> int:
    """One token per normalized name. Template tokens like PERSON:KURZ lose against PERSON:08."""
    merged = 0
    with get_db() as conn:
        rows = [
            row_to_dict(row)
            for row in conn.execute(
                """
                SELECT id, token, local_label, created
                FROM identity_mappings
                WHERE profile_id = ?
                ORDER BY created
                """,
                (profile_id,),
            ).fetchall()
        ]
        groups: dict[str, list[dict]] = {}
        for row in rows:
            core = normalize_label(row.get("local_label") or "")
            if not core or not is_maskable_label(core):
                conn.execute("DELETE FROM identity_mappings WHERE id = ?", (row["id"],))
                merged += 1
                continue
            groups.setdefault(core.casefold(), []).append({**row, "core": core})
        for items in groups.values():
            keeper = _pick_keeper(items)
            core = items[0]["core"]
            if (keeper.get("local_label") or "") != core:
                conn.execute(
                    "UPDATE identity_mappings SET local_label = ? WHERE id = ?",
                    (core, keeper["id"]),
                )
            for row in items:
                if row["id"] == keeper["id"]:
                    continue
                conn.execute("DELETE FROM identity_mappings WHERE id = ?", (row["id"],))
                merged += 1
    return merged


def _next_generic(existing: list[dict], prefix: str) -> str:
    used = set()
    for item in existing:
        token = (item.get("token") or "").upper()
        if token.startswith(prefix):
            suffix = token[len(prefix) :]
            if suffix.isdigit():
                used.add(int(suffix))
    n = 1
    while n in used:
        n += 1
    return f"{prefix}{n:02d}"


def remember_mapping(profile_id: str, local_label: str, suggested_token: str | None = None) -> dict:
    label = normalize_label(local_label)
    if not label or not is_maskable_label(label):
        raise ValueError("empty_label")
    existing = list_mappings(profile_id)
    found = find_existing_mapping(existing, label)
    if found:
        stored = (found.get("local_label") or "").strip()
        if stored != label:
            with get_db() as conn:
                conn.execute(
                    """
                    UPDATE identity_mappings
                    SET local_label = ?
                    WHERE profile_id = ? AND token = ?
                    """,
                    (label, profile_id, found["token"]),
                )
            found = {**found, "local_label": label}
        return found
    token = normalize_token(suggested_token)
    if token and any((item.get("token") or "").upper() == token for item in existing):
        token = None
    if not token:
        token = _next_generic(existing, "PERSON:")
    with get_db() as conn:
        conn.execute(
            """
            INSERT INTO identity_mappings (id, profile_id, token, local_label)
            VALUES (?, ?, ?, ?)
            """,
            (str(uuid.uuid4()), profile_id, token, label),
        )
    return {"token": token, "local_label": label}