#!/usr/bin/env python3
"""
VetNet DAPT 推理服务 — DAPT 微调 Qwen2.5-3B + FAISS RAG + LCPS 2.0

启动:
    VET_MODEL=./vet-qwen3b-v3-dapt KB_DIR=./knowledge_base_simple python server_gpu.py

API 端点:
    GET  /health              — 健康检查 + 模型状态
    POST /run                 — LCPS 推理 + SOAP (兼容旧接口)
    POST /api/v1/diagnose     — LCPS 诊断 (VetCopilot 标准接口)
    POST /chat                — 兽医自由问答
"""
import os, sys, time, json, logging
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
log = logging.getLogger('vetnet-gpu')

MODEL = None
RAG = None
MODEL_INFO = {}

class RunRequest(BaseModel):
    text: str
    species: str = "狗"

class ChatRequest(BaseModel):
    question: str
    species: str = "狗"

class DiagnoseRequest(BaseModel):
    """VetCopilot 标准诊断请求"""
    dialogue: str
    species: str = "狗"

@asynccontextmanager
async def lifespan(app: FastAPI):
    global MODEL, RAG, MODEL_INFO

    model_path = os.getenv("VET_MODEL", "./vet-qwen3b-v3-dapt")
    kb_dir = os.getenv("KB_DIR", "./knowledge_base_simple")

    log.info("=" * 60)
    log.info("VetNet DAPT 推理服务 v3.1 启动")
    log.info(f"模型: {model_path}")
    log.info(f"知识库: {kb_dir}")
    log.info(f"引擎: DAPT Qwen2.5-3B + LCPS 2.0 + FAISS RAG")
    log.info("=" * 60)

    MODEL_INFO = {
        "name": "vet-qwen3b-v3-dapt",
        "base_model": "Qwen2.5-3B-Instruct",
        "training": "DAPT (Domain-Adaptive Pretraining)",
        "engine": "LCPS 2.0",
    }

    t0 = time.time()

    # 加载 RAG
    try:
        from model.rag_faiss import FAISSRetriever
        RAG = FAISSRetriever(kb_dir=kb_dir, device="cuda")
        RAG.load()
        log.info(f"RAG 加载完成 ({RAG.get_stats()})")
    except Exception as e:
        log.warning(f"RAG 加载失败 (知识库文件缺失): {e}")
        RAG = None

    # 加载 DAPT 模型
    from model.qwen_local import QwenLocalClient
    MODEL = QwenLocalClient(
        model_path=model_path,
        temperature=float(os.getenv("VET_TEMPERATURE", "0.1")),   # 降低温度减少误诊随机性
        max_tokens=int(os.getenv("VET_MAX_TOKENS", "8192")),
        load_in_4bit=False,
        device_map="auto",
        seed=int(os.getenv("VET_SEED", "42")),   # 固定 seed 保证可复现
    )
    MODEL._ensure_loaded()

    MODEL_INFO["load_time_ms"] = int((time.time() - t0) * 1000)
    log.info(f"模型加载完成 ({MODEL_INFO['load_time_ms']}ms)")

    yield

    if MODEL:
        MODEL.unload()
    log.info("服务关闭")


app = FastAPI(title="VetNet GPU v3.0", lifespan=lifespan)
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])


@app.get("/health")
def health():
    return {
        "status": "ok",
        "model_loaded": MODEL is not None and MODEL.ping() if MODEL else False,
        "model_info": MODEL_INFO,
        "rag": RAG.get_stats() if RAG else {"ready": False, "vectors": 0, "chunks": 0},
    }


@app.post("/run")
def run(req: RunRequest):
    """LCPS 推理 + SOAP 生成 (兼容旧接口)"""
    return _do_diagnose(req.text, req.species)


@app.post("/api/v1/diagnose")
def diagnose(req: DiagnoseRequest):
    """VetCopilot 标准 LCPS 诊断接口"""
    return _do_diagnose(req.dialogue, req.species)


def _do_diagnose(text: str, species: str) -> dict:
    """统一诊断逻辑"""
    if not MODEL:
        raise HTTPException(503, "模型未加载")

    t0 = time.time()
    try:
        from model.lcps import lcps_single_shot, generate_soap_from_lcps

        rag = RAG if RAG and RAG._ready else None
        result = lcps_single_shot(text, species, MODEL, rag)
        soap = generate_soap_from_lcps(result["lcps_chain"], text, MODEL)

        elapsed = int((time.time() - t0) * 1000)
        return {
            "success": True,
            "diagnosis": result.get("primary_diagnosis", ""),
            "primary_diagnosis": result.get("primary_diagnosis", ""),
            "differential_list": result.get("differential_list", []),
            "must_not_miss": result.get("must_not_miss", []),
            "soap": soap.get("soap", {}),
            "lcps_chain": result.get("lcps_chain", {}),
            "summary": result.get("prioritisation_rationale", ""),
            "kb_used": rag is not None,
            "engine": "dapt-lcps-v3.1",
            "model": MODEL_INFO.get("name", "vet-qwen3b-v3-dapt"),
            "total_ms": elapsed,
            "elapsed_ms": elapsed,
            "error": "",
        }
    except Exception as e:
        log.exception("推理失败")
        return {
            "success": False,
            "diagnosis": "",
            "primary_diagnosis": "",
            "differential_list": [],
            "must_not_miss": [],
            "soap": {},
            "lcps_chain": {},
            "summary": "",
            "engine": "dapt-lcps-v3.1",
            "model": MODEL_INFO.get("name", ""),
            "elapsed_ms": int((time.time() - t0) * 1000),
            "error": str(e),
        }


@app.post("/chat")
def chat(req: ChatRequest):
    if not MODEL:
        raise HTTPException(503, "模型未加载")
    try:
        kb_context = RAG.retrieve(f"{req.species} {req.question[:300]}") if RAG else ""
        system = "你是一位经验丰富的兽医专家。请用专业但易懂的语言回答以下问题。"
        if kb_context:
            system += f"\n\n参考知识:\n{kb_context}"
        answer = MODEL.chat(system, f"物种: {req.species}\n问题: {req.question}")
        return {"answer": answer, "kb_used": bool(kb_context)}
    except Exception as e:
        log.exception("问答失败")
        raise HTTPException(500, str(e))


if __name__ == "__main__":
    import uvicorn
    port = int(os.getenv("PORT", "8000"))
    uvicorn.run(app, host="0.0.0.0", port=port)