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

7B 按需启动（临时停 3B，跑完恢复）：
  gpu_infer.py llm "问题" 100 7b
"""
import os
import sys
import time
import json
import subprocess
import urllib.request

LLAMA_DIR = "/tmp/llama-vulkan/llama-b10679"
MODEL_3B = os.path.expanduser("~/models/Qwen-Qwen2.5-3B-Instruct-GGUF/qwen2.5-3b-instruct-q4_k_m.gguf")
MODEL_7B = os.path.expanduser("~/models/Qwen-Qwen2.5-7B-Instruct-GGUF/qwen2.5-7b-instruct-q3_k_m.gguf")
EMBED_URL = "http://localhost:8000/v1/embeddings"
LLM_URL = "http://127.0.0.1:8080/v1/chat/completions"


def llm_chat(prompt, max_tokens=100, model="3b"):
    """LLM 对话（走 llama-server OpenAI 兼容 API，当前常驻模型）"""
    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 _run_7b_once(prompt, max_tokens):
    """7B 按需运行：临时停 3B → 跑 7B → 恢复 3B"""
    print("(7B 模式：临时停用 3B 服务...)", file=sys.stderr)
    subprocess.run(["systemctl", "--user", "stop", "llama-server"], check=False)
    try:
        env = os.environ.copy()
        env["LD_LIBRARY_PATH"] = LLAMA_DIR
        cmd = [
            os.path.join(LLAMA_DIR, "llama-cli"),
            "-m", MODEL_7B,
            "-p", prompt,
            "-n", str(max_tokens),
            "-st", "--no-display-prompt",
        ]
        r = subprocess.run(cmd, capture_output=True, text=True, timeout=600, env=env)
        # 提取生成文本
        out = r.stdout + r.stderr
        print(out)
    finally:
        subprocess.run(["systemctl", "--user", "start", "llama-server"], check=False)
        time.sleep(5)


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 (3B 或 7B)
    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 = "?"
            if os.path.exists(MODEL_3B) and os.path.exists(MODEL_7B):
                r = subprocess.run(["systemctl", "--user", "is-active", "llama-server"], capture_output=True, text=True)
                if "active" in r.stdout:
                    cur = "3B"
                else:
                    r = subprocess.run(["systemctl", "--user", "is-active", "llama-server-7b"], capture_output=True, text=True)
                    if "active" in r.stdout:
                        cur = "7B"
            print(f"llama-server {cur}: ✅ {d.get('status', 'ok')}")
    except Exception as e:
        print(f"llama-server: ❌ {e}")

    # 模型文件
    for label, path in [("3B Q4", MODEL_3B), ("7B Q3", MODEL_7B)]:
        if os.path.exists(path):
            size = os.path.getsize(path) // 1024 // 1024
            print(f"{label}: ✅ {size}MB")
        else:
            print(f"{label}: ❌ 缺失")


def bench():
    """性能基准"""
    print("=== GPU 推理性能基准 ===\n")
    print("--- LLM 3B (llama-server) ---")
    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] [3b|7b]")
            return
        prompt = sys.argv[2]
        max_tokens = int(sys.argv[3]) if len(sys.argv) > 3 else 100
        model = sys.argv[4] if len(sys.argv) > 4 else "3b"
        start = time.time()
        result = llm_chat(prompt, max_tokens, model)
        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()