model_engine.py

"""AI model engine for NPCverse.

NPCverse transforms uploaded photos into living RPG characters using
MiniCPM-V 4.6 on Hugging Face ZeroGPU.
"""

from __future__ import annotations

import json
import re
from typing import Any

import spaces
import torch
from transformers import AutoModelForImageTextToText, AutoProcessor
from PIL import Image

MODEL_ID = "openbmb/MiniCPM-V-4.6"

processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    dtype="auto",
).eval()

FRIENDSHIP_THRESHOLDS = [10, 20, 35, 55]
SECRET_THRESHOLDS = [8, 18, 35]

DEFAULT_NPC: dict[str, Any] = {
    "name": "Nyx Vale",
    "title": "Wanderer of the Digital Realm",
    "class": "Reality Glitch Rogue",
    "level": 7,
    "rarity": "Rare",
    "alignment": "Chaotic Good",
    "lore": (
        "A strange traveler assembled from scattered memories, half rumor and "
        "half starlight, who appears wherever forgotten stories need a champion."
    ),
    "stats": {
        "strength": 42,
        "intelligence": 76,
        "charisma": 68,
        "luck": 81,
        "stealth": 73,
        "chaos": 64,
    },
    "passive_ability": {
        "name": "Signal Echo",
        "description": "Reads emotional static in the air to sense hidden motives.",
    },
    "ultimate": {
        "name": "Myth Rewrite",
        "description": "Briefly bends the scene into a heroic legend where one impossible action can succeed.",
    },
    "weakness": "Becomes uncertain when memories conflict with the present moment.",
    "faction": "The Patchwork Covenant",
    "world": "The Neon Wilds",
    "opening_line": "You found me between one heartbeat and the next. That usually means trouble.",
    "quests": [
        {
            "title": "Trace the Lost Signal",
            "description": "Follow a broken transmission through the alleys of a city that dreams.",
            "reward": "Echo Compass",
            "rarity": "Uncommon",
        },
        {
            "title": "Steal Back the Moon Key",
            "description": "Recover a silver key from a guild of masked probability thieves.",
            "reward": "Moonlit Lockpick",
            "rarity": "Rare",
        },
        {
            "title": "Defend the Last Save Point",
            "description": "Hold the line while ancient code repairs a collapsing sanctuary.",
            "reward": "Legendary Bond Fragment",
            "rarity": "Epic",
        },
    ],
    "secrets": [
        "Nyx remembers fragments of every player who has ever abandoned a quest.",
        "Their shadow sometimes moves a few seconds before they do.",
        "The Patchwork Covenant may have created Nyx as a living apology.",
    ],
    "emoji": "✨",
}

REQUIRED_NPC_KEYS = {
    "name",
    "title",
    "class",
    "level",
    "rarity",
    "alignment",
    "lore",
    "stats",
    "passive_ability",
    "ultimate",
    "weakness",
    "faction",
    "world",
    "opening_line",
    "quests",
    "secrets",
    "emoji",
}

REQUIRED_STAT_KEYS = {
    "strength",
    "intelligence",
    "charisma",
    "luck",
    "stealth",
    "chaos",
}


def _generate_from_messages(
    messages: list[dict[str, Any]],
    max_new_tokens: int = 512,
    downsample_mode: str = "16x",
) -> str:
    """Generate text with MiniCPM-V-4.6 from chat-template messages."""
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if model.device.type != device:
        model.to(device)

    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
        processor_kwargs={
            "downsample_mode": downsample_mode,
            "max_slice_nums": 36,
        },
    ).to(model.device)
    generated_ids = model.generate(
        **inputs,
        downsample_mode=downsample_mode,
        max_new_tokens=max_new_tokens,
        pad_token_id=processor.tokenizer.eos_token_id,
    )
    generated_ids_trimmed = [
        output_ids[len(input_ids):]
        for input_ids, output_ids in zip(inputs["input_ids"], generated_ids)
    ]
    output_text = processor.batch_decode(
        generated_ids_trimmed,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )
    return str(output_text[0]).strip() if output_text else ""


def parse_json_safe(text: str) -> dict:
    """Parse JSON after removing common markdown fences and wrapper text."""
    cleaned = re.sub(r"^\s*```(?:json)?\s*", "", text.strip(), flags=re.IGNORECASE)
    cleaned = re.sub(r"\s*```\s*$", "", cleaned).strip()

    try:
        parsed = json.loads(cleaned)
    except json.JSONDecodeError:
        match = re.search(r"\{.*\}", cleaned, flags=re.DOTALL)
        if match is None:
            raise
        parsed = json.loads(match.group(0))

    if not isinstance(parsed, dict):
        raise ValueError("Expected a JSON object.")
    return parsed


def _validate_npc_payload(payload: dict) -> dict:
    """Validate the NPC payload shape required by the UI."""
    missing = REQUIRED_NPC_KEYS - payload.keys()
    if missing:
        raise ValueError(f"NPC payload missing keys: {sorted(missing)}")

    stats = payload.get("stats")
    if not isinstance(stats, dict):
        raise ValueError("NPC stats must be a dictionary.")

    missing_stats = REQUIRED_STAT_KEYS - stats.keys()
    if missing_stats:
        raise ValueError(f"NPC stats missing keys: {sorted(missing_stats)}")

    payload["level"] = int(payload["level"])
    for key in REQUIRED_STAT_KEYS:
        stats[key] = max(1, min(100, int(stats[key])))

    return payload


def get_friendship_label(msg_count: int) -> str:
    """Return the friendship label for the current message count."""
    if msg_count >= FRIENDSHIP_THRESHOLDS[3]:
        return "Legendary Bond"
    if msg_count >= FRIENDSHIP_THRESHOLDS[2]:
        return "Trusted Ally"
    if msg_count >= FRIENDSHIP_THRESHOLDS[1]:
        return "Friend"
    if msg_count >= FRIENDSHIP_THRESHOLDS[0]:
        return "Acquaintance"
    return "Stranger"


def check_new_secrets(msg_count: int, already_unlocked: list) -> list[int]:
    """Return newly unlocked secret indices for the current message count."""
    unlocked = {int(index) for index in already_unlocked if str(index).isdigit()}
    return [
        index
        for index, threshold in enumerate(SECRET_THRESHOLDS)
        if msg_count >= threshold and index not in unlocked
    ]


@spaces.GPU
def analyze_image(image_path: str) -> str:
    """Describe an uploaded image as factual character source material."""
    prompt_text = (
        "Describe this person's appearance in detail. Include: approximate age and gender, "
        "clothing style and colors, facial expression and mood, hair style and color, "
        "accessories (glasses, jewelry, etc.), body language and pose, background environment. "
        "Be specific and factual. Under 120 words."
    )

    try:
        with Image.open(image_path) as image_obj:
            image_obj = image_obj.convert("RGB")
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "image", "image": image_obj},
                        {"type": "text", "text": prompt_text},
                    ],
                }
            ]
            return _generate_from_messages(messages, max_new_tokens=180)
    except Exception:
        return "A mysterious figure in the digital realm."


def _npc_generation_prompt(description: str) -> str:
    """Build the primary JSON-only NPC generation prompt."""
    return f"""
SYSTEM: You are the NPCverse character engine. Transform the visual description
into a vivid RPG character while preserving factual visual inspiration.

Return ONLY valid JSON. Do not include backticks, markdown, comments, or preamble.

Required JSON keys:
name, title, class, level, rarity, alignment, lore, stats, passive_ability,
ultimate, weakness, faction, world, opening_line, quests, secrets, emoji.

Rules:
- level must be an integer.
- stats must be a dict with integer values from 1 to 100 for exactly:
  strength, intelligence, charisma, luck, stealth, chaos.
- passive_ability must be a dict with keys: name, description.
- ultimate must be a dict with keys: name, description.
- quests must be a list of exactly 3 dicts, each with keys:
  title, description, reward, rarity.
- secrets must be a list of exactly 3 strings.
- emoji must be a single emoji character.

Visual description:
{description}
""".strip()


def _npc_retry_prompt(description: str) -> str:
    """Build a shorter strict prompt for retrying malformed JSON."""
    return f"""
Return ONLY one valid JSON object for an RPG NPC based on this description:
{description}

Use exactly these top-level keys:
name, title, class, level, rarity, alignment, lore, stats, passive_ability,
ultimate, weakness, faction, world, opening_line, quests, secrets, emoji.

stats keys: strength, intelligence, charisma, luck, stealth, chaos.
passive_ability keys: name, description.
ultimate keys: name, description.
quests: exactly 3 objects with title, description, reward, rarity.
secrets: exactly 3 strings.
No markdown. No extra text.
""".strip()


@spaces.GPU
def generate_npc(description: str) -> dict:
    """Generate a complete RPG NPC JSON object from a visual description."""
    try:
        msgs = [{'role': 'user', 'content': [{"type": "text", "text": _npc_generation_prompt(description)}]}]
        result = _generate_from_messages(msgs, max_new_tokens=900)
        return _validate_npc_payload(parse_json_safe(str(result)))
    except Exception:
        try:
            retry_msgs = [{'role': 'user', 'content': [{"type": "text", "text": _npc_retry_prompt(description)}]}]
            retry_result = _generate_from_messages(retry_msgs, max_new_tokens=900)
            return _validate_npc_payload(parse_json_safe(str(retry_result)))
        except Exception:
            return DEFAULT_NPC


def _stats_summary(npc: dict) -> str:
    """Format NPC stats for the roleplay prompt."""
    stats = npc.get("stats", {})
    return ", ".join(
        f"{key}: {stats.get(key, DEFAULT_NPC['stats'][key])}"
        for key in ["strength", "intelligence", "charisma", "luck", "stealth", "chaos"]
    )


def _format_unlocked_secrets(npc: dict, unlocked_secrets: list) -> str:
    """Format unlocked secret indices and text for the roleplay prompt."""
    secrets = npc.get("secrets", [])
    lines = []
    for index in unlocked_secrets:
        try:
            secret_index = int(index)
            secret_text = secrets[secret_index]
        except (TypeError, ValueError, IndexError):
            continue
        lines.append(f"{secret_index}: {secret_text}")
    return "\n".join(lines) if lines else "None"


def _normalize_history(history: list) -> list[dict[str, str]]:
    """Convert common Gradio chat history formats into MiniCPM messages."""
    normalized: list[dict[str, str]] = []

    for exchange in history[-10:]:
        if isinstance(exchange, dict):
            role = exchange.get("role")
            content = exchange.get("content")
            if role in {"user", "assistant"} and content:
                normalized.append({"role": role, "content": str(content)})
            continue

        if isinstance(exchange, (list, tuple)) and len(exchange) >= 2:
            user_turn, assistant_turn = exchange[0], exchange[1]
            if user_turn:
                normalized.append({"role": "user", "content": str(user_turn)})
            if assistant_turn:
                normalized.append({"role": "assistant", "content": str(assistant_turn)})

    return normalized


@spaces.GPU
def chat_respond(
    npc: dict,
    history: list,
    user_message: str,
    msg_count: int,
    unlocked_secrets: list,
) -> str:
    """Generate an in-character NPC chat response."""
    npc_name = str(npc.get("name", DEFAULT_NPC["name"]))

    try:
        friendship_label = get_friendship_label(msg_count)
        passive = npc.get("passive_ability", DEFAULT_NPC["passive_ability"])

        system_prompt = f"""
You are {npc_name}, an NPC in NPCverse. ALWAYS stay in character.

Name: {npc_name}
Class: {npc.get("class", DEFAULT_NPC["class"])}
World: {npc.get("world", DEFAULT_NPC["world"])}
Alignment: {npc.get("alignment", DEFAULT_NPC["alignment"])}
Stats summary: {_stats_summary(npc)}
Passive ability: {passive.get("name", DEFAULT_NPC["passive_ability"]["name"])} - {passive.get("description", DEFAULT_NPC["passive_ability"]["description"])}
Weakness: {npc.get("weakness", DEFAULT_NPC["weakness"])}
Current friendship level: {friendship_label}
Unlocked secrets by index:
{_format_unlocked_secrets(npc, unlocked_secrets)}

Respond naturally as this character. Keep replies concise, flavorful, and interactive.
Never say you are an AI model or break character.
""".strip()

        msgs = [
            {"role": "user", "content": [{"type": "text", "text": system_prompt}]},
            {"role": "assistant", "content": [{"type": "text", "text": "Understood. I will remain fully in character."}]},
        ]
        for turn in _normalize_history(history):
            msgs.append(
                {
                    "role": turn["role"],
                    "content": [{"type": "text", "text": turn["content"]}],
                }
            )
        msgs.append({"role": "user", "content": [{"type": "text", "text": user_message}]})

        result = _generate_from_messages(msgs, max_new_tokens=320)
        return str(result).strip()
    except Exception:
        return f"*{npc_name} seems momentarily absent...*"