#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
BGEM3 inference server (FastAPI) with robust GPU / CPU management.

Launch examples:
    python server.py                       # 自动选卡 / 自动降级
    python server.py --device 1            # 固定用第 1 张 GPU
    CUDA_VISIBLE_DEVICES=0,1 python server.py
"""

import argparse
import logging
import os
import sys
import time
from typing import List, Union
import multiprocessing as mp

import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoTokenizer
from FlagEmbedding import BGEM3FlagModel

mp.set_start_method("spawn", force=True)

# -----------------------------------------------------------------------------#
# Config
# -----------------------------------------------------------------------------#
MODEL_PATH = "model/bge-m3"          # 按需改成你的权重路径
MODEL_VRAM_MB = int(os.getenv("MODEL_VRAM_MB", "8000"))   # bge-m3-large fp32 ≈ 8 GiB
POST_LOAD_GAP_MB = 192     
SAFE_MIN_FREE_MB = MODEL_VRAM_MB + POST_LOAD_GAP_MB      # == 8192 MB

# -----------------------------------------------------------------------------#
# Logging
# -----------------------------------------------------------------------------#
logger = logging.getLogger("bge-m3-server")
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)

# -----------------------------------------------------------------------------#
# GPU memory helpers (NVML → torch fallback)
# -----------------------------------------------------------------------------#
try:
    import pynvml

    pynvml.nvmlInit()
    _USE_NVML = True
except Exception:
    _USE_NVML = False
    logger.warning("pynvml 不可用，将使用 torch.cuda.mem_get_info() 探测显存")


def _gpu_mem_info(idx: int) -> tuple[int, int]:
    """Return (free_bytes, total_bytes) for GPU `idx`."""
    if _USE_NVML:
        handle = pynvml.nvmlDeviceGetHandleByIndex(idx)
        mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
        return mem.free, mem.total
    # torch fallback
    torch.cuda.set_device(idx)
    free, total = torch.cuda.mem_get_info(idx)
    return free, total


# -----------------------------------------------------------------------------#
# Precision & model loader
# -----------------------------------------------------------------------------#
def _choose_precision(device: str) -> str:
    """Return 'fp16'|'bf16'|'fp32'."""
    if not device.startswith("cuda"):
        return "fp32"

    major, _ = torch.cuda.get_device_capability(device)
    if major >= 8:  # Ampere/Hopper
        return "fp16"
    if major >= 7:
        return "bf16"
    return "fp32"


def load_model(device: str):
    precision = _choose_precision(device)
    use_fp16  = precision == "fp16"
    logger.info("Loading BGEM3 on %s (%s)", device, precision)

    if device == "cpu":
        # CPU 路径：彻底摘掉 CUDA
        os.environ["CUDA_VISIBLE_DEVICES"] = ""
        torch.cuda.is_available = lambda: False
        torch.cuda.device_count = lambda: 0
    else:
        # GPU 路径：只暴露本 worker 选中的那张卡
        # 例：device == "cuda:3" → 只让当前进程看到 GPU 3
        idx = device.split(":")[1]
        os.environ["CUDA_VISIBLE_DEVICES"] = idx
        # device_count 现在返回 1，BGEM3 只会在这张卡上建 1 个子进程
        torch.cuda.device_count = lambda: 1

    mdl = BGEM3FlagModel(MODEL_PATH, use_fp16=use_fp16, device=device)

    # 不再包 DataParallel；每个 worker 单卡即可
    return mdl, precision

# -----------------------------------------------------------------------------#
# Auto-select device (startup)
# -----------------------------------------------------------------------------#
def auto_select_and_load() -> tuple:
    """
    1. 过滤掉空闲显存 < MODEL_VRAM_MB 的 GPU
    2. 按空闲显存降序依次尝试加载
    3. 载入后再次检查：若剩余 < POST_LOAD_GAP_MB → 视为失败
    4. 若全部 GPU 不满足 → CPU
    """
    if not torch.cuda.is_available():
        logger.info("No GPU detected → CPU")
        return (*load_model("cpu"), "cpu")

    # 收集候选卡 (free_MB, idx)
    candidates = []
    for idx in range(torch.cuda.device_count()):
        free_mb = _gpu_mem_info(idx)[0] // 2**20
        if free_mb >= MODEL_VRAM_MB:                 # 至少能放下权重
            candidates.append((free_mb, idx))

    if not candidates:
        logger.warning("All GPUs free_mem < %d MB → CPU", MODEL_VRAM_MB)
        return (*load_model("cpu"), "cpu")

    # 空闲显存从高到低
    for free_mb, idx in sorted(candidates, reverse=True):
        dev = f"cuda:{idx}"
        try:
            logger.info("Trying %s (free=%d MB)", dev, free_mb)
            mdl, prec = load_model(dev)

            remain_mb = _gpu_mem_info(idx)[0] // 2**20
            if remain_mb < POST_LOAD_GAP_MB:
                raise RuntimeError(
                    f"post-load free {remain_mb} MB < {POST_LOAD_GAP_MB} MB")

            return mdl, prec, dev                    # 成功
        except RuntimeError as e:
            logger.warning("%s unusable (%s) → next", dev, e)
            torch.cuda.empty_cache()

    logger.warning("No suitable GPU left → CPU fallback")
    return (*load_model("cpu"), "cpu")

# -----------------------------------------------------------------------------#
# CLI
# -----------------------------------------------------------------------------#
parser = argparse.ArgumentParser()
parser.add_argument(
    "--device", help="GPU index (e.g. 0) or 'cpu'; overrides auto-selection"
)
args, _ = parser.parse_known_args()

if args.device is not None:
    # Forced path
    if args.device.lower() == "cpu":
        DEVICE = "cpu"
    else:
        DEVICE = f"cuda:{int(args.device)}" if torch.cuda.is_available() else "cpu"
    model, PRECISION = load_model(DEVICE)
else:
    # Auto path with VRAM check
    model, PRECISION, DEVICE = auto_select_and_load()

tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)

# -----------------------------------------------------------------------------#
# FastAPI
# -----------------------------------------------------------------------------#
app = FastAPI()
logger.info("Using SAFE_MIN_FREE_MB = %d MB", SAFE_MIN_FREE_MB)


def _warm_worker(t, q):
    try:
        _ = model.encode(t, return_dense=True)
        q.put("ok")
    except Exception as e:
        q.put(str(e))

# ② -------- FastAPI 启动预热 --------
@app.on_event("startup")
def warm_up():
    logger.info("Warm-up on %s", DEVICE)
    try:
        _ = model.encode([
            "This is a warmup sentence used to initialize CUDA kernels and avoid latency spikes."
        ], return_dense=True)
        logger.info("Warm-up complete.")
    except Exception as e:
        logger.warning("Warm-up failed: %s", e)



class EmbeddingRequest(BaseModel):
    input: Union[str, List[str]]
    model: str = "text-embedding-bge-m3"


# ③ -------- _encode() 里 worker 调用 --------
def _worker(t, q):
    try:
        # out = model.encode(t, return_dense=True)   # GPU or CPU 均安全
        out = model.encode(t, return_dense=True)
        q.put(("ok", out))
    except Exception as e:
        q.put(("err", str(e)))


def _encode(texts: List[str]):
    try:
        return model.encode(texts, return_dense=True)
    except RuntimeError as e:
        if "out of memory" in str(e).lower() or "cuda error" in str(e).lower():
            logger.warning("GPU OOM → fallback to CPU: %s", str(e))
            torch.cuda.empty_cache()
            global CPU_MODEL_CACHE
            if CPU_MODEL_CACHE is None:
                CPU_MODEL_CACHE, _ = load_model("cpu")
            return CPU_MODEL_CACHE.encode(texts, return_dense=True)
        raise



@app.post("/v1/embeddings")
def create_embedding(request: EmbeddingRequest):
    texts = [request.input] if isinstance(request.input, str) else request.input

    # Token stats
    enc = tokenizer(
        texts,
        padding=True,
        truncation=True,
        max_length=8192,
        return_tensors="pt",
    )
    prompt_tokens = int(enc["attention_mask"].sum().item())

    try:
        output = _encode(texts)
        embeddings = output["dense_vecs"]
    except Exception as e:
        logger.exception("Embedding failed")
        raise HTTPException(status_code=500, detail=str(e))

    return {
        "object": "list",
        "data": [
            {
                "object": "embedding",
                "index": i,
                "embedding": emb.tolist() if hasattr(emb, "tolist") else emb,
            }
            for i, emb in enumerate(embeddings)
        ],
        "model": request.model,
        "usage": {
            "prompt_tokens": prompt_tokens,
            "total_tokens": prompt_tokens,
        },
        "device": DEVICE,
        "precision": PRECISION,
    }


# -----------------------------------------------------------------------------#
# Entry-point for `python server.py`
# -----------------------------------------------------------------------------#
if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "server:app",
        host="0.0.0.0",
        port=int(os.getenv("PORT", 8000)),
        log_level="info",
        workers=1,  # multi-process → use gunicorn/uvicorn externally
    )
