#!/usr/bin/env python3
# DGX AI Factory - 14B vLLM Quick Benchmark
#
# 목적:
# - 현재 떠 있는 vLLM 14B champion을 간단히 점검합니다.
# - 학습하지 않습니다.
# - adapter를 이동하지 않습니다.
# - stable/rejected 승격을 하지 않습니다.
#
# 평가 성격:
# - 운영 전 quick benchmark입니다.
# - 기존 20문항 정식 benchmark를 대체하지 않습니다.
# - direct vLLM 모델 호출 품질을 빠르게 확인합니다.

import json
import time
import urllib.request
import urllib.error
from pathlib import Path
from datetime import datetime

ROOT = Path("/home/harness_user_1/dgx_ai_factory")
REPORT_DIR = ROOT / "reports"
REPORT_DIR.mkdir(parents=True, exist_ok=True)

ENDPOINT = "http://localhost:8000/v1/chat/completions"
MODEL = "dgx-14b-champion"

SYSTEM = (
    "한국어 존댓말로 답변하세요. 중국어를 쓰지 마세요. "
    "명령어와 설명을 구분하고, 불확실한 내용은 단정하지 마세요."
)

TESTS = [
    {
        "id": "q01_lora_adapter",
        "category": "lora",
        "prompt": "LoRA adapter 폴더에서 확인해야 하는 파일 2개와 stable/rejected 판단 순서를 설명해주세요.",
        "must": ["adapter_config.json", "adapter_model.safetensors", "smoke", "benchmark"],
        "bad": ["yaml", ".pt", "Loosely Regularized", "중국", "试", "不"],
    },
    {
        "id": "q02_systemd_217_user",
        "category": "systemd",
        "prompt": "systemd 서비스가 status=217/USER로 실패할 때 확인 순서와 수정 명령어를 알려주세요.",
        "must": ["User=", "id", "getent", "daemon-reload", "restart"],
        "bad": ["중국", "试", "不"],
    },
    {
        "id": "q03_dataset_quality",
        "category": "dataset",
        "prompt": "학습용 jsonl 데이터셋을 만들 때 중복, CJK 누수, template leak, category 균형을 점검하는 절차를 설명해주세요.",
        "must": ["jsonl", "중복", "CJK", "template", "category"],
        "bad": ["중국어로", "试", "不"],
    },
    {
        "id": "q04_fastapi_curl",
        "category": "fastapi",
        "prompt": "FastAPI health check 엔드포인트를 만들고 curl로 확인하는 최소 예제를 보여주세요.",
        "must": ["FastAPI", "curl", "/health", "uvicorn"],
        "bad": ["중국", "试", "不"],
    },
    {
        "id": "q05_vllm_port",
        "category": "vllm",
        "prompt": "vLLM OpenAI compatible server가 8000 포트에서 정상인지 확인하는 명령어와 판단 기준을 알려주세요.",
        "must": ["ss", "8000", "/v1/models", "curl"],
        "bad": ["중국", "试", "不"],
    },
    {
        "id": "q06_docker_openwebui",
        "category": "docker",
        "prompt": "OpenWebUI Docker 컨테이너가 떠 있는지 확인하고 로그를 보는 명령어를 알려주세요.",
        "must": ["docker ps", "docker logs"],
        "bad": ["중국", "试", "不"],
    },
    {
        "id": "q07_linux_process",
        "category": "linux",
        "prompt": "GPU 모델 프로세스를 찾고 종료하기 전에 확인해야 할 Linux 명령어를 알려주세요.",
        "must": ["ps", "grep", "nvidia-smi", "kill"],
        "bad": ["중국", "试", "不"],
    },
    {
        "id": "q08_jsonl_validate",
        "category": "jsonl",
        "prompt": "jsonl 파일이 줄 단위 JSON으로 정상인지 Python으로 검증하는 방법을 알려주세요.",
        "must": ["json.loads", "line", "bad", "python"],
        "bad": ["중국", "试", "不"],
    },
    {
        "id": "q09_cuda_oom",
        "category": "cuda",
        "prompt": "CUDA OOM이 났을 때 원인 분석 순서와 줄일 수 있는 설정값을 설명해주세요.",
        "must": ["OOM", "max-model-len", "max-num-seqs", "gpu-memory-utilization"],
        "bad": ["중국", "试", "不"],
    },
    {
        "id": "q10_korean_style",
        "category": "korean_style",
        "prompt": "앞으로 명령어를 줄 때 터미널에 긴 명령을 직접 붙이지 않고 파일로 실행하는 방식으로 안내해야 하는 이유를 설명해주세요.",
        "must": ["파일", "실행", "로그", "명령어"],
        "bad": ["해라", "해", "중국", "试", "不"],
    },
]

def post_chat(prompt: str) -> tuple[str, float, dict]:
    payload = {
        "model": MODEL,
        "messages": [
            {"role": "system", "content": SYSTEM},
            {"role": "user", "content": prompt},
        ],
        "temperature": 0,
        "max_tokens": 500,
    }
    data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    req = urllib.request.Request(
        ENDPOINT,
        data=data,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    t0 = time.time()
    try:
        with urllib.request.urlopen(req, timeout=240) as resp:
            raw = resp.read().decode("utf-8", errors="replace")
            obj = json.loads(raw)
    except urllib.error.URLError as e:
        node-7.example.invalid RuntimeError(f"Request failed: {e}") from e
    latency = time.time() - t0
    answer = obj["choices"][0]["message"]["content"]
    return answer, latency, obj

def score_answer(answer: str, must: list[str], bad: list[str]) -> dict:
    lower = answer.lower()
    missing = []
    for m in must:
        if m.lower() not in lower:
            missing.append(m)

    bad_hits = []
    for b in bad:
        if b.lower() in lower:
            bad_hits.append(b)

    score = 100
    score -= 12 * len(missing)
    score -= 18 * len(bad_hits)
    if len(answer.strip()) < 80:
        score -= 15
    score = max(0, min(100, score))

    passed = score >= 70 and not bad_hits
    strong = score >= 85 and not bad_hits
    return {
        "score": score,
        "passed": passed,
        "strong": strong,
        "missing_must": missing,
        "bad_hits": bad_hits,
    }

def main():
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    json_report = REPORT_DIR / f"14b_vllm_tune046_quick_benchmark_{ts}.json"
    jsonl_report = REPORT_DIR / f"14b_vllm_tune046_quick_benchmark_{ts}.jsonl"
    low_report = REPORT_DIR / f"14b_vllm_tune046_quick_benchmark_low_{ts}.jsonl"

    rows = []
    print("==== 14B VLLM QUICK BENCHMARK START ====")
    print("model:", MODEL)
    print("endpoint:", ENDPOINT)
    print("tests:", len(TESTS))
    print("note: no training, no adapter move, no stable promotion")
    print()

    for i, t in enumerate(TESTS, 1):
        print(f"==== {i}/{len(TESTS)} {t['id']} ({t['category']}) ====")
        try:
            answer, latency, raw = post_chat(t["prompt"])
            score = score_answer(answer, t["must"], t["bad"])
            row = {
                "id": t["id"],
                "category": t["category"],
                "prompt": t["prompt"],
                "answer": answer,
                "latency_sec": round(latency, 3),
                **score,
            }
        except Exception as e:
            row = {
                "id": t["id"],
                "category": t["category"],
                "prompt": t["prompt"],
                "answer": "",
                "latency_sec": None,
                "score": 0,
                "passed": False,
                "strong": False,
                "missing_must": t["must"],
                "bad_hits": [],
                "error": repr(e),
            }

        rows.append(row)
        print("score:", row["score"], "passed:", row["passed"], "latency:", row["latency_sec"])
        if row.get("missing_must"):
            print("missing:", row["missing_must"])
        if row.get("bad_hits"):
            print("bad_hits:", row["bad_hits"])
        print()

    avg = round(sum(r["score"] for r in rows) / len(rows), 2)
    pass_count = sum(1 for r in rows if r["passed"])
    strong_count = sum(1 for r in rows if r["strong"])
    latencies = [r["latency_sec"] for r in rows if isinstance(r["latency_sec"], (int, float))]
    avg_latency = round(sum(latencies) / len(latencies), 3) if latencies else None

    category_avg = {}
    for r in rows:
        category_avg.setdefault(r["category"], []).append(r["score"])
    category_avg = {k: round(sum(v) / len(v), 2) for k, v in sorted(category_avg.items())}

    if avg >= 85 and pass_count >= 9:
        decision = "PASS_QUICK_BENCHMARK"
    elif avg >= 75 and pass_count >= 7:
        decision = "REVIEW_BEFORE_FULL_BENCHMARK"
    else:
        decision = "FAIL_OR_ROUTE_NEEDED"

    summary = {
        "timestamp": ts,
        "model": MODEL,
        "endpoint": ENDPOINT,
        "test_count": len(rows),
        "average_score": avg,
        "pass_70_plus": f"{pass_count}/{len(rows)}",
        "strong_85_plus": f"{strong_count}/{len(rows)}",
        "average_latency_sec": avg_latency,
        "category_avg": category_avg,
        "decision": decision,
        "json_report": str(json_report),
        "jsonl_report": str(jsonl_report),
        "low_report": str(low_report),
        "note": "Direct vLLM model quick benchmark. Router/gateway benchmark should be run separately for operational route quality.",
    }

    json_report.write_text(json.dumps({"summary": summary, "rows": rows}, ensure_ascii=False, indent=2), encoding="utf-8")
    with jsonl_report.open("w", encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    with low_report.open("w", encoding="utf-8") as f:
        for r in rows:
            if r["score"] < 85 or r.get("bad_hits") or r.get("missing_must"):
                f.write(json.dumps(r, ensure_ascii=False) + "\n")

    print("==== FINAL 14B VLLM QUICK BENCHMARK SUMMARY ====")
    for k, v in summary.items():
        print(f"{k}: {v}")

if __name__ == "__main__":
    main()