"""engine.py — model inference + code-owned game rules. No Gradio in here.

The model proposes (stat + difficulty band) and narrates; this module owns all
randomness and arithmetic (DECISIONS §2): the dice, HP damage, XP, leveling,
and the day clock. It also owns the llama.cpp singleton and the lenient JSON
parser that degrades malformed output to a `say` (§13).
"""

import json
import os
import random
import re
import threading

import roster
from assemble_context import assemble_context

try:
    import spaces  # no-op decorator off Hugging Face Spaces
except ImportError:
    spaces = None

# On ZeroGPU the H200 exists only inside @spaces.GPU calls, so the model must
# load there (all layers offloaded); anywhere else we run llama.cpp on CPU.
ZEROGPU = spaces is not None and bool(os.environ.get("SPACES_ZERO_GPU"))

APP_DIR = os.path.dirname(os.path.abspath(__file__))

# Frozen serve system prompt — the same string baked into every training row.
with open(os.path.join(APP_DIR, "system-prompt.md"), encoding="utf-8") as _f:
    SYSTEM_PROMPT = _f.read().strip()

# --- Tunables ---------------------------------------------------------------

# Calibration-prone; revisit after playtesting.
TURNS_PER_DAY = 24

# How many recent log entries ride on the card (window size, DECISIONS §15 open).
LOG_WINDOW = 12

# 2d6 + stat >= DC.
BAND_DC = {"easy": 6, "medium": 8, "hard": 10}

# HP lost on a failed roll. MUST match explode_transcript.BAND_DAMAGE —
# the training data applied these exact values to the cards the model saw.
BAND_DAMAGE = {"easy": 0, "medium": 2, "hard": 3}

# XP gained on a successful roll.
BAND_XP = {"easy": 1, "medium": 2, "hard": 3}

# Cumulative XP required to *reach* each level. Level 7 is the cap.
XP_LADDER = {2: 4, 3: 10, 4: 18, 5: 28, 6: 40, 7: 54}
MAX_LEVEL = 7

DEATH_MESSAGE = (
    "Your strength gives out, and the forest gently reclaims you. "
    "Moss will grow where you fell. You were a good bug."
)

TORPOR_MESSAGE = (
    "Torpor takes hold, and your day ends. Bugs have short memories, and the "
    "forest is ever-shifting... Continue Turn to start a new day."
)

# --- World seeds & loading strings ------------------------------------------

with open(os.path.join(APP_DIR, "world_states.json"), encoding="utf-8") as _f:
    WORLD_STATES = json.load(_f)

with open(os.path.join(APP_DIR, "loading-strings.txt"), encoding="utf-8") as _f:
    LOADING_STRINGS = [line.strip() for line in _f if line.strip()]


def shuffled_world_deck():
    deck = list(WORLD_STATES)
    random.shuffle(deck)
    return deck


def shuffled_loading_strings():
    deck = list(LOADING_STRINGS)
    random.shuffle(deck)
    return deck


# --- Model loading ------------------------------------------------------------

# Two trained branches live in parallel. Flip MODEL_FAMILY to "llama" to fall
# back to the 3B-f16-v4 model (works on the HF Space CPU); "qwen3" loads the
# 8B-Q8_0 (heavier, persona-grounded, needs GPU to be tolerable).
MODEL_FAMILY = os.environ.get("BUG_MODEL_FAMILY", "qwen3")
_FAMILIES = {
    "llama": {
        "repo": "jnalv/you-are-a-bug-llama-3.2-3B",
        "file": "you-are-a-bug-3B-f16-v4.gguf",
    },
    "qwen3": {
        "repo": "jnalv/you-are-a-bug-qwen3-8B",
        "file": "you-are-a-bug-qwen3-8B-Q8_0.gguf",
    },
}
MODEL_REPO = _FAMILIES[MODEL_FAMILY]["repo"]
MODEL_FILE = _FAMILIES[MODEL_FAMILY]["file"]
N_CTX = 4096

# HF Spaces CPU Basic has 2 vCPUs, but the container reports the host's core
# count, so llama.cpp would otherwise oversubscribe and thrash. Override
# locally via BUG_N_THREADS.
N_THREADS = int(os.environ.get("BUG_N_THREADS", "2"))

_llm = None
# Spaces serve many sessions off one process; the single CPU model must
# answer one card at a time.
_llm_lock = threading.Lock()


def _model_path():
    override = os.environ.get("BUG_MODEL_PATH")
    if override:
        return override
    from huggingface_hub import hf_hub_download

    return hf_hub_download(MODEL_REPO, MODEL_FILE)


def _preload_cuda():
    """Preload the pip-packaged CUDA runtime (RTLD_GLOBAL) so the dynamic
    linker can resolve libllama.so's libcudart/libcublas deps — the ZeroGPU
    image does not put them on the default search path."""
    import ctypes
    import glob

    for pkg_glob in (
        "nvidia/cuda_runtime/lib/libcudart.so.*",
        "nvidia/cublas/lib/libcublasLt.so.*",
        "nvidia/cublas/lib/libcublas.so.*",
    ):
        for site_dir in __import__("site").getsitepackages():
            for lib in sorted(glob.glob(os.path.join(site_dir, pkg_glob))):
                try:
                    ctypes.CDLL(lib, mode=ctypes.RTLD_GLOBAL)
                except OSError:
                    pass


def get_llm():
    global _llm
    with _llm_lock:
        if _llm is None:
            # The CUDA-built wheel needs libcudart just to dlopen, even when
            # running CPU-only, so always preload (no-op if libs are absent).
            _preload_cuda()
            from llama_cpp import Llama

            _llm = Llama(
                model_path=_model_path(),
                n_ctx=N_CTX,
                n_threads=N_THREADS,
                n_threads_batch=N_THREADS,
                n_gpu_layers=-1 if ZEROGPU else 0,
                verbose=False,
            )
    return _llm


def warmup():
    """Startup warm: on ZeroGPU only prefetch the GGUF (CUDA must not be
    touched outside @spaces.GPU calls); on CPU, load the model outright."""
    if ZEROGPU:
        _model_path()
    else:
        get_llm()


# --- Inference ----------------------------------------------------------------

_ROLL_KEYS = {"stat", "difficulty", "on_success", "on_fail"}
_STATS = {"might", "speed", "smarts", "mystique"}

# Salvage patterns for the model's known shapes when narration contains
# unescaped double-quotes that break json.loads (the §13 no-quotes convention,
# violated). Greedy (.*) so interior quotes stay inside the field.
_SAY_SALVAGE = re.compile(
    r'^\s*\{\s*"action"\s*:\s*"say"\s*,\s*"text"\s*:\s*"(.*)"\s*\}\s*$', re.S
)
_ROLL_SALVAGE = re.compile(
    r'^\s*\{\s*"action"\s*:\s*"roll"\s*,\s*"stat"\s*:\s*"(\w+)"\s*,'
    r'\s*"difficulty"\s*:\s*"(\w+)"\s*,\s*"on_success"\s*:\s*"(.*)"\s*,'
    r'\s*"on_fail"\s*:\s*"(.*)"\s*\}\s*$',
    re.S,
)


# Lines the model sometimes parrots from the card (or bare action words)
# when it drops the JSON wrapper around otherwise-good narration.
_ECHO_LINES = {"PLAYER", "THE WORLD", "WHAT JUST HAPPENED", "PLAYER'S TURN", "say", "roll"}


def _unescape(value: str) -> str:
    return value.replace('\\"', '"').replace("\\n", "\n")


def _log_degraded(raw, path):
    # Surfaces in Spaces container logs so fallback turns stay diagnosable.
    print(f"[parse-degraded:{path}] {raw!r}", flush=True)


def _salvage_roll_fields(text):
    """Field-by-field salvage for roll JSON that is truncated (max_tokens) or
    quote-broken. Returns a roll dict, a say dict (when only on_success
    narration survives), or None."""
    stat = re.search(r'"stat"\s*:\s*"(\w+)"', text)
    band = re.search(r'"difficulty"\s*:\s*"(\w+)"', text)
    if not (stat and band) or stat.group(1) not in _STATS or band.group(1) not in BAND_DC:
        return None
    succ = re.search(r'"on_success"\s*:\s*"(.*?)(?:"\s*,\s*"on_fail"|"?\s*\}?\s*$)', text, re.S)
    fail = re.search(r'"on_fail"\s*:\s*"(.*?)"?\s*\}?\s*$', text, re.S)
    succ_text = _unescape(succ.group(1)).strip() if succ else ""
    fail_text = _unescape(fail.group(1)).strip() if fail else ""
    if succ_text and fail_text:
        return {
            "action": "roll",
            "stat": stat.group(1),
            "difficulty": band.group(1),
            "on_success": succ_text,
            "on_fail": fail_text,
        }
    if succ_text:
        # on_fail lost to truncation: a roll can't resolve, but the narration
        # is still better than a stock line.
        return {"action": "say", "text": succ_text}
    return None


def _strip_log_tag(line):
    """'[dm · Might check, Easy] The log gives.' -> 'The log gives.'"""
    return re.sub(r"^\s*\[[^\]]*\]\s*", "", line)


def parse_output(raw):
    """Lenient parse of the model's one-JSON-object turn (DECISIONS §13).

    Anything that isn't a well-formed roll/say degrades to a say so the game
    never stalls on a malformed turn.
    """
    text = (raw or "").strip()
    obj = None
    try:
        obj = json.loads(text)
    except (json.JSONDecodeError, ValueError):
        match = re.search(r"\{.*\}", text, re.S)
        if match:
            try:
                obj = json.loads(match.group(0))
            except (json.JSONDecodeError, ValueError):
                obj = None
    if not isinstance(obj, dict):
        say = _SAY_SALVAGE.match(text)
        if say:
            _log_degraded(raw, "say-salvage")
            return {"action": "say", "text": _unescape(say.group(1))}
        roll = _ROLL_SALVAGE.match(text)
        if roll and roll.group(1) in _STATS and roll.group(2) in BAND_DC:
            _log_degraded(raw, "roll-salvage")
            return {
                "action": "roll",
                "stat": roll.group(1),
                "difficulty": roll.group(2),
                "on_success": _unescape(roll.group(3)),
                "on_fail": _unescape(roll.group(4)),
            }
        salvaged = _salvage_roll_fields(text)
        if salvaged:
            _log_degraded(raw, "roll-fields")
            return salvaged
        # A card echo (model parroting its prompt) is worse than a stock line.
        if not text or text.startswith("PLAYER\n") or "\nTHE WORLD\n" in text:
            _log_degraded(raw, "card-echo")
            return {"action": "say", "text": "Something went wrong in the forest. Nature is like that."}
        # Headers-then-prose shape ('PLAYER'S TURN\nsay\n<narration>'): drop the
        # echoed labels; strip '[dm · ...]' log tags but keep their narration.
        kept = [
            _strip_log_tag(line)
            for line in text.splitlines()
            if _strip_log_tag(line).strip() not in _ECHO_LINES
        ]
        kept = [line for line in kept if line.strip()]
        _log_degraded(raw, "prose")
        text = "\n".join(kept).strip() or "Something went wrong in the forest. Nature is like that."
        return {"action": "say", "text": text}

    action = obj.get("action")
    if (
        action == "roll"
        and _ROLL_KEYS <= obj.keys()
        and obj.get("stat") in _STATS
        and obj.get("difficulty") in BAND_DC
    ):
        return {k: obj[k] for k in ("action", *_ROLL_KEYS)}
    if action == "say" and isinstance(obj.get("text"), str):
        return {"action": "say", "text": obj["text"]}

    # Wrong/missing keys: salvage any narration we can find.
    _log_degraded(raw, "wrong-keys")
    fallback = obj.get("text") or obj.get("on_success") or text
    return {"action": "say", "text": str(fallback)}


def _qwen3_prompt(card):
    # Qwen3's GGUF chat template defaults to thinking-on, which would wrap the
    # JSON response in <think>…</think>. We trained with enable_thinking=False,
    # so build the prompt by hand and stop at <|im_end|>. Empty <think></think>
    # blocks still sometimes leak in; parse_output handles them downstream.
    return (
        f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
        f"<|im_start|>user\n{card}<|im_end|>\n"
        f"<|im_start|>assistant\n"
    )


def _generate(card, max_tokens, temperature, top_p):
    llm = get_llm()
    with _llm_lock:
        if MODEL_FAMILY == "qwen3":
            out = llm(
                _qwen3_prompt(card),
                max_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                stop=["<|im_end|>"],
            )
            return out["choices"][0]["text"]
        out = llm.create_chat_completion(
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": card},
            ],
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
        )
    return out["choices"][0]["message"]["content"]


if spaces is not None:
    # Runs in a forked worker with the H200 attached; budget covers a cold
    # model load (GGUF -> VRAM) plus generation.
    _generate = spaces.GPU(duration=120)(_generate)


def infer(card, max_tokens=384, temperature=0.7, top_p=0.95):
    """One inference: [system, user-card] -> parsed action dict."""
    return parse_output(_generate(card, max_tokens, temperature, top_p))


# --- Cards ----------------------------------------------------------------------

def opening_card(sheet, daily_context):
    """Turn-0 card. system_prompt='' matches the training split: the frozen
    prompt rides as the system message, the card keeps its leading blank block."""
    return assemble_context(
        sheet=sheet,
        daily_context=daily_context,
        system_prompt="",
        opening=True,
    )


def turn_card(sheet, daily_context, log, message, activations):
    return assemble_context(
        sheet=sheet,
        daily_context=daily_context,
        system_prompt="",
        log=log[-LOG_WINDOW:],
        user_input={"message": message, "activations": activations},
    )


# --- Dice, damage, XP ------------------------------------------------------------

def resolve_roll(stat_value, band):
    """2d6 + stat vs the band's DC -> 'success' | 'fail'."""
    total = random.randint(1, 6) + random.randint(1, 6) + stat_value
    return "success" if total >= BAND_DC[band] else "fail"


def apply_damage(sheet, band):
    hp = sheet["hp"]
    new_cur = max(0, hp["cur"] - BAND_DAMAGE[band])
    return {**sheet, "hp": {**hp, "cur": new_cur}}


def level_for_xp(xp):
    level = 1
    for lvl, threshold in sorted(XP_LADDER.items()):
        if xp >= threshold:
            level = lvl
    return level


def maybe_level_up(sheet, xp):
    """Return (sheet, announcements). Carries damage: the level-up grants the
    +1 max HP/Moxie but does not heal (runs stay lethal, §7)."""
    new_level = min(level_for_xp(xp), MAX_LEVEL)
    if new_level <= sheet["level"]:
        return sheet, []

    old = sheet
    new = roster.load_at_level(sheet["bug"], new_level)
    for pool in ("hp", "moxie"):
        gained = new[pool]["max"] - old[pool]["max"]
        new[pool]["cur"] = min(new[pool]["max"], old[pool]["cur"] + gained)

    announcements = [f"You have grown! You are now level {new_level}."]
    old_abilities = {a["name"] for a in old.get("abilities", [])}
    for ability in new.get("abilities", []):
        if ability["name"] not in old_abilities:
            announcements.append(f"New ability unlocked: {ability['name']}!")
    return new, announcements