#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Agent cho máy GPU cộng đồng của hoclieu.org.

Chỉ gọi RA ngoài tới máy chủ (không mở cổng nào).
Chạy:  python3 hoclieu_agent.py --key hlk_xxx
"""
import argparse
import json
import re
import os
import shutil
import signal
import subprocess
import sys
import time
import urllib.error
import urllib.request

AGENT_VERSION = "0.1.0"
DEFAULT_SERVER = "https://hoclieu.org"
DEFAULT_LLM = "http://127.0.0.1:8000/v1"
DEFAULT_INTERVAL = 5
HEARTBEAT_EVERY = 30
SERVER_TIMEOUT = 30
LLM_TIMEOUT = 600
MAX_BACKOFF = 60


def _env(name, default=None):
    v = os.environ.get(name)
    return v if v not in (None, "") else default


def parse_args(argv=None):
    p = argparse.ArgumentParser(description="Agent GPU cộng đồng hoclieu.org")
    p.add_argument("--server", default=_env("HL_SERVER", DEFAULT_SERVER))
    p.add_argument("--key", default=_env("HL_KEY"))
    p.add_argument("--llm", default=_env("HL_LLM", DEFAULT_LLM))
    p.add_argument("--model", default=_env("HL_MODEL", ""))
    p.add_argument("--interval", type=float, default=_env("HL_INTERVAL", DEFAULT_INTERVAL))
    p.add_argument("--once", action="store_true",
                   help="heartbeat + xử lý tối đa 1 việc rồi thoát (kiểm thử)")
    return p.parse_args(argv)


def _run(cmd, timeout=15):
    try:
        r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
        return r.stdout or ""
    except Exception:
        return ""


def collect_info(llm_base):
    gpus = []
    driver = ""
    out = _run(["nvidia-smi", "--query-gpu=name,memory.total,driver_version",
                "--format=csv,noheader,nounits"])
    if out.strip():
        for line in out.strip().splitlines():
            parts = [x.strip() for x in line.split(",")]
            if len(parts) >= 3:
                try:
                    mem = int(parts[1])
                except ValueError:
                    mem = 0
                gpus.append({"name": parts[0][:80], "vramMb": mem})
                driver = parts[2]
    cuda = ""
    full = _run(["nvidia-smi"])
    for line in full.splitlines():
        if "CUDA Version" in line:
            cuda = line.split("CUDA Version")[-1].strip(" :|\t")
            break
    try:
        disk_free_gb = round(shutil.disk_usage(os.getcwd()).free / (1024 ** 3), 1)
    except Exception:
        disk_free_gb = 0
    models = _llm_models(llm_base)
    return {
        "agentVersion": AGENT_VERSION,
        "gpus": gpus,
        "driver": driver,
        "cuda": cuda,
        "diskFreeGb": disk_free_gb,
        "models": models,
    }


def _llm_models(llm_base):
    try:
        data = _http_json(llm_base.rstrip("/") + "/models", method="GET", timeout=10)
        out = []
        for m in (data.get("data") or []):
            mid = m.get("id")
            if mid:
                out.append(mid)
        return out
    except Exception:
        return []


def _http_json(url, method="POST", body=None, headers=None, timeout=SERVER_TIMEOUT):
    data = None
    if body is not None:
        data = json.dumps(body).encode("utf-8")
    req = urllib.request.Request(url, data=data, method=method)
    if headers:
        for k, v in headers.items():
            req.add_header(k, v)
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        raw = resp.read()
        if not raw:
            return {}
        return json.loads(raw.decode("utf-8"))


def server_call(server, key, op, extra=None):
    url = server.rstrip("/") + "/api/nodes/agent"
    body = {"op": op}
    if extra:
        body.update(extra)
    headers = {
        "Authorization": "Bearer " + key,
        "Content-Type": "application/json",
    }
    try:
        return _http_json(url, method="POST", body=body, headers=headers,
                          timeout=SERVER_TIMEOUT)
    except urllib.error.HTTPError as e:
        if e.code == 401:
            print("Key không hợp lệ hoặc đã bị thu hồi")
            sys.exit(2)
        raise
    except urllib.error.URLError as e:
        raise RuntimeError("network: %s" % e.reason)


def llm_chat(llm_base, model, messages, max_tokens, temperature):
    url = llm_base.rstrip("/") + "/chat/completions"
    body = {
        "model": model,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": temperature,
    }
    data = _http_json(url, method="POST", body=body,
                      headers={"Content-Type": "application/json"},
                      timeout=LLM_TIMEOUT)
    choices = data.get("choices") or []
    if not choices:
        raise RuntimeError("LLM không trả về lựa chọn")
    text = (choices[0].get("message") or {}).get("content", "") or ""
    # mô hình có "suy nghĩ" (Qwen3, DeepSeek-R1…): bỏ phần <think>…</think>, chỉ gửi câu trả lời
    return re.sub(r"<think>.*?</think>", "", text, flags=re.S).strip()


def resolve_model(llm_base, model):
    if model:
        return model
    models = _llm_models(llm_base)
    return models[0] if models else ""


def _sleep(seconds, stop):
    end = time.time() + seconds
    while not stop["flag"]:
        remaining = end - time.time()
        if remaining <= 0:
            break
        time.sleep(min(0.2, remaining))


def main(argv=None):
    args = parse_args(argv)
    if not args.key:
        print("Thiếu --key / HL_KEY (dạng hlk_...)")
        sys.exit(2)
    print("Agent GPU hoclieu.org v%s | server=%s | llm=%s | key=%s..."
          % (AGENT_VERSION, args.server, args.llm, args.key[:8]))

    stop = {"flag": False}

    def _sig(_s, _f):
        stop["flag"] = True

    try:
        signal.signal(signal.SIGINT, _sig)
        signal.signal(signal.SIGTERM, _sig)
    except Exception:
        pass

    model = resolve_model(args.llm, args.model)
    if not model:
        print("Cảnh báo: chưa xác định được model LLM (HL_MODEL trống và "
              "GET /models rỗng).")

    last_hb = 0.0
    backoff = 1
    jobs_done = 0

    while not stop["flag"]:
        now = time.time()
        try:
            if now - last_hb >= HEARTBEAT_EVERY:
                info = collect_info(args.llm)
                resp = server_call(args.server, args.key, "heartbeat", {"info": info})
                last_hb = now
                risk = resp.get("risk") or resp.get("riskLevel")
                if risk:
                    print("[heartbeat] mức rủi ro: %s" % risk)
            resp = server_call(args.server, args.key, "next")
            backoff = 1
            job = resp.get("job")  # {id, kind, payload: {messages, max_tokens, temperature}}
            if job:
                payload = job.get("payload") or {}
                t0 = time.time()
                ok = True
                result = ""
                try:
                    result = llm_chat(
                        args.llm,
                        model,
                        payload.get("messages") or [],
                        payload.get("max_tokens", 512),
                        payload.get("temperature", 0.2),
                    )
                except Exception as e:
                    ok = False
                    result = str(e)[:200]
                ms = int((time.time() - t0) * 1000)
                server_call(args.server, args.key, "result", {
                    "jobId": job.get("id"), "ok": ok, "result": result, "ms": ms,
                })
                jobs_done += 1
                print("[job %s] ok=%s ms=%d" % (job.get("id"), ok, ms))
                if args.once:
                    break
            else:
                if args.once:  # --once: không có việc thì thoát luôn
                    print("Không có việc.")
                    break
                _sleep(args.interval, stop)
        except SystemExit:
            raise
        except Exception as e:
            print("Lỗi: %s (thử lại sau %ds)" % (e, backoff))
            _sleep(backoff, stop)
            backoff = min(backoff * 2, MAX_BACKOFF)
            if args.once:
                break

    print("Agent dừng.")


if __name__ == "__main__":
    main()