ai/generation.py

"""Entity generation orchestration and validation."""

from __future__ import annotations

import os
import re
import time
from typing import Optional

from pydantic import BaseModel, field_validator

from ai import prompts
from ai.modal_client import extract_json, generate
from persistence.traces import log_entity_generation
from world.book_of_ages import create_entry
from world.entities import create_entity, get_legendary_entities, record_creation
from world.locations import get_location_by_name

NSFW_KEYWORDS = {
    "porn", "nude", "naked", "sex", "explicit", "nsfw", "xxx",
}

VALID_LOCATIONS = {
    "The Library of Unfinished Thoughts",
    "The Sea of Forgotten Names",
    "The Clock Forest",
    "The Moon Market",
    "The Valley of Sleeping Giants",
    "The Ember Crossroads",
    "The Mirror Bogs",
    "The Hollow Mountain",
    "Library of Unfinished Thoughts",
    "Sea of Forgotten Names",
    "Clock Forest",
    "Moon Market",
    "Valley of Sleeping Giants",
    "Ember Crossroads",
    "Mirror Bogs",
    "Hollow Mountain",
}

LOCATION_ALIASES = {
    "Library of Unfinished Thoughts": "The Library of Unfinished Thoughts",
    "Sea of Forgotten Names": "The Sea of Forgotten Names",
    "Clock Forest": "The Clock Forest",
    "Moon Market": "The Moon Market",
    "Valley of Sleeping Giants": "The Valley of Sleeping Giants",
    "Ember Crossroads": "The Ember Crossroads",
    "Mirror Bogs": "The Mirror Bogs",
    "Hollow Mountain": "The Hollow Mountain",
}

MODEL_ID = os.environ.get("MODEL_ID", "openbmb/MiniCPM3-4B")


class EntityProfile(BaseModel):
    name: str
    display_name: str
    type: str
    appearance: str
    backstory: str
    personality_traits: list[str]
    primary_goal: str
    secondary_goal: str
    primary_fear: str
    speech_style: str
    greeting: str
    suggested_location: str
    arrival_note: str

    @field_validator("type")
    @classmethod
    def validate_type(cls, v: str) -> str:
        if v not in ("character", "creature", "object", "place"):
            raise ValueError(f"Invalid entity type: {v}")
        return v

    @field_validator("personality_traits")
    @classmethod
    def validate_traits(cls, v: list[str]) -> list[str]:
        if len(v) != 4:
            raise ValueError("Must have exactly 4 personality traits")
        return v

    @field_validator("suggested_location")
    @classmethod
    def validate_location(cls, v: str) -> str:
        if v not in VALID_LOCATIONS:
            raise ValueError(f"Invalid location: {v}")
        return v


def validate_input(user_input: str) -> tuple[bool, str]:
    words = user_input.strip().split()
    if len(words) < 5:
        return False, "Describe your creation in at least 5 words. The Realm needs more to work with."
    if len(words) > 200:
        return False, "Your description is too long. Keep it under 200 words."

    lower = user_input.lower()
    for kw in NSFW_KEYWORDS:
        if kw in lower:
            return False, "The Realm cannot absorb that kind of creation. Please rephrase."

    return True, ""


def _legendary_context() -> str:
    legendaries = get_legendary_entities(limit=2)
    if not legendaries:
        return ""
    names = ", ".join(e["name"] for e in legendaries)
    return (
        f"\n\nLegendary souls already woven into the Realm: {names}. "
        "The new arrival's backstory may reference one as distant myth — "
        "not as someone they personally know unless plausible."
    )


def _clean_name(name: str, user_input: str) -> str:
    name = (name or "").strip()
    while name.lower().startswith("the the "):
        name = "The " + name[8:].lstrip()
    if re.match(r"^The A \w+", name, re.I):
        words = [w.capitalize() for w in user_input.split()[:3] if w.isalpha()]
        if words:
            name = "The " + " ".join(words)
    if len(name.split()) > 6:
        name = " ".join(name.split()[:4])
    if not name.startswith("The ") and len(name.split()) <= 4:
        name = "The " + name
    return name


def generate_entity(
    user_input: str,
    session_id: str | None = None,
) -> tuple[Optional[dict], str]:
    valid, msg = validate_input(user_input)
    if not valid:
        return None, msg

    legendary_block = _legendary_context()
    user_prompt = prompts.ENTITY_GENERATION_USER.format(
        user_input=user_input,
        legendary_context=legendary_block,
    )

    for attempt, temp in enumerate([0.75, 0.5]):
        raw = ""
        t0 = time.perf_counter()
        try:
            raw = generate(
                prompts.ENTITY_GENERATION_SYSTEM,
                user_prompt,
                temperature=temp,
                max_new_tokens=700,
            )
            data = extract_json(raw)
            if data.get("suggested_location") in LOCATION_ALIASES:
                data["suggested_location"] = LOCATION_ALIASES[data["suggested_location"]]

            data["name"] = _clean_name(data.get("name", ""), user_input)
            if not data.get("display_name"):
                data["display_name"] = data["name"].replace("The ", "")

            profile = EntityProfile(**data)

            entity_data = profile.model_dump()
            location = get_location_by_name(profile.suggested_location)

            if location and location["slug"] == "sea":
                entity_data["secret_name_generated"] = _generate_secret_name(profile.name)

            entity = create_entity(entity_data, user_input, session_id)

            # Generate portrait in background (non-blocking)
            _spawn_portrait_generation(entity)

            arrival_content = (
                f"On Day {entity.get('world_day', 1)}, {profile.name} arrived in "
                f"{profile.suggested_location}. {profile.arrival_note}"
            )
            create_entry(
                world_day=entity.get("world_day", 1),
                entry_type="arrival",
                content=arrival_content,
                entity_ids=[entity["id"]],
                location_id=location["id"] if location else None,
            )

            if session_id:
                record_creation(session_id)

            latency = (time.perf_counter() - t0) * 1000
            log_entity_generation(
                user_input=user_input,
                system_prompt=prompts.ENTITY_GENERATION_SYSTEM,
                user_prompt=user_prompt,
                raw_output=raw,
                latency_ms=latency,
                model_id=MODEL_ID,
                success=True,
                entity_id=entity["id"],
            )

            return entity, ""

        except Exception as e:
            latency = (time.perf_counter() - t0) * 1000
            log_entity_generation(
                user_input=user_input,
                system_prompt=prompts.ENTITY_GENERATION_SYSTEM,
                user_prompt=user_prompt,
                raw_output=raw,
                latency_ms=latency,
                model_id=MODEL_ID,
                success=False,
                error=str(e),
            )
            if attempt == 1:
                return None, f"The Oracle could not shape your creation. Please try again. ({e})"

    return None, "The Oracle was silent. Try again."


def _spawn_portrait_generation(entity: dict) -> None:
    """Fire-and-forget portrait generation in a background thread."""
    import threading
    def _gen():
        try:
            from ai.image_generation import generate_soul_portrait
            from world.database import db_session
            path = generate_soul_portrait(entity)
            if path:
                with db_session() as conn:
                    conn.execute(
                        "UPDATE entities SET portrait_url=? WHERE id=?",
                        (path, entity["id"])
                    )
        except Exception:
            pass
    threading.Thread(target=_gen, daemon=True).start()


def _generate_secret_name(true_name: str) -> str:
    words = true_name.replace("The ", "").split()
    if len(words) >= 2:
        return f"The {' '.join(reversed(words))}"
    return f"The Unspoken {words[0] if words else 'Name'}"