#!/usr/bin/env python3
"""本地 GPU 推理统一入口 v3 — LLM (llama-server Vulkan 4B) + 向量嵌入
用法:
  gpu_infer.py llm "问题" [max_tokens]            # LLM 对话（4B 常驻服务 llama-server-4b）
  gpu_infer.py embed "文本"                       # 向量嵌入（bge-embed 服务）
  gpu_infer.py health                             # 健康检查
  gpu_infer.py bench                              # 性能基准

2026-09-07 牧尘铁律: 3B/7B 禁止使用。本机只部署 llama-server-4b (Qwen3.5-4B Q4_K_M,
GPU -ngl 99: VRAM ~3.1GB + host ~480MB)。3B 曾 CPU 推理 RSS 7.6GB 内存黑洞。
"""
import os
import sys
import time
import json
import subprocess
import urllib.request

LLAMA_DIR = "/tmp/llama-vulkan/llama-b10679"
EMBED_URL = "http://localhost:8000/v1/embeddings"
LLM_URL = "http://127.0.0.1:8080/v1/chat/completions"
LLAMA_SVC = "llama-server-4b"


def llm_chat(prompt, max_tokens=100, model="4b"):
    """LLM 对话（走 llama-server OpenAI 兼容 API，4B 常驻服务）"""
    data = json.dumps({
        "model": "qwen" + model,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "temperature": 0.7,
    }).encode()
    req = urllib.request.Request(LLM_URL, data=data, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=300) as resp:
        result = json.loads(resp.read())
    return result["choices"][0]["message"]["content"] if result.get("choices") else None


def embed(text):
    """向量嵌入（bge-embed 服务）"""
    if isinstance(text, str):
        text = [text]
    data = json.dumps({"input": text}).encode()
    req = urllib.request.Request(EMBED_URL, data=data, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=60) as resp:
        result = json.loads(resp.read())
    return [d["embedding"] for d in result["data"]]


def health():
    """健康检查"""
    print("=== GPU 健康检查 ===")
    r = subprocess.run(["nvidia-smi", "--query-gpu=name,memory.used,memory.total,utilization.gpu", "--format=csv,noheader"], capture_output=True, text=True)
    print(f"GPU: {r.stdout.strip()}")

    # bge-embed
    try:
        req = urllib.request.Request("http://localhost:8000/health")
        with urllib.request.urlopen(req, timeout=5) as resp:
            d = json.loads(resp.read())
            print(f"bge-embed: ✅ {d['providers']}")
    except Exception as e:
        print(f"bge-embed: ❌ {e}")

    # llama-server-4b
    try:
        req = urllib.request.Request("http://127.0.0.1:8080/health")
        with urllib.request.urlopen(req, timeout=5) as resp:
            d = json.loads(resp.read())
        cur = "4B"
        print(f"llama-server-4b {cur}: ✅ {d.get('status', 'ok')}")
    except Exception as e:
        print(f"llama-server-4b: ❌ {e}")

    # 4B 服务状态
    r = subprocess.run(["systemctl", "--user", "is-active", LLAMA_SVC], capture_output=True, text=True)
    print(f"llama-server-4b service: {r.stdout.strip()}")


def bench():
    """性能基准"""
    print("=== GPU 推理性能基准 ===\n")
    print("--- LLM 4B (llama-server-4b) ---")
    prompts = ["你好", "用一句话解释什么是过拟合"]
    for p in prompts:
        llm_chat(p, max_tokens=20)
        start = time.time()
        result = llm_chat(p, max_tokens=50)
        elapsed = time.time() - start
        print(f"  {p[:15]}: {elapsed:.2f}s -> {result[:40]}...")

    print("\n--- 向量嵌入 (bge-m3 INT8) ---")
    start = time.time()
    emb = embed("测试向量嵌入性能")
    elapsed = time.time() - start
    print(f"  维度 {len(emb[0])}: {elapsed:.3f}s")


def main():
    if len(sys.argv) < 2:
        print(__doc__)
        return
    cmd = sys.argv[1]

    if cmd == "llm":
        if len(sys.argv) < 3:
            print("用法: gpu_infer.py llm \"问题\" [max_tokens]")
            return
        prompt = sys.argv[2]
        max_tokens = int(sys.argv[3]) if len(sys.argv) > 3 else 100
        start = time.time()
        result = llm_chat(prompt, max_tokens)
        elapsed = time.time() - start
        if result:
            print(f"输出: {result}")
            print(f"耗时: {elapsed:.2f}s")

    elif cmd == "embed":
        if len(sys.argv) < 3:
            print("用法: gpu_infer.py embed \"文本\"")
            return
        start = time.time()
        result = embed(sys.argv[2])
        elapsed = time.time() - start
        print(f"维度: {len(result[0])}, 耗时: {elapsed:.3f}s")

    elif cmd == "health":
        health()

    elif cmd == "bench":
        bench()

    else:
        print(__doc__)


if __name__ == "__main__":
    main()