from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import StreamingResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from transformers import AutoTokenizer, AutoModel, AutoConfig
import torch
import uvicorn
from modelscope import snapshot_download
import warnings
import json
import time
import subprocess
from threading import Lock
import os
# 屏蔽transformers冗余日志告警
warnings.filterwarnings("ignore")

# ===================== 全局监控变量 =====================
class LLMMonitorData:
    # 服务健康状态
    service_health = 1
    model_loaded = 1
    # 推理耗时指标
    last_ttft_seconds = 0.0        # 上一次首token耗时
    last_inference_seconds = 0.0   # 上一次完整推理总耗时
    # Token吞吐指标（体现4090推理性能）
    token_speed = 0.0              # tok/s，核心吞吐衡量标准
    total_input_tokens = 0
    total_output_tokens = 0
    total_inference_count = 0
    # RTX4090显存硬件指标
    gpu_mem_total_mb = 0
    gpu_mem_used_mb = 0
    gpu_mem_usage_pct = 0.0
    # 并发队列 & 异常统计
    request_queue_length = 0
    error_total = 0
# 全局单例监控对象
monitor = LLMMonitorData()
# 双锁分离：队列锁、指标锁，互不阻塞
queue_lock = Lock()
metric_lock = Lock()

# ===================== 模型基础配置 =====================
model = None
tokenizer = None
model_name = "ZhipuAI/chatglm3-6b"
cache_dir = "/home/model"
# 静态页面目录，与项目static文件夹绑定
STATIC_DIR = os.path.join(os.path.dirname(__file__), "static")
# 初始化服务实例
app = FastAPI(title="ChatGLM3-6B 生产级对话服务")
# 托管前端静态文件，访问 /static/index.html 加载页面
os.makedirs(STATIC_DIR, exist_ok=True)
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
# 全局跨域配置，允许前端、第三方工具跨接口调用
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 加载模型权重、分词器、配置文件
print("正在下载/校验模型缓存...")
local_model_path = snapshot_download(model_name, cache_dir=cache_dir)
print(f"正在加载模型: {local_model_path}")
tokenizer = AutoTokenizer.from_pretrained(local_model_path, trust_remote_code=True)
config = AutoConfig.from_pretrained(local_model_path, trust_remote_code=True)
# 修复ChatGLM3 Config无max_length属性报错
if not hasattr(config, 'max_length'):
    config.max_length = config.seq_length if hasattr(config, 'seq_length') else 8192
try:
    if not torch.cuda.is_available():
        raise RuntimeError("未检测到可用GPU设备，服务无法运行")
    # RTX4090专属加载参数：FP16半精度、低内存加载、自动设备分配
    model = AutoModel.from_pretrained(
        local_model_path,
        config=config,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
        torch_dtype=torch.float16,
        device_map="auto"
    ).eval()
    # 标记模型加载正常
    monitor.service_health = 1
    monitor.model_loaded = 1
    print("✅ ChatGLM3-6B模型加载完成，RTX4090推理就绪")
except Exception as e:
    monitor.service_health = 0
    monitor.model_loaded = 0
    print(f"❌ 模型加载失败: {e}")

# ===================== GPU显存采集工具 =====================
def update_gpu_metrics(gpu_id=0):
    try:
        res = subprocess.check_output(
            [
                "nvidia-smi", f"--id={gpu_id}",
                "--query-gpu=memory.total,memory.used",
                "--format=csv,noheader,nounits"
            ], encoding="utf-8", timeout=3
        )
        total, used = map(int, res.strip().split(", "))
        with metric_lock:
            monitor.gpu_mem_total_mb = total
            monitor.gpu_mem_used_mb = used
            monitor.gpu_mem_usage_pct = round(used / total * 100, 2)
    except Exception:
        # 采集失败不中断服务，保留上一次显存数据
        pass

@app.get("/", response_class=FileResponse)
async def index():
    return os.path.join(STATIC_DIR, "index.html")

@app.get("/monitor")
def get_monitor():
    update_gpu_metrics()
    with metric_lock:
        return {
            "llm_service_health": monitor.service_health,
            "llm_model_loaded": monitor.model_loaded,
            "llm_last_ttft_seconds": round(monitor.last_ttft_seconds, 4),
            "llm_last_inference_seconds": round(monitor.last_inference_seconds, 4),
            "llm_token_speed": round(monitor.token_speed, 2),
            "llm_gpu_mem_total_mb": monitor.gpu_mem_total_mb,
            "llm_gpu_mem_used_mb": monitor.gpu_mem_used_mb,
            "llm_gpu_mem_usage_pct": monitor.gpu_mem_usage_pct,
            "llm_request_queue_length": monitor.request_queue_length,
            "llm_error_total": monitor.error_total,
            "llm_total_input_tokens": monitor.total_input_tokens,
            "llm_total_output_tokens": monitor.total_output_tokens,
            "llm_total_inference_count": monitor.total_inference_count,
        }

@app.get("/health")
def health_check():
    return {
        "status": "healthy" if monitor.service_health else "unhealthy",
        "model_loaded": monitor.model_loaded,
        "gpu_available": torch.cuda.is_available()
    }

@app.post("/chat")
def chat(question: str):
    if model is None:
        with metric_lock:
            monitor.error_total += 1
        raise HTTPException(status_code=500, detail="模型未加载")
    # 请求入队计数
    with queue_lock:
        monitor.request_queue_length += 1
    start_time = time.time()
    try:
        input_ids = tokenizer.encode(question)
        input_tokens = len(input_ids)
        with metric_lock:
            monitor.total_input_tokens += input_tokens
        # ChatGLM原生对话生成
        response, _ = model.chat(tokenizer, question, history=[])
        output_ids = tokenizer.encode(response)
        output_tokens = len(output_ids)
        # 计算吞吐与耗时指标
        cost = time.time() - start_time
        with metric_lock:
            monitor.last_inference_seconds = cost
            monitor.last_ttft_seconds = cost * 0.2
            monitor.token_speed = output_tokens / cost if cost > 0 else 0
            monitor.total_output_tokens += output_tokens
            monitor.total_inference_count += 1
        return {"question": question, "answer": response}
    except Exception as e:
        with metric_lock:
            monitor.error_total += 1
        raise HTTPException(status_code=500, detail=f"推理出错: {str(e)}")
    finally:
        # 请求完成出队
        with queue_lock:
            monitor.request_queue_length -= 1

import uuid
# 生成OpenAI格式对话ID
def _build_id():
    return "chatcmpl-" + uuid.uuid4().hex[:24]
# 构造非流式标准返回体
def _build_response(content, req_id, model_name, input_tokens, output_tokens, created):
    return {
        "id": req_id,
        "object": "chat.completion",
        "created": created,
        "model": model_name,
        "choices": [{
            "index": 0,
            "message": {"role": "assistant", "content": content},
            "finish_reason": "stop"
        }],
        "usage": {
            "prompt_tokens": input_tokens,
            "completion_tokens": output_tokens,
            "total_tokens": input_tokens + output_tokens
        }
    }
@app.post("/v1/chat/completions")
def chat_completions(payload: dict):
    if model is None:
        with metric_lock:
            monitor.error_total += 1
        raise HTTPException(status_code=500, detail="模型未加载")
    messages = payload.get("messages", [])
    stream = payload.get("stream", False)
    if not messages:
        raise HTTPException(status_code=400, detail="messages不能为空")
    # 拼接对话上下文Prompt
    prompt = ""
    for msg in messages:
        role = msg.get("role", "user")
        content = msg.get("content", "")
        if role == "user":
            prompt += f"用户：{content}\n"
        elif role == "assistant":
            prompt += f"助手：{content}\n"
    with queue_lock:
        monitor.request_queue_length += 1
    start_time = time.time()
    req_id = _build_id()
    created = int(start_time)
    input_ids = tokenizer.encode(prompt)
    input_tokens = len(input_ids)
    with metric_lock:
        monitor.total_input_tokens += input_tokens
    # ========== 普通模式：一次性返回完整结果 ==========
    if not stream:
        try:
            resp, _ = model.chat(tokenizer, prompt, history=[])
            output_ids = tokenizer.encode(resp)
            output_tokens = len(output_ids)
            cost = time.time() - start_time
            with metric_lock:
                monitor.last_inference_seconds = cost
                monitor.last_ttft_seconds = cost * 0.2
                monitor.token_speed = output_tokens / cost if cost > 0 else 0
                monitor.total_output_tokens += output_tokens
                monitor.total_inference_count += 1
            return _build_response(resp, req_id, model_name, input_tokens, output_tokens, created)
        except Exception as e:
            with metric_lock:
                monitor.error_total += 1
            raise HTTPException(500, str(e))
        finally:
            with queue_lock:
                monitor.request_queue_length -= 1
    # ========== SSE流式模式：逐token实时输出 ==========
    def stream_generator():
        total_output_token = 0
        ttft_recorded = False
        full_resp = ""
        try:
            for partial, _ in model.stream_chat(tokenizer, prompt, history=[]):
                # 首次输出记录首Token耗时TTFT
                if not ttft_recorded:
                    with metric_lock:
                        monitor.last_ttft_seconds = time.time() - start_time
                    ttft_recorded = True
                delta_text = partial.replace(full_resp, "")
                full_resp = partial
                token_cnt = len(tokenizer.encode(delta_text))
                total_output_token += token_cnt
                # SSE分片输出
                chunk = {
                    "id": req_id,
                    "object": "chat.completion.chunk",
                    "created": created,
                    "model": model_name,
                    "choices": [{
                        "index": 0,
                        "delta": {"content": delta_text},
                        "finish_reason": None
                    }]
                }
                yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n"
            # 流式结束统计整体吞吐指标
            total_cost = time.time() - start_time
            with metric_lock:
                monitor.last_inference_seconds = total_cost
                monitor.token_speed = total_output_token / total_cost if total_cost > 0 else 0
                monitor.total_output_tokens += total_output_token
                monitor.total_inference_count += 1
            # 结束分片
            final_chunk = {
                "id": req_id,
                "object": "chat.completion.chunk",
                "created": created,
                "model": model_name,
                "choices": [{
                    "index": 0,
                    "delta": {},
                    "finish_reason": "stop"
                }],
                "usage": {
                    "prompt_tokens": input_tokens,
                    "completion_tokens": total_output_token,
                    "total_tokens": input_tokens + total_output_token
                }
            }
            yield f"data: {json.dumps(final_chunk, ensure_ascii=False)}\n\n"
            yield "data: [DONE]\n\n"
        except Exception as e:
            with metric_lock:
                monitor.error_total += 1
            err_chunk = {"error": {"message": str(e)}}
            yield f"data: {json.dumps(err_chunk, ensure_ascii=False)}\n\n"
        finally:
            with queue_lock:
                monitor.request_queue_length -= 1
    return StreamingResponse(stream_generator(), media_type="text/event-stream")

if __name__ == "__main__":
    print("="*60)
    print("ChatGLM3-6B 对话服务启动成功")
    print(f"前端页面地址：http://0.0.0.0:8000")
    print(f"监控指标地址：http://0.0.0.0:8000/monitor")
    print(f"健康检查地址：http://0.0.0.0:8000/health")
    print("="*60)
    # 监听所有网卡，端口8000，日志级别info
    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")