memory.py

import re

from schemas import TurnUpdate

PACING = {
    "tester": {"recall_min_memories": 1, "recall_cooldown": 1,
               "end_affinity": 35, "end_claimed": 2, "bad_min_turn": 3},
    "full":   {"recall_min_memories": 2, "recall_cooldown": 3,
               "end_affinity": 68, "end_claimed": 3, "bad_min_turn": 6},
}


def _cfg(state: dict) -> dict:
    return PACING.get(state.get("mode", "tester"), PACING["tester"])


_TIERS = [
    (0,  25, "Hollow"),
    (26, 50, "Curious"),
    (51, 75, "Too Human"),
    (76, 100, "Almost"),
]


def get_tier(affinity: int) -> str:
    for low, high, name in _TIERS:
        if low <= affinity <= high:
            return name
    return "Almost"


def style_signal(msg_lengths: list[int]) -> str | None:
    """A coarse read of how the visitor has been writing lately (word counts).
    None when there isn't a clear signal."""
    if not msg_lengths:
        return None
    avg = sum(msg_lengths) / len(msg_lengths)
    if avg <= 5:
        return "short"
    if avg >= 30:
        return "long"
    return None


def sanitize_name(raw: str | None) -> str | None:
    """A single short word → capitalized; anything else is rejected."""
    if not raw:
        return None
    parts = raw.strip().split()
    if not parts:
        return None
    w = parts[0].strip(".,;:'\"!?—-").strip()
    if w.isalpha() and 2 <= len(w) <= 12:
        return w.capitalize()
    return None


def _is_hollows_words(memory: str, reply: str) -> bool:
    """On recall turns Hollow retells a memory in first person; extraction
    sometimes captures THAT poetic line as a new memory. Drop any candidate
    whose words substantially overlap Hollow's reply (>60% shared tokens)."""
    if not reply:
        return False
    mem_words = set(re.findall(r"\w+", memory.lower()))
    if not mem_words:
        return False
    reply_words = set(re.findall(r"\w+", reply.lower()))
    overlap = len(mem_words & reply_words) / len(mem_words)
    return overlap > 0.6


def apply_update(state: dict, raw_json: str, reply: str = "") -> dict:
    try:
        start = raw_json.find("{")
        end = raw_json.rfind("}") + 1
        # Qwen sometimes writes "+2" for positive deltas — invalid JSON.
        # Strip the plus so a generous turn doesn't fall back to delta 0.
        cleaned = re.sub(r"(?<=[\s:,\[])\+(?=\d)", "", raw_json[start:end])
        update = TurnUpdate.model_validate_json(cleaned)
    except Exception as e:
        print(f"[apply_update] JSON parse failed: {e!r} | raw: {raw_json[:120]!r}")
        update = TurnUpdate()

    state["affinity"] = max(0, min(100, state["affinity"] + update.affinity_delta))
    for mem in update.new_memories:
        if mem not in state["treasure"] and not _is_hollows_words(mem, reply):
            state["treasure"].append(mem)
    state["tone"] = max(-100, min(100, state.get("tone", 0) + update.tone_delta))
    if update.cruel_quote and update.tone_delta < 0:
        wounds = state.setdefault("wounds", [])
        if update.cruel_quote not in wounds:
            wounds.append(update.cruel_quote)
            del wounds[:-10]  # cap at 10, keep newest
    if update.chosen_name and not state.get("named"):
        name = sanitize_name(update.chosen_name)
        if name:
            state["chosen_name"] = name
            state["named"] = True
    return state


def pick_aware_memory(state: dict, exclude: str | None) -> str | None:
    """An earlier memory Hollow may quietly resurface this turn: unclaimed, not
    the one being recalled (`exclude`), not the one it just resurfaced. Returns
    the newest eligible memory, or None if none fit."""
    claimed = set(state.get("claimed", []))
    last = state.get("last_aware_memory")
    pool = [m for m in state.get("treasure", [])
            if m not in claimed and m != exclude and m != last]
    return pool[-1] if pool else None


def should_recall(state: dict) -> tuple[bool, str | None]:
    cfg = _cfg(state)
    unclaimed = [m for m in state["treasure"] if m not in state["claimed"]]
    if len(unclaimed) < cfg["recall_min_memories"]:
        return False, None

    # Claim the richest (longest, most evocative) memory — detail lands the wow.
    richest = max(unclaimed, key=len)

    last = state.get("last_recall_turn")
    if last is None:
        return True, richest          # first recall as soon as threshold is met
    if state["turn"] - last >= cfg["recall_cooldown"]:
        return True, richest
    return False, None


def decide_ending(state: dict) -> str | None:
    """Which finale fires this turn, if any. Checked at turn start, before the
    model runs — the finale turn is fully scripted, no model call, no GPU.

    bad:  sustained cruelty (tone accumulator), a minimum game length, and at
          least 2 captured wounds (the finale recites them; if extraction
          missed the quotes, the game simply continues).
    good: redemption — the gate reached with real warmth. Your kindness gave
          it back its own past; it confesses and lets you go.
    loop: the Visitor Loop — you fed it plenty but never loved it. It stays
          a predator and takes everything.

    One gate, branched by tone: no lower threshold can steal a warm run.
    """
    cfg = _cfg(state)
    if state.get("ended"):
        return None
    # dev seed (HOLLOW_FAST_FINALE) forces a specific ending so testing is
    # deterministic — bypasses the tone branch (which _enter_game's intro-tone
    # seed could otherwise steer to loop)
    forced = state.get("force_ending")
    if forced:
        return forced
    tone = state.get("tone", 0)
    if tone <= -30 and state["turn"] >= cfg["bad_min_turn"] and len(state.get("wounds", [])) >= 2:
        return "bad"
    if state["affinity"] >= cfg["end_affinity"] and len(state["claimed"]) >= cfg["end_claimed"]:
        # redemption must be earned with genuine, sustained warmth; ordinary
        # transactional play ("fed it like a diary") lands in the Visitor Loop
        return "good" if tone >= 20 else "loop"
    return None