#!/usr/bin/env python3
"""
Uncensored chat LLM API — loads an instruct model once, serves an
OpenAI-compatible /v1/chat/completions endpoint. Runs on the AMD MI50 (gfx906)
via ROCm 5.7 + torch 2.3.1, mirroring the joycaption service pattern.

Model fits in fp16 inside the 32GB VRAM (≈12B ceiling). Override with env MODEL_ID.

Endpoints:
  POST /v1/chat/completions  — OpenAI-compatible chat, supports stream=true (SSE)
  POST /v1/completions       — legacy text completion, supports stream=true (SSE)
  GET  /v1/models            — list the loaded model
  GET  /health               — health check + GPU info

Env:
  MODEL_ID   HuggingFace repo id (default below)
  PORT       listen port (default 8001)
  HSA_OVERRIDE_GFX_VERSION=9.0.6  set by start script for gfx906
"""

import asyncio
import json
import os
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
import subprocess
import threading
import time
import uuid
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager

import torch
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

torch.set_num_threads(1)
torch.set_num_interop_threads(1)

gpu_executor = ThreadPoolExecutor(max_workers=1)

# Mistral [INST] template — fallback for models without chat_template in tokenizer_config
_MISTRAL_TEMPLATE = (
    "{{ bos_token }}"
    "{% for message in messages %}"
    "{% if message['role'] == 'system' %}"
    "{{ '[INST] ' + message['content'] + '\n\n' }}"
    "{% elif message['role'] == 'user' %}"
    "{% if not loop.first or messages[0]['role'] != 'system' %}{{ '[INST] ' }}{% endif %}"
    "{{ message['content'] + ' [/INST]' }}"
    "{% elif message['role'] == 'assistant' %}"
    "{{ ' ' + message['content'] + eos_token }}"
    "{% endif %}"
    "{% endfor %}"
)

MODEL_ID = os.environ.get("MODEL_ID", "dphn/Dolphin3.0-Mistral-24B")
# For 24B fp16 (~48GB): split across 32GB VRAM + CPU RAM via device_map="auto".
# Override with MAX_GPU_MEM / MAX_CPU_MEM env vars if needed.
MAX_GPU_MEM = os.environ.get("MAX_GPU_MEM", "27GiB")   # 32GiB total - 4GiB sys - 1GiB headroom
MAX_CPU_MEM = os.environ.get("MAX_CPU_MEM", "100GiB")  # 113GB available

state: dict = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    # Boost GPU to high performance mode (avoids power-saving clock throttle)
    subprocess.run(["/opt/rocm/bin/rocm-smi", "--setperflevel", "high"], capture_output=True)
    print(f"Loading model {MODEL_ID} (gpu≤{MAX_GPU_MEM} cpu≤{MAX_CPU_MEM})...", flush=True)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, local_files_only=True)
    if not tokenizer.chat_template:
        print("No chat_template found — applying Mistral [INST] fallback.", flush=True)
        tokenizer.chat_template = _MISTRAL_TEMPLATE
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float16,
        device_map="auto",
        max_memory={0: MAX_GPU_MEM, "cpu": MAX_CPU_MEM},
        local_files_only=True,
    )
    model.eval()
    state["tokenizer"] = tokenizer
    state["model"] = model

    # Warm up: tiny generation to trigger kernel compilation
    print("Warming up...", flush=True)
    _run_chat([{"role": "user", "content": "hi"}], max_tokens=4, temperature=0.0, top_p=1.0, stop=None)
    print(f"Model ready on: {next(model.parameters()).device}", flush=True)
    yield
    gpu_executor.shutdown(wait=False, cancel_futures=True)
    state.clear()


app = FastAPI(title="Uncensored Chat LLM API", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


def _run_chat(messages, max_tokens, temperature, top_p, stop):
    tokenizer = state["tokenizer"]
    model = state["model"]

    input_ids = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt"
    ).to(model.device)
    prompt_tokens = input_ids.shape[1]

    do_sample = temperature is not None and temperature > 0
    gen_kwargs = dict(
        max_new_tokens=max_tokens,
        do_sample=do_sample,
        use_cache=True,
        pad_token_id=tokenizer.eos_token_id,
    )
    if do_sample:
        gen_kwargs["temperature"] = temperature
        gen_kwargs["top_p"] = top_p

    t0 = time.perf_counter()
    with torch.inference_mode():
        output_ids = model.generate(input_ids, **gen_kwargs)
    dt = time.perf_counter() - t0

    new_ids = output_ids[0][prompt_tokens:]
    text = tokenizer.decode(new_ids, skip_special_tokens=True)

    finish = "stop"
    if stop:
        stops = [stop] if isinstance(stop, str) else list(stop)
        cut = min((text.find(s) for s in stops if s and s in text), default=-1)
        if cut != -1:
            text = text[:cut]
    completion_tokens = int(new_ids.shape[0])
    if completion_tokens >= max_tokens:
        finish = "length"

    del input_ids, output_ids
    return {
        "text": text.strip(),
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "finish_reason": finish,
        "generate_s": round(dt, 2),
    }


async def run_chat(*args):
    return await asyncio.get_running_loop().run_in_executor(gpu_executor, _run_chat, *args)


def _iter_stream(messages, max_tokens, temperature, top_p, stop):
    """Generator that yields SSE lines for a streaming chat completion."""
    tokenizer = state["tokenizer"]
    model = state["model"]
    cid = f"chatcmpl-{uuid.uuid4().hex[:24]}"
    created = int(time.time())

    input_ids = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt"
    ).to(model.device)

    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    do_sample = temperature is not None and temperature > 0
    gen_kwargs = dict(
        input_ids=input_ids,
        max_new_tokens=max_tokens,
        do_sample=do_sample,
        use_cache=True,
        pad_token_id=tokenizer.eos_token_id,
        streamer=streamer,
    )
    if do_sample:
        gen_kwargs["temperature"] = temperature
        gen_kwargs["top_p"] = top_p

    stops = []
    if stop:
        stops = [stop] if isinstance(stop, str) else list(stop)

    def _generate():
        with torch.inference_mode():
            model.generate(**gen_kwargs)

    t = threading.Thread(target=_generate, daemon=True)
    t.start()

    buffer = ""
    for token in streamer:
        buffer += token
        # Check stop sequences across the accumulated buffer
        if stops:
            cut = min((buffer.find(s) for s in stops if s and s in buffer), default=-1)
            if cut != -1:
                token = buffer[len(buffer) - len(token):cut] if cut < len(buffer) else ""
                chunk = {
                    "id": cid, "object": "chat.completion.chunk", "created": created,
                    "model": MODEL_ID,
                    "choices": [{"index": 0, "delta": {"content": token}, "finish_reason": None}],
                }
                yield f"data: {json.dumps(chunk)}\n\n"
                break
        chunk = {
            "id": cid, "object": "chat.completion.chunk", "created": created,
            "model": MODEL_ID,
            "choices": [{"index": 0, "delta": {"content": token}, "finish_reason": None}],
        }
        yield f"data: {json.dumps(chunk)}\n\n"

    t.join()
    done = {
        "id": cid, "object": "chat.completion.chunk", "created": created,
        "model": MODEL_ID,
        "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
    }
    yield f"data: {json.dumps(done)}\n\n"
    yield "data: [DONE]\n\n"
    del input_ids


class ChatMessage(BaseModel):
    role: str
    content: str


class ChatRequest(BaseModel):
    model: str | None = None
    messages: list[ChatMessage]
    max_tokens: int = 512
    temperature: float = 0.8
    top_p: float = 0.95
    stop: list[str] | str | None = None
    stream: bool = False


class CompletionRequest(BaseModel):
    model: str | None = None
    prompt: str
    max_tokens: int = 512
    temperature: float = 0.8
    top_p: float = 0.95
    stop: list[str] | str | None = None
    stream: bool = False


@app.post("/v1/chat/completions")
async def chat_completions(req: ChatRequest):
    if not req.messages:
        raise HTTPException(400, "messages must not be empty")
    msgs = [{"role": m.role, "content": m.content} for m in req.messages]

    if req.stream:
        gen = _iter_stream(msgs, req.max_tokens, req.temperature, req.top_p, req.stop)
        return StreamingResponse(gen, media_type="text/event-stream")

    r = await run_chat(msgs, req.max_tokens, req.temperature, req.top_p, req.stop)
    return {
        "id": f"chatcmpl-{uuid.uuid4().hex[:24]}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": MODEL_ID,
        "choices": [
            {
                "index": 0,
                "message": {"role": "assistant", "content": r["text"]},
                "finish_reason": r["finish_reason"],
            }
        ],
        "usage": {
            "prompt_tokens": r["prompt_tokens"],
            "completion_tokens": r["completion_tokens"],
            "total_tokens": r["prompt_tokens"] + r["completion_tokens"],
        },
        "timing": {"generate_s": r["generate_s"]},
    }


@app.post("/v1/completions")
async def completions(req: CompletionRequest):
    """Legacy text completion endpoint — wraps prompt as a user message."""
    msgs = [{"role": "user", "content": req.prompt}]

    if req.stream:
        cid = f"cmpl-{uuid.uuid4().hex[:24]}"
        created = int(time.time())

        def _legacy_stream():
            for sse in _iter_stream(msgs, req.max_tokens, req.temperature, req.top_p, req.stop):
                # Re-wrap chat chunk format as legacy completion chunk
                if sse.startswith("data: [DONE]"):
                    yield sse
                    return
                if sse.startswith("data: "):
                    chat_chunk = json.loads(sse[6:])
                    delta = chat_chunk["choices"][0]["delta"].get("content", "")
                    finish = chat_chunk["choices"][0]["finish_reason"]
                    chunk = {
                        "id": cid, "object": "text_completion", "created": created,
                        "model": MODEL_ID,
                        "choices": [{"index": 0, "text": delta, "finish_reason": finish}],
                    }
                    yield f"data: {json.dumps(chunk)}\n\n"

        return StreamingResponse(_legacy_stream(), media_type="text/event-stream")

    r = await run_chat(msgs, req.max_tokens, req.temperature, req.top_p, req.stop)
    return {
        "id": f"cmpl-{uuid.uuid4().hex[:24]}",
        "object": "text_completion",
        "created": int(time.time()),
        "model": MODEL_ID,
        "choices": [
            {
                "index": 0,
                "text": r["text"],
                "finish_reason": r["finish_reason"],
            }
        ],
        "usage": {
            "prompt_tokens": r["prompt_tokens"],
            "completion_tokens": r["completion_tokens"],
            "total_tokens": r["prompt_tokens"] + r["completion_tokens"],
        },
    }


@app.get("/v1/models")
async def list_models():
    return {
        "object": "list",
        "data": [{"id": MODEL_ID, "object": "model", "owned_by": "local", "permission": []}],
    }


@app.get("/health")
async def health():
    gpu = {}
    if torch.cuda.is_available():
        gpu = {
            "name": torch.cuda.get_device_name(0),
            "vram_total_gb": round(torch.cuda.get_device_properties(0).total_memory / 1e9, 1),
            "vram_used_gb": round(torch.cuda.memory_allocated(0) / 1e9, 1),
        }
    return {"status": "ok", "model": MODEL_ID, "loaded": "model" in state, "gpu": gpu}


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 8001)))