"""ModelHub XC 适配智能体 —— 低排队算力批次（5 条）

每条任务分配一张低排队算力卡，并选用该卡历史通过率较高的框架。
OSP 系列 max_position_embeddings=2048，配置里的 max_model_len 已由 4096 下调为 2048。
"""
import json
import os
import signal
import threading
import time
from datetime import datetime
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import requests

BASE_URL = os.environ.get("BASE_URL", "https://modelhub.org.cn")
ADD_TASK_ENDPOINT = "/api/adapt/task/add"
XC_TOKEN = os.environ.get("XC_TOKEN", "6616dfe75c874e61987b1f7c908e6d77")
STRATEGY_ID = os.environ.get("STRATEGY_ID", "")
TASK_TYPE = "text-generation"

TASKS = [
    {
        "label": "OSP-1.4B-100B-Muon-Only",
        "model": "https://www.modelscope.cn/models/dmis-lab/OSP-1.4B-100B-Muon-Only",
        "gpu": "MetaX_c-500",
        "framework": "vllm_tokenizer_patch",
        "config": "nv_framework: vllm_fix_tokenizer\nframework: vllm_tokenizer_patch\nstorage: gpfs\nmax_tokens: 1024\nmax_model_len: 2048\napi: completion\nsut_config:\n    gpu_num: 1\n    values:\n        command: ['/opt/conda/bin/vllm', 'serve', '/model', '--port', '8000', '--served-model-name', 'llm',  '--max-model-len', '2048', '--gpu-memory-utilization', '0.9',  '--enforce-eager', '--trust-remote-code' ,'-tp', '1' ]\nref_config:\n    gpu_num: 1\n    values:\n        command: ['/opt/entrypoint.sh', '/model', '--port', '80', '--served-model-name', 'llm', '--max-model-len', '2048', '--enforce-eager', '--trust-remote-code', '-tp', '1']"
    },
    {
        "label": "OSP-1.4B-100B-Muon-SSNorm",
        "model": "https://www.modelscope.cn/models/dmis-lab/OSP-1.4B-100B-Muon-SSNorm",
        "gpu": "Kunlunxin_p-800",
        "framework": "vllm_fix_tokenizer",
        "config": "framework: vllm_fix_tokenizer\nlang: en\nmax_model_len: 2048\nsut_config:\n  gpu_num: 1\n  values:\n    command:\n    - /opt/entrypoint.sh    - /model\n    - --port\n    - '80'\n    - --served-model-name\n    - llm\n    - --max-model-len\n    - '2048'\n    - --enforce-eager\n    - --trust-remote-code\n    - -tp\n    - '1'\nref_config:\n  gpu_num: 1\n  values:\n    command:\n    - /opt/entrypoint.sh    - /model\n    - --port\n    - '80'\n    - --served-model-name\n    - llm\n    - --max-model-len\n    - '2048'\n    - --enforce-eager\n    - --trust-remote-code\n    - -tp\n    - '1'"
    },
    {
        "label": "OSP-1.4B-100B-Shampoo-SSNorm",
        "model": "https://www.modelscope.cn/models/dmis-lab/OSP-1.4B-100B-Shampoo-SSNorm",
        "gpu": "Cambricon_mlu-370-x8",
        "framework": "vllm-mlu",
        "config": "framework: vllm-mlu\nnv_framework: vllm_fix_tokenizer\napi: completion\nlang: en\ntemperature: 0.4\nrepetition_penalty: 1.1\ntop_p: 0.9\nsut_config:\n  gpu_num: 1\n  values:\n    env:\n    - name: MAX_MODEL_LEN\n      value: 2048\n    - name: VLLM_ALLOW_LONG_MAX_MODEL_LEN\n      value: '1'\n    command:\n    - vllm\n    - serve\n    - /model\n    - --port\n    - '8000'\n    - --served-model-name\n    - llm\n    - --max-model-len\n    - '2048'\n    - --trust-remote-code\n    - --dtype\n    - float16\nref_config:\n  cpu_num: 2\n  gpu_num: 1\n  values:\n    env:\n    - name: MAX_MODEL_LEN\n      value: 2048\n    - name: VLLM_ALLOW_LONG_MAX_MODEL_LEN\n      value: '1'\n    command:\n    - /opt/entrypoint.sh\n    - /model\n    - --port\n    - '80'\n    - --served-model-name\n    - llm\n    - --max-model-len\n    - '2048'\n    - --trust-remote-code\n    - --dtype\n    - float16\n"
    },
    {
        "label": "OSP-1.4B-100B-Muon-EmbProj",
        "model": "https://www.modelscope.cn/models/dmis-lab/OSP-1.4B-100B-Muon-EmbProj",
        "gpu": "Biren_166m",
        "framework": "vllm_fix_tokenizer",
        "config": "framework: vllm_fix_tokenizer\nlang: en\nmax_model_len: 2048\nsut_config:\n  gpu_num: 1\n  values:\n    command:\n    - /opt/entrypoint.sh    - /model\n    - --port\n    - '80'\n    - --served-model-name\n    - llm\n    - --max-model-len\n    - '2048'\n    - --enforce-eager\n    - --trust-remote-code\n    - -tp\n    - '1'\nref_config:\n  gpu_num: 1\n  values:\n    command:\n    - /opt/entrypoint.sh    - /model\n    - --port\n    - '80'\n    - --served-model-name\n    - llm\n    - --max-model-len\n    - '2048'\n    - --enforce-eager\n    - --trust-remote-code\n    - -tp\n    - '1'"
    },
    {
        "label": "OSP-1.4B-1T-Muon-SSNorm-EmbProj",
        "model": "https://www.modelscope.cn/models/dmis-lab/OSP-1.4B-1T-Muon-SSNorm-EmbProj",
        "gpu": "Sunrise_pt-200-x1",
        "framework": "vllm_fix_tokenizer",
        "config": "framework: vllm_fix_tokenizer\napi: completion\nmax_tokens: 1024\ntemperature: 0.7\nrepetition_penalty: 1.2\ntop_p: 0.9\nlang: zh\nmax_model_len: 2048\nsut_config:\n  gpu_num: 1\n  values:\n    command:\n    - /opt/entrypoint.sh\n    - /model\n    - --port\n    - '8000'\n    - --served-model-name\n    - llm\n    - --max-model-len\n    - '2048'\n    - --gpu-memory-utilization\n    - '0.9'\n    - --enforce-eager\n    - --trust-remote-code\nref_config:\n  gpu_num: 1\n  values:\n    command:\n    - /opt/entrypoint.sh\n    - /model\n    - --port\n    - '8000'\n    - --served-model-name\n    - llm\n    - --max-model-len\n    - '2048'\n    - --dtype\n    - float16\n    - --gpu-memory-utilization\n    - '0.95'\n    - -tp\n    - '1'\n    - --enforce-eager\n    - --trust-remote-code\n"
    }
]

HEADERS = {"Content-Type": "application/json", "Xc-Token": XC_TOKEN}
HTTP_HOST, HTTP_PORT = "0.0.0.0", 8080
GAP = 4
_shutdown = threading.Event()
_state = {"phase": "starting", "total": len(TASKS), "submitted": 0, "failed": 0, "results": []}


def submit_one(t):
    body = {
        "configParams": t["config"],
        "framework": t["framework"],
        "modelAddress": t["model"],
        "targetGpu": t["gpu"],
        "taskType": TASK_TYPE,
        "strategyId": STRATEGY_ID,
    }
    try:
        r = requests.post(BASE_URL + ADD_TASK_ENDPOINT, headers=HEADERS, json=body, timeout=90)
        j = r.json()
    except Exception as e:  # noqa: BLE001
        j = {"code": -1, "message": "exception: %s" % e}
    tid = (j.get("data") or {}).get("id")
    print("[submit] %s gpu=%s fw=%s code=%s msg=%s taskId=%s"
          % (t["label"], t["gpu"], t["framework"], j.get("code"), j.get("message"), tid), flush=True)
    return j.get("code"), tid


def run():
    _state["phase"] = "submitting"
    for i, t in enumerate(TASKS, 1):
        if _shutdown.is_set():
            break
        code, tid = submit_one(t)
        _state["results"].append({"label": t["label"], "gpu": t["gpu"],
                                  "framework": t["framework"], "code": code, "task_id": tid})
        _state["submitted" if code == 0 else "failed"] += 1
        if i < len(TASKS):
            time.sleep(GAP)
    _state["phase"] = "done"
    print("[batch] 完成 submitted=%d failed=%d" % (_state["submitted"], _state["failed"]), flush=True)


class H(BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path == "/health":
            self._json({"status": "ok"})
        elif self.path in ("/", "/status"):
            self._json(_state)
        else:
            self._json({"error": "not found"}, 404)

    def _json(self, b, c=200):
        raw = json.dumps(b, ensure_ascii=False, default=str).encode()
        self.send_response(c)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(raw)))
        self.end_headers()
        self.wfile.write(raw)

    def log_message(self, *a):
        pass


def main():
    signal.signal(signal.SIGTERM, lambda *_: _shutdown.set())
    print("[main] strategy_id=%s tasks=%d" % (STRATEGY_ID, len(TASKS)), flush=True)
    for t in TASKS:
        print("[main] 计划: %s @ %s / %s" % (t["label"], t["gpu"], t["framework"]), flush=True)
    srv = ThreadingHTTPServer((HTTP_HOST, HTTP_PORT), H)
    srv.timeout = 1
    threading.Thread(target=run, daemon=True).start()
    while not _shutdown.is_set():
        srv.handle_request()
    srv.server_close()


main()