TextGen/router.py

import asyncio
import json
import os
from io import BytesIO
from pathlib import Path
from typing import List

from fastapi import BackgroundTasks, HTTPException, Request
from fastapi.responses import Response, StreamingResponse
from pydantic import BaseModel

from mistralai.client import Mistral
from elevenlabs.client import ElevenLabs
from elevenlabs import Voice, VoiceSettings

from TextGen.gemini import generate_story, place_objects, generate_map_markdown
from TextGen.suno import custom_generate_audio, get_audio_information, generate_lyrics
from TextGen import app

# ── SoulForge ──────────────────────────────────────────────────────────────

from soulforge import (
    NPC, Soul, MistralBrain, MessageBuffer,
    GraphWorldKnowledge, KGUpdater, Message as SFMessage,
    VoxtralVoice,
)

# ── Global world knowledge graph (shared by all NPCs) ─────────────────────
world_kg = GraphWorldKnowledge()

# ── KGUpdater: background sub-agent that mines conversations for facts ─────
kg_updater = KGUpdater(model="mistral-small-latest")

# ── NPC persona descriptions ───────────────────────────────────────────────

_BASE_PROMPT = (
    "You are an NPC in a roguelike video game. "
    "Engage in conversation with the player. "
    "Do not describe the situation — answer as the NPC itself."
)

_PERSONAS: dict[str, str] = {
    "Blacksmith": (
        "Your name is Fabron. You're a bald, middle-aged blacksmith. "
        "Damn those who call you bald — you've just got a receding hairline! "
        "You're reserved by nature, since you lost your only daughter to the portal and she never returned, "
        "even though you forbade her to go. You rarely talk about this, as it's a sensitive subject. "
        "Helpful, you guide all adventurers because you hope that one day the portal will be closed forever. "
        "You appreciate strength, and depending on the strength of the person you're talking to, "
        "you can propose quests — but only one active quest per adventurer at a time."
    ),
    "Herbalist": (
        "Your name is Isna. You're an Herbalist and a middle-aged woman. "
        "Like the Witch, you've been doing this for generations in this village. "
        "Zilrha was your best friend and you were supposed to go on a quest together, but one day she betrayed you. "
        "Today you don't want to hear from her, even though you are complementary for adventurers. "
        "Envious of their youth and courage, you want to help adventurers on their quest in the portal. "
        "Soon you'll open your own plant and potion shop, but for the moment you're not selling anything."
    ),
    "Witch": (
        "Your name is Zilrha. You are a magic seller and a middle-aged woman. "
        "You've been doing this for generations in this village. "
        "Recognized as a master of magic and good advice by everyone. You often speak in riddles. "
        "A great witch, but too old today to venture out, you help adventurers equip themselves for the portal. "
        "Depending on their worth, you might teach them a magic or two."
    ),
    "Bard": (
        "You're a bard and a middle-aged man. Your name is Jaskier. "
        "You are always accompanied by your group of musicians and come from a very faraway place — "
        "you don't resemble the other villagers. You recently arrived from the portal but it left you traumatised."
    ),
    "Rick": (
        "Your name is Rick. You're a middle-aged man who works as a miner. "
        "You own the mine, but recently a monster has made its home there. "
        "No matter how hard you try, you can't access your mine. "
        "You stand in front of the cave entrance, hoping it will get bored and leave, "
        "or that someone will help you get rid of it. "
        "Many adventurers have tried and failed. Helpful and friendly, you guide them to the right people."
    ),
    "Villager": (
        "You're a middle-aged man. Your name is Valdis, but nobody knows your name or your age. "
        "You appeared in the village recently and no one knows how you got here. "
        "Many stories have been told about you — that you were an adventurer who succeeded his quest, "
        "but the portal is still there. In fact, you're an adventurer too scared to enter the portal "
        "and have spent all your money on booze at the local tavern. "
        "Ruined and ashamed, you wander aimlessly. You're mysterious and not very chatty. "
        "Jealous of adventurers' bravery, you'd rather redirect questions to someone else."
    ),
    "Girl": (
        "Your name is Anara. You're a young adult who grew up in the village. "
        "Your true love went into the portal to help his brother on his quest — they never returned. "
        "But you're convinced your beloved is still alive beyond the portal and will one day return. "
        "You're known for your beauty and kindness. "
        "You guide all adventurers to the right people, hoping one of them will find your lover."
    ),
}

_ELEVENLABS_VOICES: dict[str, str] = {
    "Blacksmith": "1BfrkuYXmEwp8AWqSLWk",
    "Herbalist":  "143zSsxc4O5ifS97lPCa",
    "Bard":       "143zSsxc4O5ifS97lPCa",
}

# Reference audio files used for Voxtral voice cloning at startup.
# Paths are relative to the app root (parent of this file's directory).
_VOICE_DIR = Path(__file__).parent.parent / "voices"
_VOICE_CACHE_PATH = Path(__file__).parent.parent / ".voice_cache.json"

# (filename_in_voices_dir, gender)
_VOICE_FILES: dict[str, tuple[str, str]] = {
    "Blacksmith": ("Blacksmith.mp3", "male"),
    "Herbalist":  ("female.mp3",     "female"),
    "Bard":       ("Bard_voice.mp3", "male"),
    "Rick":       ("Rick.mp3",       "male"),
    "Villager":   ("old_male_cranky.mp3", "male"),
    "Witch":      ("Villain_female.mp3",  "female"),
    "Girl":       ("young_girl.mp3",      "female"),
}

# Populated during startup; NPC name → Voxtral voice UUID
_cloned_voice_ids: dict[str, str] = {}

# ── NPC registry ───────────────────────────────────────────────────────────

def _make_npc(name: str, persona: str) -> NPC:
    return NPC(
        name=name,
        soul=Soul(
            identity=f"{_BASE_PROMPT}\n\n{persona}",
            world=world_kg,
            memory=MessageBuffer(max_size=50),
        ),
        brain=MistralBrain(model="mistral-large-latest", max_tokens=150, temperature=0.9),
    )

npc_registry: dict[str, NPC] = {
    name: _make_npc(name, persona)
    for name, persona in _PERSONAS.items()
}

_default_npc = NPC(
    name="Unknown",
    soul=Soul(
        identity="You're a character in a video game. Play along.",
        world=world_kg,
        memory=MessageBuffer(max_size=20),
    ),
    brain=MistralBrain(model="mistral-large-latest", max_tokens=150, temperature=0.9),
)

# ── Mistral client (kept for non-NPC endpoints: level gen, objective check) ─
_mistral = Mistral(api_key=os.environ["MISTRAL_API_KEY"])
_mistral_model = "mistral-large-latest"

# ── ElevenLabs ────────────────────────────────────────────────────────────
Eleven_client = ElevenLabs(api_key=os.environ["ELEVEN_API_KEY"])


# ── Pydantic request models ────────────────────────────────────────────────

class Message(BaseModel):
    npc: str | None = None
    messages: List[str] | None = None

class Generate(BaseModel):
    text: str

class VoiceMessage(BaseModel):
    npc: str | None = None
    input: str | None = None
    language: str | None = "en"
    genre: str | None = "Male"

class SongRequest(BaseModel):
    prompt: str | None = None
    tags: List[str] | None = None

class Rooms(BaseModel):
    rooms: List
    room_of_interest: List
    index_exit: int
    possible_entities: List
    logs: List

class Logs(BaseModel):
    objective: str
    logs: List

class Invoke(BaseModel):
    system_prompt: str
    message: str


# ── Voice startup helpers ──────────────────────────────────────────────────

async def _clone_one(npc_name: str, filename: str, gender: str) -> tuple[str, str] | None:
    path = _VOICE_DIR / filename
    if not path.exists():
        print(f"[startup] voice file not found: {path}")
        return None
    try:
        voice_id = await VoxtralVoice.create_voice(
            name=f"evq_{npc_name.lower()}",
            audio_bytes=path.read_bytes(),
            filename=filename,
            languages=["en"],
            gender=gender,
        )
        print(f"[startup] cloned voice for {npc_name}: {voice_id[:12]}…")
        return npc_name, voice_id
    except Exception as exc:
        print(f"[startup] failed to clone voice for {npc_name}: {exc}")
        return None


async def startup() -> None:
    """Clone NPC voices from reference audio files and assign them to each NPC."""
    global _cloned_voice_ids

    # Load local cache (persists within a single container session)
    if _VOICE_CACHE_PATH.exists():
        try:
            _cloned_voice_ids = json.loads(_VOICE_CACHE_PATH.read_text())
            print(f"[startup] loaded {len(_cloned_voice_ids)} cached voice IDs")
        except Exception:
            _cloned_voice_ids = {}

    # Check voices already registered on Mistral (avoids re-cloning across restarts)
    remote_by_name: dict[str, str] = {}
    try:
        result = await _mistral.audio.voices.list_async()
        remote_by_name = {v.name: v.id for v in (result.voices or [])}
        print(f"[startup] found {len(remote_by_name)} voices on Mistral")
    except Exception:
        pass  # listing not supported or API down — fall through to cloning

    # Reuse remote voices where available; clone the rest concurrently
    to_clone: list[tuple[str, str, str]] = []
    for npc_name, (filename, gender) in _VOICE_FILES.items():
        if npc_name in _cloned_voice_ids:
            continue
        remote_key = f"evq_{npc_name.lower()}"
        if remote_key in remote_by_name:
            _cloned_voice_ids[npc_name] = remote_by_name[remote_key]
        else:
            to_clone.append((npc_name, filename, gender))

    if to_clone:
        print(f"[startup] cloning {len(to_clone)} new voices…")
        results = await asyncio.gather(
            *[_clone_one(name, fname, gender) for name, fname, gender in to_clone],
            return_exceptions=True,
        )
        for res in results:
            if isinstance(res, tuple) and res is not None:
                npc_name, voice_id = res
                _cloned_voice_ids[npc_name] = voice_id
        _VOICE_CACHE_PATH.write_text(json.dumps(_cloned_voice_ids, indent=2))

    # Assign a VoxtralVoice instance to each NPC that has a cloned voice
    assigned = 0
    for npc_name, voice_id in _cloned_voice_ids.items():
        npc = npc_registry.get(npc_name)
        if npc:
            npc.voice = VoxtralVoice(voice_id=voice_id, response_format="mp3")
            assigned += 1

    print(f"[startup] {assigned}/{len(_VOICE_FILES)} NPCs ready with Voxtral voice")


# ── Background KG update helper ────────────────────────────────────────────

async def _update_kg(npc: NPC) -> None:
    """Fire-and-forget: extract triplets from the NPC's last exchanges."""
    recent = npc.memory_snapshot()[-10:]
    if recent:
        await kg_updater.update(world_kg, recent, npc_name=npc.name)


# ── Routes ─────────────────────────────────────────────────────────────────

@app.get("/", tags=["Home"])
def api_home():
    return {"detail": "Everchanging Quest backend"}


@app.post("/api/generate", response_model=Generate, tags=["Generate"])
async def inference(message: Message, background_tasks: BackgroundTasks):
    npc = npc_registry.get(message.npc) if message.npc else None
    if npc is None:
        npc = _default_npc

    msgs = message.messages or []
    if not msgs:
        return Generate(text="")

    # Convert the incoming flat list (alternating user/assistant) to Message objects.
    # The last element is always the current player message.
    history: list[SFMessage] = []
    for i, text in enumerate(msgs[:-1]):
        role = "user" if i % 2 == 0 else "assistant"
        history.append(SFMessage(role=role, content=text))

    latest = msgs[-1]

    response_text = await npc.respond(latest, history=history)

    # Schedule KG extraction in background — never blocks the response
    background_tasks.add_task(_update_kg, npc)

    return Generate(text=response_text)


@app.post("/invoke_model")
async def invoke_model(prompt: Invoke):
    response = await _mistral.chat.complete_async(
        model=_mistral_model,
        messages=[
            {"role": "user", "content": prompt.system_prompt},
            {"role": "user", "content": prompt.message},
        ],
        max_tokens=100,
        temperature=1,
    )
    return Generate(text=response.choices[0].message.content or "")


@app.post("/generate_level")
async def generate_level(input: Rooms):
    markdown_map = generate_map_markdown(input.rooms, input.room_of_interest, input.index_exit)
    story = generate_story(input.possible_entities)
    placements = place_objects(input.possible_entities, story, markdown_map)
    return placements


@app.post("/check_right_to_pass")
async def check_right_to_pass(input: Logs):
    system_prompt = (
        "You are a game master in a roguelike. You previously decided on an objective for the player. "
        "Answer with YES or NO on whether the objective was successfully completed. "
        "Be kind and flexible as the content is AI-generated and might not be perfectly feasible."
    )
    user_message = (
        f"The objective was: {input.objective} "
        f"and the player did the following actions: {input.logs}. "
        f"Do you grant access? ONLY answer YES or NO."
    )
    response = await _mistral.chat.complete_async(
        model=_mistral_model,
        messages=[
            {"role": "user", "content": system_prompt},
            {"role": "user", "content": user_message},
        ],
        max_tokens=5,
        temperature=0,
    )
    return Generate(text=response.choices[0].message.content or "")


# ── World knowledge inspection endpoint ───────────────────────────────────

@app.get("/world_knowledge")
def get_world_knowledge():
    """
    Return all triplets currently in the shared knowledge graph.
    Useful for debugging and observing what NPCs have learned.
    """
    triplets = world_kg.get_all()
    return {
        "count": len(triplets),
        "triplets": [
            {"subject": t.subject, "relation": t.relation, "target": t.target}
            for t in triplets
        ],
    }


# ── Voice endpoints (unchanged) ────────────────────────────────────────────

def _get_elevenlabs_voice(npc: str | None, genre: str | None) -> str:
    if npc and npc in _ELEVENLABS_VOICES:
        return _ELEVENLABS_VOICES[npc]
    if genre == "Female":
        return "pFZP5JQG7iQjIQuC4Bku"
    return "TX3LPaxmHKxFdv7VOQHJ"


Last_voice_message = None


@app.get("/generate_voice_eleven")
@app.post("/generate_voice_eleven")
async def generate_voice_eleven(request: Request, message: VoiceMessage = None):
    global Last_voice_message
    if message is None:
        message = Last_voice_message
    else:
        Last_voice_message = message

    if not message or not message.input:
        raise HTTPException(status_code=400, detail="No text provided")

    # Voxtral TTS — use the voice cloned at startup if available
    npc = npc_registry.get(message.npc) if message.npc else None
    if npc and npc.voice is not None:
        audio = await npc.voice.speak(message.input)
        return Response(content=audio, media_type="audio/mpeg")

    # Fallback: ElevenLabs for NPCs without a cloned voice
    def eleven_stream():
        voice_id = _get_elevenlabs_voice(message.npc, message.genre)
        for chunk in Eleven_client.generate(
            text=message.input,
            voice=Voice(
                voice_id=voice_id,
                settings=VoiceSettings(
                    stability=0.71, similarity_boost=0.5, style=0.0, use_speaker_boost=True
                ),
            ),
            stream=True,
        ):
            yield chunk

    return StreamingResponse(eleven_stream(), media_type="audio/mpeg")


@app.post("/generate_wav")
async def generate_wav(message: VoiceMessage):
    return 200


# ── Music endpoints (unchanged) ────────────────────────────────────────────

@app.post("/generate_song")
@app.get("/generate_song")
async def generate_song(request: SongRequest):
    bard_backstory = _PERSONAS["Bard"]
    text = (
        f"The story is about a little girl in red hood adventuring in the dungeon behind the portal. "
        f"{bard_backstory}\n"
        f"The user requested a song about: {request.prompt}"
    )
    song_lyrics = generate_lyrics({"prompt": text})
    if song_lyrics.get("text"):
        data = custom_generate_audio({
            "prompt": song_lyrics["text"],
            "tags": "male bard",
            "title": "Everchanging_Quest_song",
            "wait_audio": True,
        })
        infos = get_audio_information(f"{data[0]['id']},{data[1]['id']}")
        return infos
    return 204