#!/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

# -----------------------------------------------------------------------------#
# Config
# -----------------------------------------------------------------------------#
MODEL_PATH = "model/bge-m3"          # 按需改成你的权重路径
SAFE_MIN_FREE_MB = int(
    os.getenv("SAFE_MIN_FREE_MB", "16384")
)  # 启动时要求的最小空闲显存，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):
    """Instantiate model on `device`; return (model, precision)."""
    precision = _choose_precision(device)
    use_fp16 = precision == "fp16"
    logger.info("Loading BGEM3 on %s (%s)", device, precision)

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

    if device == "cpu":
        # 屏蔽 GPU，让后续 torch / BGEM3 都认不出 CUDA
        os.environ["CUDA_VISIBLE_DEVICES"] = ""

    # Simple DataParallel for multi-GPU inference
    if device.startswith("cuda") and torch.cuda.device_count() > 1:
        logger.info(
            "Wrapping model with torch.nn.DataParallel (%d GPUs)",
            torch.cuda.device_count(),
        )
        mdl = torch.nn.DataParallel(mdl)
    return mdl, precision


# -----------------------------------------------------------------------------#
# Auto-select device (startup)
# -----------------------------------------------------------------------------#
def auto_select_and_load(min_free_mb: int = 4096):
    """
    1. Gather GPUs with free_mem ≥ min_free_mb
    2. Sort by free_mem desc, attempt to load
    3. All fail → CPU
    Return (model, device, precision)
    """
    if not torch.cuda.is_available():
        logger.info("No GPU detected → CPU")
        return (*load_model("cpu"), "cpu")

    # Build candidate list
    candidates: list[tuple[int, int]] = []  # (free_MB, idx)
    for idx in range(torch.cuda.device_count()):
        free, _ = _gpu_mem_info(idx)
        free_mb = free // 2**20
        candidates.append((free_mb, idx))

    candidates = [c for c in candidates if c[0] >= min_free_mb]
    if not candidates:
        logger.warning("All GPUs free_mem < %d MB → CPU", min_free_mb)
        return (*load_model("cpu"), "cpu")

    candidates.sort(reverse=True)  # high free_mem first
    for free_mb, idx in candidates:
        dev = f"cuda:{idx}"
        try:
            logger.info("Trying %s (free=%d MB)", dev, free_mb)
            mdl, prec = load_model(dev)
            return mdl, prec, dev
        except RuntimeError as e:
            if "out of memory" in str(e).lower():
                logger.warning("%s OOM → next GPU", dev)
                torch.cuda.empty_cache()
                continue
            raise  # non-OOM error

    logger.warning("All GPUs failed → 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(SAFE_MIN_FREE_MB)

tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)

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

@app.on_event("startup")
async def warm_up_mp_pool():
    try:
        if DEVICE.startswith("cuda"):
            logger.info("Warm-up (GPU) → 建多进程池")
            _ = model.encode(["warmup"], return_dense=True)
        else:
            logger.info("Warm-up (CPU) → 单进程初始化")
            # ── 临时让库“以为”没有 GPU ────────────────────────────
            orig_cnt = torch.cuda.device_count
            torch.cuda.device_count = lambda: 0
            _ = model.encode(["warmup"], return_dense=True)  # 不传 num_processes
            torch.cuda.device_count = orig_cnt
            # ──────────────────────────────────────────────────────
    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"


def _encode(texts: List[str]):
    """
    单次请求：
      1. 子进程跑 GPU 推理；成功→返回
      2. 若子进程 OOM / CUDA Error → 同一次请求 fallback 到 CPU
    绝不改全局状态，其他并发请求不受影响
    """
    def _worker(t, q):
        try:
            if DEVICE.startswith("cuda"):
                out = model.encode(t, return_dense=True)          # GPU 正常跑
            else:
                # 临时屏蔽 GPU，单进程 CPU 推理
                orig_cnt = torch.cuda.device_count
                torch.cuda.device_count = lambda: 0
                out = model.encode(t, return_dense=True)          # 不传 num_processes
                torch.cuda.device_count = orig_cnt
            q.put(("ok", out))
        except Exception as e:
            q.put(("err", str(e)))

    q = mp.Queue()
    p = mp.Process(target=_worker, args=(texts, q))
    p.start()
    p.join(timeout=60)

    if not q.empty():
        status, payload = q.get()
        if status == "ok":
            return payload
        if "out of memory" in payload.lower() or "cuda error" in payload.lower():
            logger.warning("GPU OOM → 本次请求改走 CPU：%s", payload)
            torch.cuda.empty_cache()
            cpu_model, _ = load_model("cpu")
            return cpu_model.encode(texts, return_dense=True)
        raise RuntimeError(payload)

    raise RuntimeError("子进程异常退出，无返回")

# fallback_done = False  # prevent endless downgrade loop

# def _encode(texts: List[str]):
#     """Encode with single downgrade to CPU on OOM / CUDA failure."""
#     global model, DEVICE, PRECISION, fallback_done

#     try:
#         return model.encode(texts, return_dense=True)

#     except RuntimeError as err:
#         is_oom = "out of memory" in str(err).lower()
#         is_cuda_fail = "cuda error" in str(err).lower() or "device-side assert" in str(
#             err
#         ).lower()

#         if (is_oom or is_cuda_fail) and not fallback_done:
#             logger.error("GPU failure (%s). Falling back to CPU…", err)
#             fallback_done = True
#             torch.cuda.empty_cache()
#             DEVICE = "cpu"
#             model, PRECISION = load_model(DEVICE)
#             return model.encode(texts, return_dense=True)

#         raise  # second failure → propagate


@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
    )
