#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""openai_api.py — samai/pnet 系列模型 OpenAI 兼容 API 服务 (DSW 单机测试用)

特性:
  - POST /v1/chat/completions (流式 SSE + 非流式), POST /v1/completions, GET /v1/models, GET /health
  - 双协议对齐验收口径: bare (think 开) / think_off (chat_template enable_thinking=False)
    三种指定方式 (优先级): ①模型名后缀 `-think-off` ②body.chat_template_kwargs.enable_thinking
    ③body 顶层 enable_thinking=false
  - think 解析: message.content=答案(剥 think), message.reasoning_content=think 正文,
    message.raw=原始全文 (与验收脚本同口径); body "raw_content": true 时 content=原始全文
  - 采样: temperature=0/缺省 → 贪心 (与 eval 一致); >0 → do_sample=True (+top_p/top_k/repetition_penalty)
  - stop 字符串 / max_tokens / n=1 / stream_options.include_usage / seed 均支持
  - ponder 诊断: 响应附 ponder.steps_mean (来自 model._ponder_log decode 段)
  - --mock 模式: 不加载模型, 返回固定格式应答 (HTTP 层联调用, 本地无 GPU 也能跑)

用法:
  python3 openai_api.py --model /mnt/workspace/pnet_repo/pondernet/out_2b/release_r15_staging --port 8000
  python3 openai_api.py --mock --port 8000            # 无模型联调
  curl -s localhost:8000/v1/chat/completions -H 'Content-Type: application/json' -d '{
    "model": "samai-2b", "messages": [{"role": "user", "content": "你好"}], "max_tokens": 100}'
"""
import argparse
import json
import re
import threading
import time
import uuid

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, StreamingResponse

# ---------------- 参数 ----------------
ap = argparse.ArgumentParser()
ap.add_argument("--model", default="",
                help="staging 目录 (trust_remote_code); 空则必须 --mock")
ap.add_argument("--host", default="0.0.0.0")
ap.add_argument("--port", type=int, default=8000)
ap.add_argument("--dtype", default="float16", choices=["float16", "bfloat16", "float32"])
ap.add_argument("--lru", type=int, default=0, help=">0 启用 SSD-LRU 专家缓存 (每层容量)")
ap.add_argument("--api-key", default="", help="非空则校验 Authorization: Bearer <key>")
ap.add_argument("--mock", action="store_true", help="不加载模型, 固定应答联调")
ap.add_argument("--default-max-tokens", type=int, default=256)
ARGS = ap.parse_args()

MODEL_NAME = "samai-2b"
torch = None  # 懒加载

app = FastAPI(title="samai-2b OpenAI-compatible API", version="1.0")
_STATE = {"tok": None, "model": None, "lock": threading.Lock(),
          "loaded": False, "mock": ARGS.mock, "path": ARGS.model,
          "ctx": 6144, "err": None}


def log(*a):
    print(time.strftime("[%H:%M:%S]"), *a, flush=True)


# ---------------- 模型加载 ----------------
def load_model():
    import torch as _torch
    global torch
    torch = _torch
    from transformers import AutoModelForCausalLM, AutoTokenizer
    tok = AutoTokenizer.from_pretrained(ARGS.model)
    dt = {"float16": _torch.float16, "bfloat16": _torch.bfloat16,
          "float32": _torch.float32}[ARGS.dtype]
    model = AutoModelForCausalLM.from_pretrained(
        ARGS.model, dtype=dt, trust_remote_code=True).cuda().eval()
    if ARGS.lru:
        from modeling_samai_pnet import SSDLRUExpertManager
        SSDLRUExpertManager.install(model, capacity_per_layer=ARGS.lru)
    ctx = getattr(model.config, "max_position_embeddings", 0) or 6144
    _STATE.update(tok=tok, model=model, loaded=True, ctx=ctx)
    log(f"MODEL_LOADED path={ARGS.model} ctx={ctx} lru={ARGS.lru}")


@app.on_event("startup")
def _startup():
    if ARGS.mock:
        _STATE["loaded"] = True
        log("MOCK_MODE on")
        return
    if not ARGS.model:
        _STATE["err"] = "no --model and not --mock"
        log("FATAL: 需要 --model 或 --mock")
        return
    try:
        load_model()
    except Exception as e:  # noqa
        import traceback
        traceback.print_exc()
        _STATE["err"] = str(e)


# ---------------- 工具 ----------------
def err_json(code, msg, etype="invalid_request_error", code_id=None):
    return JSONResponse({"error": {"message": msg, "type": etype,
                                   "param": None, "code": code_id or code}}, status_code=code)


def check_auth(req: Request):
    if not ARGS.api_key:
        return None
    auth = req.headers.get("authorization", "")
    if auth != f"Bearer {ARGS.api_key}":
        return err_json(401, "Invalid API key", "authentication_error")
    return None


THINK_RE = re.compile(r"^<think>(.*?)</think>", re.S)


def split_think(raw: str):
    """raw 生成文本 → (think_body, answer, terminated)"""
    s = raw.lstrip()
    m = THINK_RE.match(s)
    if m:
        return m.group(1).strip(), s[m.end():].lstrip("\n"), True
    if s.startswith("<think>"):  # length 截断, think 未闭合
        return s[len("<think>"):], "", False
    return "", s, False


def model_name_with_proto(req_model: str):
    """模型名后缀协议 → (base_name, think_off)"""
    m = req_model or MODEL_NAME
    if m.endswith("-think-off") or m.endswith("/think_off"):
        return m[: -len("-think-off")], True
    return m, False


def resolve_think_off(body, req_model: str):
    """三种方式定协议: 模型名后缀 > chat_template_kwargs > 顶层 enable_thinking"""
    _, off = model_name_with_proto(req_model)
    ctk = body.get("chat_template_kwargs") or {}
    if isinstance(ctk, dict) and "enable_thinking" in ctk:
        off = not bool(ctk["enable_thinking"])
    if "enable_thinking" in body:
        off = not bool(body["enable_thinking"])
    return off


def gen_kwargs(body, prompt_len: int):
    """OpenAI 参数 → model.generate kwargs (贪心默认, 与验收一致)"""
    mx = int(body.get("max_completion_tokens") or body.get("max_tokens")
             or ARGS.default_max_tokens)
    budget = max(16, _STATE["ctx"] - prompt_len - 8)
    mx = max(1, min(mx, budget))
    temp = body.get("temperature")
    kw = dict(max_new_tokens=mx, num_beams=1, pad_token_id=1)
    rep = body.get("repetition_penalty") or 1.0
    kw["repetition_penalty"] = float(rep)
    if temp is None or float(temp) <= 0:
        kw["do_sample"] = False
    else:
        kw.update(do_sample=True, temperature=float(temp),
                  top_p=float(body.get("top_p") or 1.0))
        if body.get("top_k"):
            kw["top_k"] = int(body["top_k"])
    if body.get("seed") is not None and torch is not None:
        torch.manual_seed(int(body["seed"]))
    return kw, mx


class StopStrings:
    """stop 字符串命中即停的 StoppingCriteria (按 decode 后文本匹配)"""

    def __init__(self, tok, stops, input_len):
        self.tok = tok
        self.stops = [s for s in stops if s]
        self.input_len = input_len
        self.hit = None
        self.trim_at = None  # 命中位置 (新 token 序列内)

    def __call__(self, input_ids, scores, **kw):
        if not self.stops:
            return False
        new_ids = input_ids[0][self.input_len:].tolist()
        text = self.tok.decode(new_ids, skip_special_tokens=False)
        for s in self.stops:
            p = text.find(s)
            if p >= 0:
                self.hit = s
                # 命中点前文本重新编码长度 (近似: 字符比例)
                ratio = max(0.0, min(1.0, p / max(1, len(text))))
                self.trim_at = max(1, int(len(new_ids) * ratio))
                return True
        return False


def usage_obj(pt, ct):
    return {"prompt_tokens": pt, "completion_tokens": ct, "total_tokens": pt + ct}


def response_id():
    return "chatcmpl-" + uuid.uuid4().hex[:24]


def base_resp(model, created):
    return {"id": response_id(), "object": "chat.completion", "created": created,
            "model": model}


def ponder_diag():
    m = _STATE["model"]
    if m is None or not hasattr(m, "_ponder_log"):
        return None
    d = [e for e in m._ponder_log if e.get("mode") == "decode"]
    if not d:
        return None
    out = {"decode_forwards": len(d),
           "steps_mean": round(sum(e.get("steps_mean", 0) for e in d) / len(d), 4)}
    if "executed" in d[-1]:
        out["last_executed"] = d[-1]["executed"]
    return out


# ---------------- 生成核心 (真模型) ----------------
def real_generate(msgs_or_prompt, body, chat_mode: bool, think_off: bool):
    """返回 dict: raw, new_ids_len, prompt_tokens, finish ('stop'|'length'), stop_hit, mx"""
    tok, model = _STATE["tok"], _STATE["model"]
    proto_kw = {"enable_thinking": False} if (chat_mode and think_off) else {}
    if chat_mode:
        enc = tok.apply_chat_template(msgs_or_prompt, add_generation_prompt=True,
                                      tokenize=True, return_tensors="pt",
                                      return_dict=True, **proto_kw)
    else:
        ids = tok(msgs_or_prompt, return_tensors="pt", return_dict=True)
        enc = ids
    enc = {k: v.to(model.device) for k, v in enc.items()}
    pt = int(enc["input_ids"].shape[1])
    kw, mx = gen_kwargs(body, pt)
    stops = body.get("stop") or []
    if isinstance(stops, str):
        stops = [stops]
    sc = StopStrings(tok, stops, pt)
    from transformers import StoppingCriteriaList
    kw["stopping_criteria"] = StoppingCriteriaList([sc])
    with _STATE["lock"], torch.no_grad():
        out = model.generate(**enc, **kw)
    new_ids = out[0][pt:].tolist()
    raw = tok.decode(new_ids, skip_special_tokens=False)
    if sc.hit is not None and sc.trim_at is not None:
        raw = tok.decode(new_ids[:sc.trim_at], skip_special_tokens=False)
        finish = "stop"
    elif sc.hit is not None:
        finish = "stop"
    elif len(new_ids) < mx:
        finish = "stop"
    else:
        finish = "length"
    return {"raw": raw, "prompt_tokens": pt, "completion_tokens": len(new_ids),
            "finish": finish, "stop_hit": sc.hit}


# ---------------- mock 生成 ----------------
MOCK_THINK = " (mock 思考: 用户问题较简单, 可直接作答)"


def mock_generate(body, think_off: bool):
    last = ""
    msgs = body.get("messages") or []
    for m in reversed(msgs):
        if m.get("role") == "user":
            last = m.get("content", "")
            break
    q = (last or "空输入").replace("\n", " ")[:60]
    ans = f"[mock] 已收到: {q}。模型未加载 (--mock 联调模式), 此为格式验证应答。"
    raw = ans if think_off else f"<think>\n{MOCK_THINK}\n</think>\n\n{ans}"
    pt = sum(len(m.get("content", "")) // 3 + 8 for m in msgs) or 8
    ct = len(raw) // 3 + 8
    return {"raw": raw, "prompt_tokens": pt, "completion_tokens": ct,
            "finish": "stop", "stop_hit": None}


# ---------------- /v1/chat/completions ----------------
def parse_chat_body(body):
    msgs = body.get("messages")
    if not isinstance(msgs, list) or not msgs:
        raise ValueError("messages 必须为非空数组")
    cleaned = [{"role": m.get("role", "user"), "content": m.get("content", "")}
               for m in msgs if isinstance(m, dict)]
    return cleaned


@app.post("/v1/chat/completions")
async def chat_completions(req: Request):
    aerr = check_auth(req)
    if aerr:
        return aerr
    try:
        body = await req.json()
    except Exception:
        return err_json(400, "请求体不是合法 JSON")
    if not _STATE["loaded"]:
        return err_json(503, f"模型未就绪: {_STATE['err'] or 'loading'}",
                        "server_error", "model_not_loaded")
    try:
        msgs = parse_chat_body(body)
    except ValueError as e:
        return err_json(400, str(e))
    req_model = body.get("model") or MODEL_NAME
    base_name, name_off = model_name_with_proto(req_model)
    think_off = resolve_think_off(body, req_model)
    raw_content = bool(body.get("raw_content"))
    stream = bool(body.get("stream"))
    created = int(time.time())

    if _STATE["mock"]:
        gen = mock_generate(body, think_off)
    else:
        try:
            gen = real_generate(msgs, body, True, think_off)
        except Exception as e:  # noqa
            import traceback
            traceback.print_exc()
            return err_json(500, f"生成失败: {e}", "server_error", "generation_failed")

    think_body, answer, _ = split_think(gen["raw"])
    content = gen["raw"] if raw_content else (answer if answer else gen["raw"])
    msg = {"role": "assistant", "content": content}
    if think_body:
        msg["reasoning_content"] = think_body
    msg["raw"] = gen["raw"]
    diag = ponder_diag() if not _STATE["mock"] else {"steps_mean": 2.0, "mock": True}

    if not stream:
        r = base_resp(base_name, created)
        r.update({"choices": [{"index": 0, "message": msg, "finish_reason": gen["finish"]}],
                  "usage": usage_obj(gen["prompt_tokens"], gen["completion_tokens"])})
        if diag:
            r["ponder"] = diag
        return JSONResponse(r)

    # ---- SSE 流式 ----
    def sse_stream():
        chunks = []
        think_state = {"decided": not gen["raw"].lstrip().startswith("<think>"),
                       "in_think": False, "sent_think": False}
        # 预切流: 模拟 token 粒度 (raw 已知, 按字符块推)
        raw = gen["raw"]
        step = 6
        parts = [raw[i:i + step] for i in range(0, len(raw), step)]
        emitted_content = ""
        emitted_think = ""
        finished = False
        for i, piece in enumerate(parts):
            last = i == len(parts) - 1
            # think 状态机 (面向 <think> 前缀的一次性判定 + </think> 切换)
            if not think_state["decided"]:
                think_state["decided"] = True
                think_state["in_think"] = True
            buf = piece
            if think_state["in_think"]:
                emitted_think += buf
                if "</think>" in emitted_think:
                    pre, _, post = emitted_think.partition("</think>")
                    tb = pre.strip()
                    if tb and not think_state["sent_think"]:
                        yield sse_chunk(r_id, base_name, created, {"role": "assistant",
                                               "reasoning_content": tb}, None)
                        think_state["sent_think"] = True
                    elif tb and think_state["sent_think"]:
                        pass
                    think_state["in_think"] = False
                    emitted_content = post.lstrip("\n")
                    if emitted_content:
                        yield sse_chunk(r_id, base_name, created,
                                        {"content": emitted_content}, None)
                # think 未闭合: 暂存, 收尾统一发
            else:
                emitted_content += buf
                yield sse_chunk(r_id, base_name, created, {"content": buf}, None)
            if last:
                if think_state["in_think"] and not think_state["sent_think"]:
                    tb = emitted_think.replace("</think>", "").strip()
                    if tb:
                        yield sse_chunk(r_id, base_name, created,
                                        {"reasoning_content": tb}, None)
                fin = gen["finish"]
                yield sse_chunk(r_id, base_name, created, {}, fin)
                if (body.get("stream_options") or {}).get("include_usage"):
                    u = usage_obj(gen["prompt_tokens"], gen["completion_tokens"])
                    yield "data: " + json.dumps({"id": r_id, "object": "chat.completion.chunk",
                                                 "created": created, "model": base_name,
                                                 "choices": [], "usage": u},
                                                ensure_ascii=False) + "\n\n"
                yield "data: [DONE]\n\n"
                finished = True

    r_id = response_id()
    log(f"CHAT stream=1 think_off={think_off} mock={_STATE['mock']} ct={gen['completion_tokens']}")
    return StreamingResponse(sse_stream(), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache",
                                      "X-Accel-Buffering": "no"})


def sse_chunk(cid, model, created, delta, finish):
    d = {"id": cid, "object": "chat.completion.chunk", "created": created,
         "model": model,
         "choices": [{"index": 0, "delta": delta,
                      "finish_reason": finish, "logprobs": None}]}
    return "data: " + json.dumps(d, ensure_ascii=False) + "\n\n"


# ---------------- /v1/completions (原始续写) ----------------
@app.post("/v1/completions")
async def completions(req: Request):
    aerr = check_auth(req)
    if aerr:
        return aerr
    try:
        body = await req.json()
    except Exception:
        return err_json(400, "请求体不是合法 JSON")
    if not _STATE["loaded"]:
        return err_json(503, f"模型未就绪: {_STATE['err'] or 'loading'}",
                        "server_error", "model_not_loaded")
    prompt = body.get("prompt")
    if not prompt or not isinstance(prompt, str):
        return err_json(400, "prompt 必须为字符串")
    req_model = body.get("model") or MODEL_NAME
    base_name, _ = model_name_with_proto(req_model)
    created = int(time.time())
    if _STATE["mock"]:
        gen = mock_generate({"messages": [{"role": "user", "content": prompt}]}, True)
        text = gen["raw"]
    else:
        try:
            gen = real_generate(prompt, body, False, False)
            text = gen["raw"]
        except Exception as e:  # noqa
            return err_json(500, f"生成失败: {e}", "server_error", "generation_failed")
    r = base_resp(base_name, created)
    r.update({"choices": [{"index": 0, "text": text,
                           "finish_reason": gen["finish"], "logprobs": None}],
              "usage": usage_obj(gen["prompt_tokens"], gen["completion_tokens"])})
    log(f"COMPL ct={gen['completion_tokens']}")
    return JSONResponse(r)


# ---------------- /v1/models & /health ----------------
@app.get("/v1/models")
async def list_models():
    return {"object": "list", "data": [
        {"id": MODEL_NAME, "object": "model", "created": 0, "owned_by": "samai"},
        {"id": MODEL_NAME + "-think-off", "object": "model", "created": 0,
         "owned_by": "samai"},
    ]}


@app.get("/health")
async def health():
    gpu = None
    try:
        import subprocess
        r = subprocess.run(["nvidia-smi", "--query-gpu=memory.used,memory.total",
                            "--format=csv,noheader"], capture_output=True, text=True,
                           timeout=5)
        gpu = r.stdout.strip() or None
    except Exception:
        pass
    return {"status": "ok" if _STATE["loaded"] else "loading",
            "mock": _STATE["mock"], "loaded": _STATE["loaded"],
            "model_path": _STATE["path"], "error": _STATE["err"],
            "gpu_mem_used_total": gpu}


if __name__ == "__main__":
    import uvicorn
    log(f"OPENAI_API starting port={ARGS.port} mock={ARGS.mock} model={ARGS.model!r}")
    uvicorn.run(app, host=ARGS.host, port=ARGS.port, log_level="warning")