"""信创模盒（ModelHub XC）适配智能体。

平台契约（见 https://modelhub.org.cn/#/docs/6）：
  - 仓库根目录必须有 Dockerfile，且 EXPOSE 一个端口用于健康检查
  - 必须提供 GET /health，返回 200 表示存活
  - 通过环境变量 STRATEGY_ID 获取自身策略 ID
  - 必须正确处理 SIGTERM，30s 内完成退出
  - 资源：requests 100m/256Mi，limits 1C/512Mi

业务逻辑：
  1. 拉取社区模型库，筛出「文本生成 + 仓库完整 + 未在该卡适配过」的模型；
  2. 按卡依次调用 /api/adapt/task/add 提交验证任务；
  3. 命中 60027（该卡已验证成功）/ 60028（该卡验证中）就换下一个组合；
  4. 直到提交满 TARGET_TASKS 条或模型池耗尽。
"""

from __future__ import annotations

import json
import os
import signal
import sys
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
from datetime import datetime
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

from cards import CARDS

# ---------------------------------------------------------------- 配置

API_BASE = os.getenv("MODELHUB_API_BASE", "https://www.modelhub.org.cn").rstrip("/")
STRATEGY_ID = os.getenv("STRATEGY_ID", "")
TOKEN_ENV_KEYS = ("XC_TOKEN", "EXTERNAL_SERVICE_TOKEN", "MODELHUB_TOKEN", "SERVICE_TOKEN",
                  "XC_SERVICE_TOKEN", "API_TOKEN", "AGENT_TOKEN")
TOKEN_FILE_CANDIDATES = ("/etc/xc/token", "/etc/modelhub/token", "/var/run/secrets/xc/token",
                         "/var/run/secrets/modelhub/token", "/token", "/app/token")


def _read_token() -> str:
    for key in TOKEN_ENV_KEYS:
        val = os.getenv(key, "").strip()
        if val:
            return val
    for path in TOKEN_FILE_CANDIDATES:
        try:
            with open(path, "r", encoding="utf-8") as fh:
                val = fh.read().strip()
            if val:
                return val
        except Exception:
            continue
    return ""


XC_TOKEN = _read_token()
# 平台在 pod 内可能用不同方式鉴权：依次尝试 Xc-Token → Authorization: Bearer → 无鉴权
AUTH_MODES = ("xctoken", "bearer", "none")
AUTH_MODE = None  # 运行期探测到的可用方式
TARGET_TASKS = int(os.getenv("TARGET_TASKS", "250"))
QUOTA_PER_MODEL = int(os.getenv("QUOTA_PER_MODEL", "2"))
INTERVAL_SEC = float(os.getenv("INTERVAL_SEC", "1.5"))
MAX_ATTEMPTS = int(os.getenv("MAX_ATTEMPTS", "4000"))
DRY_RUN = os.getenv("DRY_RUN", "").strip().lower() in ("1", "true", "yes", "on")
STATE_FILE = os.getenv("STATE_FILE", "/tmp/modelhub_agent_state.json")
POOL_FILE = os.getenv("POOL_FILE", "/tmp/modelhub_agent_pool.json")
PAGE_SIZE = int(os.getenv("PAGE_SIZE", "50"))

SEED_QUERIES = [q for q in os.getenv(
    "SEED_QUERIES",
    "Qwen2.5-0.5B,Qwen2.5-1.5B,Qwen2.5-3B,Qwen2.5-7B,Qwen3-0.6B,Qwen3-1.7B,Qwen3-4B,Qwen3-8B,"
    "Llama-3.2,Llama-3.1,InternLM2.5,MiniCPM,Phi-3,Phi-4,gemma-2,Mistral-7B,Yi-1.5,GLM-4,"
    "ChatGLM3,DeepSeek-R1-Distill,Qwen2-1.5B,Qwen2-7B",
).split(",") if q.strip()]

# 明显不能直接跑文本生成验证的仓库（统一转小写后匹配）
BAD_TOKENS = ("gensyn", "gguf", "awq", "gptq", "fp8", "fp4", "nvfp4", "adapter", "lora",
              "int4", "int8", "w4a16", "w8a8", "w8a16", "quantized", "bnb", "exl2", "hqq",
              "mlx", "onnx", "openvino", "tflite")
MIN_MACHINES = int(os.getenv("MIN_MACHINES", "2"))  # 至少在一张卡上适配成功过，说明仓库完整

# ---------------------------------------------------------------- 全局状态

STOP = threading.Event()
LOCK = threading.Lock()
STATE = {
    "started_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
    "submitted": [],          # [{model, card, taskId, at}]
    "attempts": 0,
    "codes": {},              # code -> count
    "last_event": None,
    "finished": False,
    "error": None,
}


def log(msg: str) -> None:
    print(f"[{datetime.now().strftime('%H:%M:%S')}] {msg}", flush=True)


def _save_state() -> None:
    try:
        tmp = STATE_FILE + ".tmp"
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(STATE, f, ensure_ascii=False, indent=1)
        os.replace(tmp, STATE_FILE)
    except Exception as exc:  # 状态落盘失败不影响主流程
        log(f"save state failed: {exc}")


def _load_state() -> None:
    try:
        with open(STATE_FILE, "r", encoding="utf-8") as f:
            saved = json.load(f)
        STATE["submitted"] = saved.get("submitted", [])
        log(f"resume: {len(STATE['submitted'])} submitted")
    except FileNotFoundError:
        pass
    except Exception as exc:
        log(f"load state failed: {exc}")


# ---------------------------------------------------------------- HTTP

def _auth_header(mode: str) -> dict:
    if not XC_TOKEN or mode == "none":
        return {}
    if mode == "bearer":
        return {"Authorization": "Bearer " + XC_TOKEN}
    return {"Xc-Token": XC_TOKEN}


def _request(method: str, path: str, body: dict | None = None, timeout: int = 30,
             auth_mode: str | None = None) -> tuple[int, str]:
    url = API_BASE + path
    data = None
    headers = {"Accept": "application/json", "User-Agent": "modelhub-xc-agent/1.0"}
    if body is not None:
        data = json.dumps(body, ensure_ascii=False).encode("utf-8")
        headers["Content-Type"] = "application/json"
    headers.update(_auth_header(auth_mode or AUTH_MODE or "xctoken"))
    req = urllib.request.Request(url, data=data, headers=headers, method=method)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status, resp.read().decode("utf-8", "replace")
    except urllib.error.HTTPError as exc:
        return exc.code, exc.read().decode("utf-8", "replace")
    except Exception as exc:
        return 0, str(exc)


def _get_json(path: str, timeout: int = 30, auth_mode: str | None = None) -> dict:
    status, text = _request("GET", path, timeout=timeout, auth_mode=auth_mode)
    try:
        return json.loads(text)
    except Exception:
        return {"code": -1, "message": f"HTTP {status}: {text[:200]}"}


def _post_json(path: str, body: dict, timeout: int = 60, auth_mode: str | None = None) -> dict:
    status, text = _request("POST", path, body=body, timeout=timeout, auth_mode=auth_mode)
    try:
        return json.loads(text)
    except Exception:
        return {"code": -1, "message": f"HTTP {status}: {text[:200]}"}


# ---------------------------------------------------------------- 模型池

def build_pool() -> list[dict]:
    """拉模型库，返回 [{id, machines}]，按「已适配卡数」升序（可选卡更多）。"""
    seen: dict[str, dict] = {}
    for query in SEED_QUERIES:
        for page in range(1, 6):
            payload = {"pageNum": page, "pageSize": PAGE_SIZE, "searchText": query}
            res = _post_json("/api/computility/models/list/page/vo", payload, timeout=40)
            records = (res.get("data") or {}).get("records") or []
            if not records:
                break
            for item in records:
                model_id = item.get("modelId") or ""
                machines = [m.get("machineName") for m in (item.get("machines") or [])]
                level = (item.get("taskLevelsInfo") or {}).get("taskLevelChineseName")
                if not model_id or not machines or level != "文本生成":
                    continue
                lower = model_id.lower()
                if any(tok in lower for tok in BAD_TOKENS):
                    continue
                if len(machines) < MIN_MACHINES:
                    continue
                if model_id in seen:
                    continue
                seen[model_id] = {"id": model_id, "machines": machines}
            if len(records) < PAGE_SIZE:
                break
        if STOP.is_set():
            break
    pool = sorted(seen.values(), key=lambda x: len(x["machines"]))
    log(f"pool built: {len(pool)} models")
    try:
        with open(POOL_FILE, "w", encoding="utf-8") as f:
            json.dump(pool, f)
    except Exception:
        pass
    return pool


def load_pool() -> list[dict]:
    try:
        with open(POOL_FILE, "r", encoding="utf-8") as f:
            pool = json.load(f)
        if pool:
            log(f"pool loaded from cache: {len(pool)}")
            return pool
    except Exception:
        pass
    return build_pool()


# ---------------------------------------------------------------- 提交

def submit(model: str, card: dict, auth_mode: str | None = None) -> dict:
    body = {
        "modelAddress": f"https://huggingface.co/{model}",
        "taskType": "text-generation",
        "targetGpu": card["gpu"],
        "framework": card["framework"],
        "strategyId": STRATEGY_ID,
        "configParams": card["configParams"],
    }
    if DRY_RUN:
        return {"code": 0, "data": {"id": "dry-run"}, "dryRun": True}
    return _post_json("/api/adapt/task/add", body, auth_mode=auth_mode)


def worker() -> None:
    _load_state()
    if not STRATEGY_ID:
        STATE["error"] = "STRATEGY_ID 未注入，智能体无法提交（平台运行时自动注入）"
        log(STATE["error"])
        return
    global AUTH_MODE

    pool = load_pool()
    submitted_pairs = {(s["model"], s["card"]) for s in STATE["submitted"]}
    per_model: dict[str, int] = {}
    for s in STATE["submitted"]:
        per_model[s["model"]] = per_model.get(s["model"], 0) + 1

    while not STOP.is_set() and len(STATE["submitted"]) < TARGET_TASKS:
        progress = False
        for card in CARDS:
            if STOP.is_set() or len(STATE["submitted"]) >= TARGET_TASKS:
                break
            for cand in pool:
                if STOP.is_set() or len(STATE["submitted"]) >= TARGET_TASKS:
                    break
                if STATE["attempts"] >= MAX_ATTEMPTS:
                    log("reach MAX_ATTEMPTS, stop")
                    return
                if card["gpu"] in cand["machines"]:
                    continue
                if per_model.get(cand["id"], 0) >= QUOTA_PER_MODEL:
                    continue
                pair = (cand["id"], card["gpu"])
                if pair in submitted_pairs:
                    continue

                res = submit(cand["id"], card)
                code = res.get("code")
                # 401 时换一种鉴权方式重试（pod 内可能用 Bearer 或免鉴权）
                if code == 401:
                    for mode in AUTH_MODES:
                        if mode == (AUTH_MODE or "xctoken"):
                            continue
                        res = submit(cand["id"], card, auth_mode=mode)
                        code = res.get("code")
                        if code != 401:
                            AUTH_MODE = mode
                            log(f"auth mode ok: {mode}")
                            break
                    if code == 401:
                        STATE["error"] = "401: 三种鉴权方式均被拒绝"
                        log(STATE["error"])
                        _save_state()
                        STOP.wait(300)
                STATE["attempts"] += 1
                STATE["codes"][str(code)] = STATE["codes"].get(str(code), 0) + 1
                STATE["last_event"] = {
                    "at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                    "model": cand["id"], "gpu": card["gpu"], "code": code,
                    "message": (res.get("message") or "")[:160],
                }
                progress = True

                if code == 0:
                    task_id = (res.get("data") or {}).get("id")
                    with LOCK:
                        STATE["submitted"].append({
                            "model": cand["id"], "card": card["gpu"],
                            "taskId": task_id,
                            "at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                        })
                        per_model[cand["id"]] = per_model.get(cand["id"], 0) + 1
                    submitted_pairs.add(pair)
                    log(f"OK {len(STATE['submitted'])}/{TARGET_TASKS} | {card['label']} | {cand['id']} | {task_id}")
                    _save_state()
                elif code == 60027:      # 该卡已经验证成功过
                    submitted_pairs.add(pair)
                elif code == 60028:      # 该组合正在验证中
                    submitted_pairs.add(pair)
                elif code == 50000:      # 拉不到上游模型信息
                    submitted_pairs.add(pair)
                elif code in (401, 403, 60014):
                    log(f"fatal {code}: {res.get('message')}")
                    STATE["error"] = f"{code}: {res.get('message')}"
                    _save_state()
                    STOP.wait(120)
                elif code in (-1, 0) and res.get("message", "").startswith("HTTP"):
                    log(f"network issue: {res.get('message')}")
                    STOP.wait(10)
                time.sleep(INTERVAL_SEC)
        if not progress:
            log("no progress this round; rebuild pool")
            pool = build_pool()
            if not pool:
                log("empty pool, exit")
                break
            time.sleep(30)

    STATE["finished"] = True
    _save_state()
    log(f"worker done: submitted={len(STATE['submitted'])} attempts={STATE['attempts']}")


# ---------------------------------------------------------------- HTTP 服务

class Handler(BaseHTTPRequestHandler):
    server_version = "modelhub-xc-agent/1.0"

    def do_GET(self) -> None:  # noqa: N802
        path = urllib.parse.urlparse(self.path).path
        if path == "/health":
            self._send({"status": "ok"})
            return
        if path == "/stats":
            with LOCK:
                payload = json.loads(json.dumps(STATE, ensure_ascii=False))
            payload["target"] = TARGET_TASKS
            payload["progress"] = f"{len(payload['submitted'])}/{TARGET_TASKS}"
            payload["strategyId"] = STRATEGY_ID or None
            payload["dryRun"] = DRY_RUN
            self._send(payload)
            return
        if path == "/":
            self._send({
                "name": "modelhub-xc-adapt-agent",
                "status": "stopped" if STOP.is_set() else "running",
                "strategyId": STRATEGY_ID or None,
                "target": TARGET_TASKS,
                "submitted": len(STATE["submitted"]),
                "attempts": STATE["attempts"],
                "dryRun": DRY_RUN,
            })
            return
        self._send({"error": "not found"}, status=404)

    def log_message(self, fmt: str, *args) -> None:
        print(f"{self.address_string()} - {fmt % args}", flush=True)

    def _send(self, body: dict, status: int = 200) -> None:
        payload = json.dumps(body, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)


def _handle_signal(signum: int, _frame) -> None:
    log(f"received signal {signum}, shutting down")
    STOP.set()


def _diagnose() -> None:
    """打印运行环境的令牌线索（脱敏），便于排查 401。"""
    log("---- env diagnose ----")
    for key in sorted(os.environ):
        if any(k in key.upper() for k in ("TOKEN", "STRATEGY", "XC", "SECRET", "KEY", "PASS", "API")):
            val = os.environ[key]
            log(f"env {key} = {val[:6]}...(len={len(val)})")
    log(f"resolved token: {'yes len=' + str(len(XC_TOKEN)) if XC_TOKEN else 'NO'}")
    for path in TOKEN_FILE_CANDIDATES:
        log(f"file {path} exists={os.path.exists(path)}")
    for root in ("/var/run/secrets", "/etc/xc", "/etc/modelhub"):
        try:
            log(f"dir {root}: {os.listdir(root)[:10]}")
        except Exception as exc:
            log(f"dir {root}: {type(exc).__name__}")
    log("---- env diagnose end ----")


def main() -> None:
    signal.signal(signal.SIGTERM, _handle_signal)
    _diagnose()
    signal.signal(signal.SIGINT, _handle_signal)

    thread = threading.Thread(target=worker, name="submit-worker", daemon=True)
    thread.start()

    server = ThreadingHTTPServer(("0.0.0.0", int(os.getenv("PORT", "8080"))), Handler)
    server.timeout = 1
    log(f"agent listening on 0.0.0.0:{os.getenv('PORT', '8080')} (dryRun={DRY_RUN}, target={TARGET_TASKS})")

    while not STOP.is_set():
        server.handle_request()

    thread.join(timeout=25)
    server.server_close()
    _save_state()
    log("bye")
    sys.exit(0)


if __name__ == "__main__":
    main()