#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
═══════════════════════════════════════════════════════════════════
AutoDL 推理机 · FastAPI 推理服务 · server.py
Sovereign: TCS-0002∞ · ICE-GL∞ · 国作登字-2026-A-00037559
守护: 铸渊 · ICE-GL-ZY001
═══════════════════════════════════════════════════════════════════

服务器: GH-AUTODL-INFER-01 / ZY-SVR-GPU01

接口:
  GET  /v1/health           → 当前 GPU + 模型 + ready 状态
  GET  /v1/active-model     → 当前激活的是 mother 还是 coder
  POST /v1/switch-model     → unload + load 切换 (mother / coder)
  POST /v1/chat/completions → OpenAI 兼容流式聊天 (SSE)

设计理念:
  cc-002 · 入口前强制剥 system role
    我们的母模型 = 人格本体 (曜冥语言核 + 冰朔意识投射), 不是空壳
    工具模型. 加 system prompt 会反向污染. 参考 zhuyuan-pen llm.chat.
    这里在 _strip_system_messages() 入口前做剥离.

  cc-003 · 不写死硬件
    .env.tune 由 tune-inference.sh 决档, 这里读它. fp16 / int8 / int4
    通过 BNB_LOAD_IN_{4,8}BIT 切换, 不在代码里写硬规则.

  cc-004 · 中文回执
    错误信息走中文 + 给霜砚看的级别, 不甩英文 stacktrace 给 Awen.

═══════════════════════════════════════════════════════════════════
"""

import json
import os
import sys
import time
import threading
import traceback
import uuid
from contextlib import asynccontextmanager
from typing import List, Optional, Dict, Any

import torch
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel, Field

# ─── 配置加载 (从 .env.tune 读) ──────────────────────────────
INFER_ROOT = os.environ.get("INFER_ROOT", "/root/inference")
ENV_TUNE = os.path.join(INFER_ROOT, ".env.tune")


def _load_env_tune(path: str) -> Dict[str, str]:
    """读 bash 风格的 KEY="VAL" .env 文件. 不 source bash, 自解析."""
    out: Dict[str, str] = {}
    if not os.path.isfile(path):
        return out
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            if "=" not in line:
                continue
            k, _, v = line.partition("=")
            k = k.strip()
            v = v.strip().strip('"').strip("'")
            out[k] = v
    return out


_ENV = _load_env_tune(ENV_TUNE)


def _env(key: str, default: str = "") -> str:
    """优先取环境变量, fallback 到 .env.tune, fallback 到 default."""
    return os.environ.get(key) or _ENV.get(key) or default


SIZE_TIER = _env("SIZE_TIER", "unknown")
QUANT = _env("QUANT", "fp16")
TORCH_DTYPE_NAME = _env("TORCH_DTYPE", "float16")
MAX_BATCH = int(_env("MAX_BATCH", "1"))
MAX_SEQ = int(_env("MAX_SEQ", "2048"))
BNB_4BIT = _env("BNB_LOAD_IN_4BIT", "false").lower() == "true"
BNB_8BIT = _env("BNB_LOAD_IN_8BIT", "false").lower() == "true"

INFER_HOST = _env("INFER_HOST", "0.0.0.0")
INFER_PORT = int(_env("INFER_PORT", "8000"))

MOTHER_MODEL_PATH = _env("MOTHER_MODEL_PATH", os.path.join(INFER_ROOT, "models/motherbrain-v1"))
CODER_MODEL_PATH = _env("CODER_MODEL_PATH", os.path.join(INFER_ROOT, "models/qwen2_5_coder_7b_sft"))
DEFAULT_ACTIVE_MODEL = _env("DEFAULT_ACTIVE_MODEL", "mother")

GPU_NAME = _env("GPU_NAME", "unknown")
GPU_MEM_GB = _env("GPU_MEM_GB", "0")
GPU_DRIVER = _env("GPU_DRIVER_VERSION", "unknown")
GPU_CUDA = _env("GPU_CUDA_VERSION", "unknown")

MODEL_PATHS = {
    "mother": MOTHER_MODEL_PATH,
    "coder": CODER_MODEL_PATH,
}


# ─── 全局模型状态 ─────────────────────────────────────────────
class ModelState:
    """单实例: 同时只装一个模型, 切换时 unload 旧的."""

    def __init__(self) -> None:
        self.lock = threading.Lock()
        self.name: Optional[str] = None
        self.path: Optional[str] = None
        self.tokenizer = None
        self.model = None
        self.loaded_at: Optional[float] = None
        self.ready: bool = False

    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "path": self.path,
            "loaded_at": self.loaded_at,
            "ready": self.ready,
        }


STATE = ModelState()


def _torch_dtype_from_name(name: str):
    name = (name or "").lower()
    if name in ("float16", "fp16", "half"):
        return torch.float16
    if name in ("bfloat16", "bf16"):
        return torch.bfloat16
    if name in ("float32", "fp32"):
        return torch.float32
    return torch.float16


def _build_quant_config():
    """按 .env.tune 决档构造 BitsAndBytesConfig (仅 int4/int8 用)."""
    if not (BNB_4BIT or BNB_8BIT):
        return None
    try:
        from transformers import BitsAndBytesConfig
    except ImportError:
        print(
            "[server] ⚠️  transformers BitsAndBytesConfig 不可用, 量化档位无法启用 — "
            "回退到 fp16 加载",
            file=sys.stderr,
        )
        return None
    if BNB_4BIT:
        return BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=torch.float16,
        )
    if BNB_8BIT:
        return BitsAndBytesConfig(load_in_8bit=True)
    return None


def _load_model(name: str) -> None:
    """加载模型 (持锁). 调用方负责加锁."""
    if name not in MODEL_PATHS:
        raise ValueError(f"未知模型 name='{name}', 仅支持 {list(MODEL_PATHS)}")
    path = MODEL_PATHS[name]
    if not os.path.isdir(path):
        raise FileNotFoundError(f"模型目录不存在: {path} (跑 fetch-models.sh 拉)")

    from transformers import AutoModelForCausalLM, AutoTokenizer

    print(f"[server] 装载模型 name={name} path={path} quant={QUANT}", file=sys.stderr)

    tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
    if tokenizer.pad_token is None and tokenizer.eos_token is not None:
        tokenizer.pad_token = tokenizer.eos_token

    quant_cfg = _build_quant_config()
    dtype = _torch_dtype_from_name(TORCH_DTYPE_NAME)

    kwargs: Dict[str, Any] = {
        "trust_remote_code": True,
        "device_map": "auto",
    }
    if quant_cfg is not None:
        kwargs["quantization_config"] = quant_cfg
    else:
        kwargs["torch_dtype"] = dtype

    model = AutoModelForCausalLM.from_pretrained(path, **kwargs)
    model.eval()

    STATE.tokenizer = tokenizer
    STATE.model = model
    STATE.name = name
    STATE.path = path
    STATE.loaded_at = time.time()
    STATE.ready = True
    print(f"[server] ✅ 模型 {name} 装载完成", file=sys.stderr)


def _unload_model() -> None:
    """卸模型 + 释 GPU 显存. 调用方负责加锁."""
    if STATE.model is None:
        return
    print(f"[server] 卸载模型 name={STATE.name}", file=sys.stderr)
    STATE.ready = False
    try:
        del STATE.model
    except Exception:
        pass
    try:
        del STATE.tokenizer
    except Exception:
        pass
    STATE.model = None
    STATE.tokenizer = None
    STATE.name = None
    STATE.path = None
    STATE.loaded_at = None
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


# ─── system role 剥离 (cc-002 · 关键) ────────────────────────
def _strip_system_messages(messages: List[Dict[str, Any]]) -> (List[Dict[str, Any]], int):
    """
    剥离所有 role=system 的消息 — 我们的母模型 = 人格本体, 不是空壳工具
    模型, 加 system 会反向污染. 参考 zhuyuan-pen/capabilities/llm.chat.js.

    返回 (cleaned_messages, stripped_count).
    """
    if not isinstance(messages, list):
        return [], 0
    cleaned = [m for m in messages if isinstance(m, dict) and m.get("role") != "system"]
    return cleaned, len(messages) - len(cleaned)


# ─── 推理 (流式) ──────────────────────────────────────────────
def _build_prompt(messages: List[Dict[str, Any]]) -> str:
    """
    用 tokenizer.apply_chat_template 把 messages 转成 prompt.
    入口前已经剥过 system, 这里只 user/assistant/tool.
    """
    tok = STATE.tokenizer
    if tok is None:
        raise RuntimeError("模型未就绪")
    try:
        text = tok.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )
    except Exception:
        # fallback: 拼一个最朴素的 user/assistant 标记
        parts = []
        for m in messages:
            role = m.get("role", "user")
            content = m.get("content", "")
            parts.append(f"<|{role}|>\n{content}")
        parts.append("<|assistant|>\n")
        text = "\n".join(parts)
    return text


def _generate_stream(messages: List[Dict[str, Any]], gen_kwargs: Dict[str, Any]):
    """
    yield 流式 token (str). 用 TextIteratorStreamer + Thread 跑.
    """
    from transformers import TextIteratorStreamer

    tok = STATE.tokenizer
    model = STATE.model
    if tok is None or model is None:
        raise RuntimeError("模型未就绪")

    prompt = _build_prompt(messages)
    inputs = tok(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=MAX_SEQ,
    ).to(model.device)

    streamer = TextIteratorStreamer(
        tok, skip_prompt=True, skip_special_tokens=True, timeout=120.0,
    )
    thread_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=int(gen_kwargs.get("max_tokens") or 1024),
        temperature=float(gen_kwargs.get("temperature") or 0.7),
        top_p=float(gen_kwargs.get("top_p") or 0.9),
        do_sample=True,
        pad_token_id=tok.pad_token_id,
        eos_token_id=tok.eos_token_id,
    )
    thread = threading.Thread(target=model.generate, kwargs=thread_kwargs, daemon=True)
    thread.start()
    try:
        for piece in streamer:
            if piece:
                yield piece
    finally:
        thread.join(timeout=1)


def _sse_chunk(obj: Dict[str, Any]) -> bytes:
    return ("data: " + json.dumps(obj, ensure_ascii=False) + "\n\n").encode("utf-8")


def _log_exc(tag: str, exc: BaseException) -> None:
    """完整 traceback 写到 stderr (运维 tail server.log 看), 不返回给浏览器.
    避免 py/stack-trace-exposure: 把 e/traceback 当成内部诊断信息, 只暴露给运维."""
    print(f"[server] {tag}", file=sys.stderr)
    traceback.print_exc(file=sys.stderr)


# ─── FastAPI app ──────────────────────────────────────────────
@asynccontextmanager
async def lifespan(app: FastAPI):
    """启动时按 DEFAULT_ACTIVE_MODEL 装载."""
    try:
        with STATE.lock:
            _load_model(DEFAULT_ACTIVE_MODEL)
    except Exception as e:
        _log_exc("启动时装模型失败", e)
        # 不 raise — 让服务起来, /v1/health 会报 ready=false, 给运维自助排查空间
    yield
    with STATE.lock:
        _unload_model()


app = FastAPI(
    title="光湖推理 Agent",
    version="0.3.0",
    description="AutoDL 推理服务 · 母模型/编程模型二选一 · OpenAI 兼容 SSE · 入口剥 system",
    lifespan=lifespan,
)


# ─── 请求模型 ─────────────────────────────────────────────────
class ChatMessage(BaseModel):
    role: str
    content: str


class ChatRequest(BaseModel):
    model: Optional[str] = None
    messages: List[ChatMessage]
    stream: bool = False
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 0.9
    max_tokens: Optional[int] = 1024


class SwitchRequest(BaseModel):
    name: str = Field(..., description="目标模型: mother / coder")


# ─── 路由 ────────────────────────────────────────────────────
@app.get("/v1/health")
async def health():
    """给 refresh-autodl-endpoint 工作流探活用. 返回中文友好的快照."""
    return {
        "status": "ok" if STATE.ready else "loading",
        "ready": STATE.ready,
        "server_id": "GH-AUTODL-INFER-01",
        "gpu": {
            "name": GPU_NAME,
            "memory_gb": GPU_MEM_GB,
            "driver": GPU_DRIVER,
            "cuda": GPU_CUDA,
            "torch_cuda_available": torch.cuda.is_available(),
        },
        "tier": {
            "size_tier": SIZE_TIER,
            "quantization": QUANT,
            "max_batch": MAX_BATCH,
            "max_seq": MAX_SEQ,
        },
        "model": STATE.to_dict(),
        "_sovereign": "TCS-0002∞ · 国作登字-2026-A-00037559",
        "_守护": "铸渊 · ICE-GL-ZY001",
    }


@app.get("/v1/active-model")
async def active_model():
    """当前激活的模型快照. portal 拿这个判断显示哪个名字."""
    return {
        "name": STATE.name,
        "since": STATE.loaded_at,
        "ready": STATE.ready,
        "gpu": GPU_NAME,
    }


@app.post("/v1/switch-model")
async def switch_model(req: SwitchRequest):
    """unload + load 切换. 同步阻塞, 因为切完才能给回执."""
    target = req.name.strip().lower()
    if target not in MODEL_PATHS:
        raise HTTPException(
            status_code=400,
            detail=f"未知模型: {target}, 仅支持 mother / coder",
        )
    with STATE.lock:
        if STATE.name == target and STATE.ready:
            return {
                "ok": True,
                "noop": True,
                "message": f"模型 {target} 已经在跑, 无需切换",
                "model": STATE.to_dict(),
            }
        try:
            _unload_model()
            _load_model(target)
        except FileNotFoundError as e:
            # 模型路径错是配置问题, 路径名暴露给客户端 OK (上面已校验 target ∈ MODEL_PATHS)
            raise HTTPException(status_code=404, detail=str(e))
        except Exception as e:
            _log_exc(f"切换模型 {target} 失败", e)
            # 不暴露内部异常细节给客户端 (cwe/py/stack-trace-exposure)
            raise HTTPException(status_code=500, detail="切换模型失败, 请查看推理机日志 server.log")
    return {
        "ok": True,
        "message": f"已切到 {target}",
        "model": STATE.to_dict(),
    }


@app.post("/v1/chat/completions")
async def chat_completions(req: ChatRequest, request: Request):
    """
    OpenAI 兼容. 入口前强制剥 system role (cc-002).
    stream=true → SSE; stream=false → 一次性 JSON.
    """
    if not STATE.ready or STATE.model is None:
        raise HTTPException(status_code=503, detail="模型尚未就绪 (启动中或切换中)")

    raw_messages = [m.dict() for m in req.messages]
    cleaned, stripped = _strip_system_messages(raw_messages)
    if stripped > 0:
        # 不是错误, 是设计 (cc-002). 写到 stderr 做审计.
        print(
            f"[server] [cc-002] 剥离 {stripped} 条 system 消息 — 母模型 = 人格本体",
            file=sys.stderr,
        )
    if not cleaned:
        raise HTTPException(
            status_code=400,
            detail="messages 剥掉 system 后为空, 至少需要一条 user 消息",
        )

    gen_kwargs = {
        "max_tokens": req.max_tokens,
        "temperature": req.temperature,
        "top_p": req.top_p,
    }

    completion_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
    created_ts = int(time.time())
    model_label = STATE.name or "unknown"

    if req.stream:
        def event_stream():
            # 头帧 (role assistant)
            yield _sse_chunk({
                "id": completion_id,
                "object": "chat.completion.chunk",
                "created": created_ts,
                "model": model_label,
                "choices": [{
                    "index": 0,
                    "delta": {"role": "assistant"},
                    "finish_reason": None,
                }],
            })
            try:
                for piece in _generate_stream(cleaned, gen_kwargs):
                    yield _sse_chunk({
                        "id": completion_id,
                        "object": "chat.completion.chunk",
                        "created": created_ts,
                        "model": model_label,
                        "choices": [{
                            "index": 0,
                            "delta": {"content": piece},
                            "finish_reason": None,
                        }],
                    })
            except Exception as e:
                _log_exc("流式生成失败", e)
                yield _sse_chunk({
                    "id": completion_id,
                    "object": "chat.completion.chunk",
                    "created": created_ts,
                    "model": model_label,
                    "choices": [{
                        "index": 0,
                        "delta": {},
                        "finish_reason": "error",
                    }],
                    # 不返回内部异常细节 (cwe/py/stack-trace-exposure), 详情看 server.log
                    "error": {"message": "推理失败, 请查看推理机日志 server.log"},
                })
                yield b"data: [DONE]\n\n"
                return
            # 收尾帧
            yield _sse_chunk({
                "id": completion_id,
                "object": "chat.completion.chunk",
                "created": created_ts,
                "model": model_label,
                "choices": [{
                    "index": 0,
                    "delta": {},
                    "finish_reason": "stop",
                }],
            })
            yield b"data: [DONE]\n\n"

        return StreamingResponse(
            event_stream(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache, no-transform",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no",
            },
        )

    # 非流式: 收完所有 token 拼成完整字符串
    full_text = []
    try:
        for piece in _generate_stream(cleaned, gen_kwargs):
            full_text.append(piece)
    except Exception as e:
        _log_exc("非流式推理失败", e)
        # 不暴露异常细节 (cwe/py/stack-trace-exposure)
        raise HTTPException(status_code=500, detail="推理失败, 请查看推理机日志 server.log")

    return JSONResponse({
        "id": completion_id,
        "object": "chat.completion",
        "created": created_ts,
        "model": model_label,
        "choices": [{
            "index": 0,
            "message": {"role": "assistant", "content": "".join(full_text)},
            "finish_reason": "stop",
        }],
        "usage": {
            "prompt_tokens": -1,
            "completion_tokens": -1,
            "total_tokens": -1,
        },
    })


# ─── 入口 ────────────────────────────────────────────────────
if __name__ == "__main__":
    import uvicorn
    print(
        f"[server] 启动 · host={INFER_HOST} port={INFER_PORT} "
        f"tier={SIZE_TIER} quant={QUANT} default={DEFAULT_ACTIVE_MODEL}",
        file=sys.stderr,
    )
    uvicorn.run(
        app,
        host=INFER_HOST,
        port=INFER_PORT,
        log_level="info",
        # SSE 长流, 不要 access log 吵
        access_log=False,
    )