#!/usr/bin/env python
"""Minimal OpenAI-compatible server to serve Qwen3-4B on CPU for fingerprinting.

Exposes what llm-fingerprint-detector needs: GET /v1/models, POST /v1/chat/completions.
Hidden thinking is disabled server-side (enable_thinking=False) so single-token
probes return visible text immediately and the detector uses the high-confidence
'none' reasoning strategy.

Run:  python serve_qwen_cpu.py [--model /data1/models/Qwen3-4B] [--port 30002]
"""
import argparse
import os
import threading

os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")

import torch
import uvicorn
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "Qwen3-4B"
_lock = threading.Lock()
_model = None
_tokenizer = None

app = FastAPI(title="Qwen3-4B CPU server (OpenAI-compatible)")


def load(ModelDir):
    global _model, _tokenizer
    tok = AutoTokenizer.from_pretrained(ModelDir, trust_remote_code=True)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    model = AutoModelForCausalLM.from_pretrained(
        ModelDir,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
    )
    model.eval()
    _tokenizer, _model = tok, model
    print(f"[server] loaded {ModelDir} on backend:", model.device)


def _generate(messages, temperature, max_tokens):
    global _model, _tokenizer
    tok, model = _tokenizer, _model
    text = tok.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,  # disable hidden thinking
    )
    in_toks = tok(text, return_tensors="pt", return_token_type_ids=False).to(model.device)
    gen_kwargs = dict(
        max_new_tokens=max(1, int(max_tokens or 16)),
        pad_token_id=tok.eos_token_id,
        eos_token_id=tok.eos_token_id,
    )
    temp = float(temperature) if temperature is not None else 1.0
    if temp > 0:
        gen_kwargs.update(do_sample=True, temperature=temp, top_p=0.95)
    else:
        gen_kwargs.update(do_sample=False)
    with torch.no_grad():
        out = model.generate(**in_toks, **gen_kwargs)
    gen = out[0, in_toks.input_ids.shape[1]:]
    return tok.decode(gen, skip_special_tokens=True)


@app.get("/health")
async def health():
    return {"status": "ok"}


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


@app.post("/v1/chat/completions")
async def chat(request: Request):
    body = await request.json()
    messages = [
        {"role": m["role"], "content": m["content"]}
        for m in body.get("messages", [])
        if m.get("role") in ("system", "user", "assistant")
    ]
    max_tokens = body.get("max_tokens", 16)
    temperature = body.get("temperature", 1.0)

    def run():
        with _lock:
            return _generate(messages, temperature, max_tokens)

    import asyncio
    loop = asyncio.get_event_loop()
    content = await loop.run_in_executor(None, run)
    model_name = body.get("model", MODEL_ID)
    import time
    base = {
        "id": "chatcmpl-local",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": model_name,
        "system_fingerprint": "qwen-cpu-local",
    }
    usage = {"prompt_tokens": 0, "completion_tokens": len(content),
             "total_tokens": len(content)}

    # 流式请求：按 OpenAI SSE 协议分块吐出（evalscope 默认 stream=True）
    if body.get("stream"):
        import json as _json
        from fastapi.responses import StreamingResponse

        def sse():
            def chunk(delta, finish=None, usage=None):
                payload = {**base, "object": "chat.completion.chunk",
                           "choices": [{"index": 0,
                                        "delta": delta,
                                        "finish_reason": finish}]}
                if usage:
                    payload["usage"] = usage
                return "data: " + _json.dumps(payload) + "\n\n"
            yield chunk({"role": "assistant"})
            step = max(1, len(content) // 8)
            for i in range(0, len(content), step):
                yield chunk({"content": content[i:i + step]})
            yield chunk({}, "stop", usage=usage)
            yield "data: [DONE]\n\n"

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

    return JSONResponse({
        **base,
        "choices": [{
            "index": 0,
            "message": {"role": "assistant", "content": content},
            "finish_reason": "stop",
        }],
        "usage": usage,
    })


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="/data1/models/Qwen3-4B")
    ap.add_argument("--port", type=int, default=30002)
    ap.add_argument("--host", default="0.0.0.0")
    args = ap.parse_args()

    import asyncio
    load(args.model)
    loop = asyncio.new_event_loop()
    app.state.loop = loop
    config = uvicorn.Config(app, host=args.host, port=args.port, log_level="info")
    server = uvicorn.Server(config)
    loop.run_until_complete(server.serve())


if __name__ == "__main__":
    main()