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,
        do_sample=True,
        temperature=0.8,
        top_p=0.9,
    )
    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, filling in missing keys from DEFAULT_NPC."""
    # Fill any missing top-level keys with defaults so a partial response still works
    for key in REQUIRED_NPC_KEYS:
        if key not in payload:
            print(f"[NPCverse] NPC payload missing key '{key}', using default.")
            payload[key] = DEFAULT_NPC[key]

    stats = payload.get("stats")
    if not isinstance(stats, dict):
        print("[NPCverse] NPC stats invalid, using defaults.")
        payload["stats"] = dict(DEFAULT_NPC["stats"])
        stats = payload["stats"]

    # Fill any missing stat keys with defaults
    for key in REQUIRED_STAT_KEYS:
        if key not in stats:
            print(f"[NPCverse] NPC stat '{key}' missing, using default.")
            stats[key] = DEFAULT_NPC["stats"][key]

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

    # Ensure quests is a list of dicts
    if not isinstance(payload.get("quests"), list) or len(payload["quests"]) == 0:
        payload["quests"] = DEFAULT_NPC["quests"]

    # Ensure secrets is a list of strings
    if not isinstance(payload.get("secrets"), list) or len(payload["secrets"]) == 0:
        payload["secrets"] = DEFAULT_NPC["secrets"]

    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"""You are the NPCverse character engine. Transform the visual description below into a vivid RPG character.

Return ONLY a valid JSON object. No backticks, no markdown, no extra text before or after.

You MUST include every single key shown in the template below. Do not skip any key.

JSON template (replace all placeholder values):
{{
  "name": "<unique fantasy name inspired by appearance>",
  "title": "<short evocative title, e.g. 'Wanderer of the Neon Wilds'>",
  "class": "<creative RPG class name>",
  "level": <integer 1-20>,
  "rarity": "<one of: Common, Uncommon, Rare, Epic, Legendary>",
  "alignment": "<e.g. Chaotic Good, Lawful Neutral, True Neutral>",
  "lore": "<2-3 sentence backstory inspired by their appearance>",
  "stats": {{
    "strength": <1-100>,
    "intelligence": <1-100>,
    "charisma": <1-100>,
    "luck": <1-100>,
    "stealth": <1-100>,
    "chaos": <1-100>
  }},
  "passive_ability": {{
    "name": "<passive skill name>",
    "description": "<what it does>"
  }},
  "ultimate": {{
    "name": "<ultimate skill name>",
    "description": "<what it does>"
  }},
  "weakness": "<one sentence describing their weakness or flaw>",
  "faction": "<name of the group or faction they belong to>",
  "world": "<name of the realm or world they come from>",
  "opening_line": "<first thing they say when summoned, in character, 1-2 sentences>",
  "quests": [
    {{"title": "<quest 1 title>", "description": "<quest 1 desc>", "reward": "<item name>", "rarity": "Common"}},
    {{"title": "<quest 2 title>", "description": "<quest 2 desc>", "reward": "<item name>", "rarity": "Rare"}},
    {{"title": "<quest 3 title>", "description": "<quest 3 desc>", "reward": "<item name>", "rarity": "Epic"}}
  ],
  "secrets": [
    "<secret 1 about this character>",
    "<secret 2 about this character>",
    "<secret 3 about this character>"
  ],
  "emoji": "<single emoji that represents this character>"
}}

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


def _npc_retry_prompt(description: str) -> str:
    """Build a shorter strict prompt for retrying malformed JSON."""
    return f"""Output ONLY a valid JSON object. No markdown. No extra text.

Create an RPG character from this description: {description}

Required JSON — include ALL these keys:
{{"name":"","title":"","class":"","level":1,"rarity":"Common","alignment":"","lore":"",
"stats":{{"strength":50,"intelligence":50,"charisma":50,"luck":50,"stealth":50,"chaos":50}},
"passive_ability":{{"name":"","description":""}},
"ultimate":{{"name":"","description":""}},
"weakness":"","faction":"","world":"","opening_line":"",
"quests":[{{"title":"","description":"","reward":"","rarity":"Common"}},{{"title":"","description":"","reward":"","rarity":"Rare"}},{{"title":"","description":"","reward":"","rarity":"Epic"}}],
"secrets":["","",""],"emoji":"⚔"}}

Fill in all empty string values with creative content based on the description.""".strip()


@spaces.GPU
def generate_npc(description: str) -> dict:
    """Generate a complete RPG NPC JSON object from a visual description."""
    import traceback
    try:
        msgs = [{'role': 'user', 'content': [{"type": "text", "text": _npc_generation_prompt(description)}]}]
        result = _generate_from_messages(msgs, max_new_tokens=1200)
        return _validate_npc_payload(parse_json_safe(str(result)))
    except Exception as e:
        print(f"[NPCverse] generate_npc first attempt failed: {e}")
        traceback.print_exc()
        try:
            retry_msgs = [{'role': 'user', 'content': [{"type": "text", "text": _npc_retry_prompt(description)}]}]
            retry_result = _generate_from_messages(retry_msgs, max_new_tokens=1200)
            return _validate_npc_payload(parse_json_safe(str(retry_result)))
        except Exception as e2:
            print(f"[NPCverse] generate_npc retry also failed: {e2}")
            traceback.print_exc()
            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[-20:]:  # keep last 20 turns for richer memory
        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)}

IMPORTANT memory rules:
- You have access to the full conversation history below. Read it carefully.
- Remember everything the user has told you (their name, choices, quests, preferences).
- Reference past exchanges naturally — if they told you their name, use it; if they completed a quest, acknowledge it.
- Never contradict something you said earlier in this conversation.
- Stay consistent with your established personality throughout.

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...*"