import asyncio
import json
import os
import time
import re
from threading import Thread
from typing import Any

import modal


APP_NAME = "virtual-characters-llm"
MODEL_ID = os.environ.get("VC_LLM_MODEL", "google/gemma-4-12B-it")
MODEL_REVISION = os.environ.get("VC_LLM_REVISION")
GPU = os.environ.get("VC_LLM_GPU", "L40S")
MODEL_DIR = "/root/.cache/huggingface"
HF_SECRET_NAME = os.environ.get("VC_HF_SECRET_NAME", "hf-token")
HF_SECRETS = [] if os.environ.get("VC_SKIP_HF_SECRET") == "1" else [modal.Secret.from_name(HF_SECRET_NAME)]


image = (
    modal.Image.from_registry("nvidia/cuda:12.9.0-devel-ubuntu22.04", add_python="3.12")
    .entrypoint([])
    .uv_pip_install(
        "accelerate>=1.8.0",
        "fastapi[standard]>=0.115.0",
        "huggingface-hub>=0.36.0",
        "librosa>=0.10.2",
        "pillow>=11.0.0",
        "torch>=2.7.0",
        "torchvision>=0.22.0",
        "transformers>=4.57.0",
    )
    .env(
        {
            "HF_HUB_CACHE": MODEL_DIR,
            "HF_XET_HIGH_PERFORMANCE": "1",
        }
    )
)

hf_cache = modal.Volume.from_name("vc-hf-cache", create_if_missing=True)
app = modal.App(APP_NAME, image=image)


def _sse(event: dict[str, Any]) -> bytes:
    return f"data: {json.dumps(event, ensure_ascii=False)}\n\n".encode("utf-8")


def _build_messages(payload: dict[str, Any]) -> list[dict[str, Any]]:
    character = payload.get("character", {})
    character_name = character.get("display_name") or character.get("id") or "角色"
    user_text = payload.get("text") or payload.get("message") or ""
    vision_note = payload.get("vision_note")

    system_prompt = f"""
你是虚拟角色系统的对话引擎。当前角色：{character_name}。
你必须保持角色设定，并用中文优先回答。
输出时自然、有角色感、简短。不要声称自己是真实商业角色或官方角色。
后端会把你的文本切成事件流，所以你只需要输出角色要说的话，不要输出 JSON。
""".strip()

    content: list[dict[str, Any]] = []
    if vision_note:
        content.append({"type": "text", "text": f"视觉观察：{vision_note}"})
    for image_url in payload.get("image_urls", []) or []:
        content.append({"type": "image", "url": image_url})
    for audio_url in payload.get("audio_urls", []) or []:
        content.append({"type": "audio", "audio": audio_url})
    for video_url in payload.get("video_urls", []) or []:
        content.append({"type": "video", "video": video_url})
    content.append({"type": "text", "text": user_text})

    return [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": content},
    ]


def _guess_stage_event(text: str) -> dict[str, Any] | None:
    if any(word in text for word in ["累", "难过", "害怕", "担心", "压力"]):
        return {"type": "stage", "expression": "worried", "motion": "gentle_blink", "intensity": 0.7}
    if any(word in text for word in ["开心", "喜欢", "太好了", "谢谢"]):
        return {"type": "stage", "expression": "smile", "motion": "soft_sway", "intensity": 0.6}
    if any(word in text for word in ["为什么", "怎么", "什么", "？", "?"]):
        return {"type": "stage", "expression": "thinking", "motion": "look_at_user", "intensity": 0.5}
    return None


def _strip_empty_thought_prefix(text: str) -> str:
    return re.sub(r"^\s*(?:<\|channel\>)?thought\s*(?:<channel\|>)?", "", text, count=1)


@app.cls(
    gpu=GPU,
    scaledown_window=60 * 3,
    timeout=60 * 20,
    secrets=HF_SECRETS,
    volumes={MODEL_DIR: hf_cache},
)
class PersonaLLM:
    def _ensure_loaded(self):
        if getattr(self, "model", None) is not None:
            return
        import torch
        from transformers import AutoModelForMultimodalLM, AutoProcessor

        kwargs: dict[str, Any] = {"dtype": "auto", "device_map": "auto"}
        if MODEL_REVISION:
            kwargs["revision"] = MODEL_REVISION
        self.processor = AutoProcessor.from_pretrained(MODEL_ID, revision=MODEL_REVISION)
        self.model = AutoModelForMultimodalLM.from_pretrained(MODEL_ID, **kwargs)
        self.model.eval()
        self.torch = torch

    @modal.method()
    def health(self) -> dict[str, Any]:
        return {"ok": True, "model": MODEL_ID, "gpu": GPU, "loaded": getattr(self, "model", None) is not None}

    @modal.method()
    def generate_text(self, user_text: str, character: dict[str, Any] | None = None, max_new_tokens: int = 120) -> dict[str, Any]:
        self._ensure_loaded()
        started = time.perf_counter()
        payload = {
            "text": user_text,
            "character": character or {"display_name": "星萤"},
            "max_new_tokens": max_new_tokens,
            "temperature": 0.75,
            "top_p": 0.95,
        }
        messages = _build_messages(payload)
        try:
            inputs = self.processor.apply_chat_template(
                messages,
                tokenize=True,
                return_dict=True,
                return_tensors="pt",
                add_generation_prompt=True,
                enable_thinking=False,
            ).to(self.model.device)
        except TypeError:
            inputs = self.processor.apply_chat_template(
                messages,
                tokenize=True,
                return_dict=True,
                return_tensors="pt",
                add_generation_prompt=True,
            ).to(self.model.device)

        input_len = int(inputs["input_ids"].shape[-1])
        with self.torch.inference_mode():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=True,
                temperature=0.75,
                top_p=0.95,
            )
        new_token_ids = outputs[0][input_len:]
        text = self.processor.decode(new_token_ids, skip_special_tokens=True)
        text = _strip_empty_thought_prefix(text).strip()
        elapsed = time.perf_counter() - started
        return {
            "text": text,
            "output_tokens": int(new_token_ids.shape[-1]),
            "remote_s": round(elapsed, 3),
        }

    @modal.fastapi_endpoint(method="POST")
    async def persona_events(self, request):
        from fastapi.responses import StreamingResponse

        payload = await request.json()
        return StreamingResponse(
            self._event_stream(payload),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
        )

    def _event_stream(self, payload: dict[str, Any]):
        from transformers import TextIteratorStreamer

        self._ensure_loaded()
        max_new_tokens = int(payload.get("max_new_tokens", 160))
        messages = _build_messages(payload)

        yield _sse({"type": "stage", "expression": "listening", "motion": "look_at_user", "intensity": 0.4})
        yield _sse({"type": "voice", "style": "soft", "speed": 0.96, "energy": 0.45})
        yield _sse({"type": "skill", "name": payload.get("skill") or "daily_chat"})

        try:
            inputs = self.processor.apply_chat_template(
                messages,
                tokenize=True,
                return_dict=True,
                return_tensors="pt",
                add_generation_prompt=True,
                enable_thinking=False,
            ).to(self.model.device)
        except TypeError:
            inputs = self.processor.apply_chat_template(
                messages,
                tokenize=True,
                return_dict=True,
                return_tensors="pt",
                add_generation_prompt=True,
            ).to(self.model.device)

        streamer = TextIteratorStreamer(self.processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
        generate_kwargs = {
            **inputs,
            "streamer": streamer,
            "max_new_tokens": max_new_tokens,
            "do_sample": True,
            "temperature": float(payload.get("temperature", 0.8)),
            "top_p": float(payload.get("top_p", 0.95)),
        }

        thread = Thread(target=self.model.generate, kwargs=generate_kwargs)
        thread.start()

        buffer = ""
        sentence_buffer = ""
        last_stage_at = 0.0
        first_chunk = True
        terminators = "。！？!?；;\n"

        for chunk in streamer:
            if not chunk:
                continue
            if first_chunk:
                chunk = _strip_empty_thought_prefix(chunk)
                first_chunk = False
                if not chunk:
                    continue
            buffer += chunk
            sentence_buffer += chunk
            yield _sse({"type": "text_delta", "text": chunk})

            now = time.monotonic()
            if now - last_stage_at > 1.0:
                stage = _guess_stage_event(buffer[-80:])
                if stage:
                    last_stage_at = now
                    yield _sse(stage)

            if any(mark in chunk for mark in terminators):
                text = sentence_buffer.strip()
                if text:
                    yield _sse({"type": "sentence_end", "text": text})
                sentence_buffer = ""

        thread.join(timeout=2)
        if sentence_buffer.strip():
            yield _sse({"type": "sentence_end", "text": sentence_buffer.strip()})
        yield _sse({"type": "done"})


@app.local_entrypoint()
def main(prompt: str = "你好，今天有点累。请用角色口吻回复我。", max_new_tokens: int = 80):
    print("Checking remote health...")
    print(PersonaLLM().health.remote())
    print("Use `modal deploy modal_apps/modal_llm.py` to deploy the SSE endpoint.")