import json
import os
import re
import time
from typing import Any

import modal


APP_NAME = "virtual-characters-gemma-benchmark"
MODEL_ID = os.environ.get("VC_BENCH_MODEL", "google/gemma-4-12B-it")
GPU = os.environ.get("VC_BENCH_GPU", "L40S")
MODEL_DIR = "/root/.cache/huggingface"
HF_SECRET_NAME = os.environ.get("VC_HF_SECRET_NAME", "hf-token")


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",
        "huggingface-hub>=0.36.0",
        "librosa>=0.10.2",
        "pillow>=11.0.0",
        "safetensors>=0.5.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 _strip_empty_thought_prefix(text: str) -> str:
    return re.sub(r"^\s*(?:<\|channel\>)?thought\s*(?:<channel\|>)?", "", text, count=1).strip()


@app.function(
    gpu=GPU,
    timeout=60 * 45,
    scaledown_window=60,
    secrets=[modal.Secret.from_name(HF_SECRET_NAME)],
    volumes={MODEL_DIR: hf_cache},
)
def run_benchmark(prompt: str, max_new_tokens: int = 64, enable_thinking: bool = False) -> dict[str, Any]:
    import torch
    from transformers import AutoModelForMultimodalLM, AutoProcessor

    result: dict[str, Any] = {
        "model_id": MODEL_ID,
        "gpu": GPU,
        "max_new_tokens": max_new_tokens,
        "enable_thinking": enable_thinking,
    }
    remote_started = time.perf_counter()

    t0 = time.perf_counter()
    processor = AutoProcessor.from_pretrained(MODEL_ID)
    result["processor_load_s"] = round(time.perf_counter() - t0, 3)

    t0 = time.perf_counter()
    model = AutoModelForMultimodalLM.from_pretrained(MODEL_ID, dtype="auto", device_map="auto")
    model.eval()
    if torch.cuda.is_available():
        torch.cuda.synchronize()
        torch.cuda.reset_peak_memory_stats()
    result["model_load_s"] = round(time.perf_counter() - t0, 3)

    messages = [
        {"role": "system", "content": "你是一个中文虚拟角色对话模型。回答要自然、简短、有角色感。"},
        {"role": "user", "content": prompt},
    ]

    t0 = time.perf_counter()
    try:
        inputs = processor.apply_chat_template(
            messages,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
            add_generation_prompt=True,
            enable_thinking=enable_thinking,
        ).to(model.device)
    except TypeError:
        inputs = processor.apply_chat_template(
            messages,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
            add_generation_prompt=True,
        ).to(model.device)
    input_len = int(inputs["input_ids"].shape[-1])
    result["prompt_tokens"] = input_len
    result["prepare_input_s"] = round(time.perf_counter() - t0, 3)

    generate_kwargs = {
        **inputs,
        "max_new_tokens": max_new_tokens,
        "do_sample": False,
        "pad_token_id": getattr(processor.tokenizer, "eos_token_id", None),
    }

    if torch.cuda.is_available():
        torch.cuda.synchronize()
    t0 = time.perf_counter()
    outputs = model.generate(**generate_kwargs)
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    generation_s = time.perf_counter() - t0

    new_token_ids = outputs[0][input_len:]
    output_tokens = int(new_token_ids.shape[-1])
    decoded = processor.decode(new_token_ids, skip_special_tokens=True)
    preview = _strip_empty_thought_prefix(decoded)
    result.update(
        {
            "output_tokens": output_tokens,
            "generation_s": round(generation_s, 3),
            "tokens_per_s": round(output_tokens / generation_s, 3) if generation_s > 0 else None,
            "remote_function_s": round(time.perf_counter() - remote_started, 3),
            "response_preview": preview[:300],
        }
    )
    if torch.cuda.is_available():
        result["cuda_name"] = torch.cuda.get_device_name(0)
        result["cuda_peak_memory_gb"] = round(torch.cuda.max_memory_allocated() / 1024**3, 3)
    return result


@app.local_entrypoint()
def main(
    prompt: str = "我今天有点累，请你用虚拟角色的语气安慰我，并给我一个很短的建议。",
    max_new_tokens: int = 64,
    enable_thinking: bool = False,
):
    started = time.perf_counter()
    result = run_benchmark.remote(prompt=prompt, max_new_tokens=max_new_tokens, enable_thinking=enable_thinking)
    result["client_wall_s"] = round(time.perf_counter() - started, 3)
    print(json.dumps(result, ensure_ascii=False, indent=2))