---

#!/usr/bin/env python3
"""
deploy_zcode.py — разворачивает GGUF на Lium, автоматически выбирая
САМЫЙ ДЕШЁВЫЙ подходящий узел (RTX 5090), и регистрирует провайдера на zcode.

Только stdlib. Настройка:
  export LIUM_API_KEY=sk_...        # ключ показывается один раз на https://lium.io
  export SSH_PUBKEY="$(cat ~/.ssh/id_ed25519.pub)"   # или оставить — возьмёт из ~/.ssh

Использование:
  python3 deploy_zcode.py                          # dry-run: показать выбранный узел и цену
  python3 deploy_zcode.py --rent                   # реально арендовать и развернуть
  python3 deploy_zcode.py --rent --hours 2
  python3 deploy_zcode.py --rent --quant Q5_K_M    # другой quant
  python3 deploy_zcode.py --rent --gpu "RTX 4090"  # другая модель GPU (по умолчанию 5090)
"""
import argparse, json, os, ssl, subprocess, sys, time, urllib.request, urllib.error

API_BASE = "https://lium.io/api"
TEMPLATE_ID = "587dd7ff-d68e-43bd-a54e-b7c9bb261f0c"   # light sshd, проверен
HF_REPO = "DavidAU/Qwen3.8-27B-TURBO-Fable-Cold-Fusion-735-882-Heretic-Uncensored-NEO-CODER-MAX-MTP-GGUF"
DEFAULT_QUANT = "Q4_K_M"
DEFAULT_GPU = "RTX 5090"


# ---------- helpers ----------
def api_key():
    k = os.environ.get("LIUM_API_KEY")
    if k:
        return k.strip()
    here = os.path.dirname(os.path.abspath(__file__))
    for name in ("deploy.env", ".env"):
        p = os.path.join(here, name)
        if os.path.exists(p):
            for line in open(p):
                if line.strip().startswith("LIUM_API_KEY="):
                    return line.split("=", 1)[1].strip().strip('"').strip("'")
    return None


def api(method, path, payload=None):
    url = API_BASE + "/" + path.lstrip("/")
    data = json.dumps(payload).encode() if payload is not None else None
    req = urllib.request.Request(
        url, data=data, method=method,
        headers={"Content-Type": "application/json", "X-API-Key": api_key()})
    try:
        with urllib.request.urlopen(req, timeout=30) as r:
            return r.status, json.loads(r.read().decode())
    except urllib.error.URLError as e:
        # fallback для хостов со stale CA bundle
        if isinstance(getattr(e, "reason", None), ssl.SSLError):
            ctx = ssl.create_default_context()
            ctx.check_hostname = False
            ctx.verify_mode = ssl.CERT_NONE
            with urllib.request.urlopen(req, timeout=30, context=ctx) as r:
                return r.status, json.loads(r.read().decode())
        return 0, str(e)
    except urllib.error.HTTPError as e:
        return e.code, e.read().decode()[:300]


def ssh(host, port, cmd, timeout=60):
    try:
        r = subprocess.run(
            ["ssh", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=15",
             "-p", str(port), "root@%s" % host, cmd],
            capture_output=True, text=True, timeout=timeout)
        return (r.stdout + r.stderr)[-3000:]
    except Exception as e:
        return "ERR %s" % e


# ---------- 1. авто-выбор самого дешёвого узла ----------
def pick_cheapest_executor(gpu_name):
    """
    Возвращает (executor_id, info) — самый дешёвый узел, удовлетворяющий:
      * доступен хотя бы 1 GPU нужного типа,
      * узел single-GPU (gpu_count == 1) ИЛИ полностью свободен
        (available_gpu_count == gpu_count) — иначе Lium отдаёт 400 "GPU splitting".
    Сортировка — по price_per_gpu (возрастание).
    """
    st, execs = api("GET", "executors?size=1000")
    if st != 200 or not isinstance(execs, list):
        print("Не удалось получить список executors:", st, execs); sys.exit(1)

    target = gpu_name.lower()
    candidates = []
    for e in execs:
        gname = (e.get("gpu_name") or e.get("gpu_type") or "").lower()
        if target not in gname:
            continue
        avail = int(e.get("available_gpu_count") or 0)
        total = int(e.get("gpu_count") or 0)
        if avail < 1:
            continue
        # правило против 400 "GPU splitting"
        if not (total == 1 or avail == total):
            continue
        price = float(e.get("price_per_gpu") or 1e9)
        # бонус за быстрые гео (US/CA) — но выбираем именно по цене
        country = (e.get("location") or {}).get("country_code", "??")
        speed = e.get("effective_download_speed_mbps") or 0
        candidates.append({
            "id": e.get("id"),
            "machine": e.get("machine_name"),
            "country": country,
            "price": price,
            "speed": speed,
            "total": total,
            "avail": avail,
            "gpu": e.get("gpu_name") or e.get("gpu_type"),
        })

    if not candidates:
        print("Нет доступных узлов с GPU '%s'." % gpu_name); sys.exit(1)

    candidates.sort(key=lambda c: (c["price"], -c["speed"]))
    return candidates[0], candidates[:8]


# ---------- 2. рент ----------
def rent(exec_id, pubkey, hours):
    st, body = api("POST", "executors/%s/rent" % exec_id, {
        "pod_name": "zcode-deploy",
        "template_id": TEMPLATE_ID,
        "gpu_count": 1,
        "user_public_key": pubkey,        # ВАЖНО: строка, не массив
        "termination_hours": hours,       # авто-удаление; продлить нельзя
        "initial_port_count": 2,          # 22 (ssh) + 11434 (ollama)
        "enable_jupyter": False,
        "enable_volume_encryption": False,
    })
    print("rent ->", st, str(body)[:200])
    if st in (200, 201) and isinstance(body, dict):
        return body.get("pod_id")
    return None


def wait_running(pod_id):
    for i in range(36):
        st, body = api("GET", "pods/%s" % pod_id)
        if st != 200:
            print("pod resp", st, str(body)[:120]); return None, None, None
        status = body.get("status")
        print("[%3ds] %s" % (i * 10, status))
        if status == "RUNNING":
            raw = body.get("ssh_connect_cmd", "")
            # "ssh root@1.2.3.4 -p 50009"
            parts = raw.replace("ssh root@", "").split(" -p ")
            host, port = parts[0].strip(), int(parts[1].strip())
            return host, port, body
        if status in ("CREATION_FAILED", "ERROR", "REMOVED"):
            print("FAILED:", status); return None, None, None
        time.sleep(10)
    return None, None, None


# ---------- 3. zcode settings.yaml ----------
def register_zcode(model_id, context, provider_key="lium-zcode-qwen"):
    zcode_dir = os.path.expanduser("~/.zcode")
    os.makedirs(zcode_dir, exist_ok=True)
    path = os.path.join(zcode_dir, "settings.yaml")
    block = f"""
  {provider_key}:
    displayName: Qwen3.8-27B TURBO Heretic (Lium cheapest 5090)
    api: openai-completions
    baseURL: http://127.0.0.1:11435/v1
    streamIdleTimeoutMs: 900000
    timeoutMs: 1800000
    reasoning: max
    headers:
      Authorization: Bearer local
    compat:
      thinkingFormat: qwen-chat-template
      supportsReasoningEffort: false
    models:
      - id: '{model_id}'
        name: Qwen3.8-27B TURBO NEO-CODER (RTX 5090)
        contextWindow: {context}
        maxTokens: 32768
        reasoningEfforts:
          'off': null
          low: low
          medium: medium
          high: high
          xhigh: xhigh
          max: max
"""
    existing = open(path).read() if os.path.exists(path) else ""
    if provider_key in existing:
        print("zcode: провайдер '%s' уже есть — пропускаю запись." % provider_key)
        return path
    with open(path, "a") as f:
        if "providers:" not in existing:
            f.write("llm-pi-ai:\n  providers:\n")
        f.write(block)
    return path


# ---------- main ----------
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--rent", action="store_true", help="реально арендовать (тратит деньги)")
    ap.add_argument("--hours", type=int, default=1)
    ap.add_argument("--context", type=int, default=200000)
    ap.add_argument("--quant", default=DEFAULT_QUANT, help="Q4_K_M / Q5_K_M / Q6_K ...")
    ap.add_argument("--gpu", default=DEFAULT_GPU, help="имя GPU, напр. 'RTX 5090'")
    ap.add_argument("--model", default=None, help="полный HF-репо (по умолчанию DavidAU/...)")
    o = ap.parse_args()

    if not api_key():
        print("Нет LIUM_API_KEY. Экспортируйте переменную или создайте deploy.env рядом со скриптом.")
        sys.exit(1)

    repo = o.model or HF_REPO
    model_id = f"hf.co/{repo}:{o.quant}"

    # --- выбор узла ---
    best, top = pick_cheapest_executor(o.gpu)
    print("\n=== Топ-8 самых дешёвых узлов '%s' (с учётом правил splitting) ===" % o.gpu)
    for c in top:
        print("  $%5.3f/h  %-30s  %-3s  %4.0f MB/s  avail=%d/%d  id=%s"
              % (c["price"], c["machine"][:30], c["country"], c["speed"],
                 c["avail"], c["total"], c["id"][:8]))
    print("\nВЫБРАН: $%.3f/h  %s (%s)  id=%s"
          % (best["price"], best["machine"], best["country"], best["id"]))
    print("Модель: %s" % model_id)
    print("Контекст: %d | Время аренды: %d ч | Оценка стоимости: ~$%.2f"
          % (o.context, o.hours, best["price"] * o.hours))

    if not o.rent:
        print("\n[dry-run] Ничего не арендовано. Запусти с --rent, чтобы выполнить.")
        return

    pub = os.environ.get("SSH_PUBKEY")
    if not pub:
        for k in ("id_ed25519.pub", "id_rsa.pub"):
            p = os.path.expanduser("~/.ssh/" + k)
            if os.path.exists(p):
                pub = open(p).read().strip(); break
    if not pub:
        print("Нет SSH-ключа (~/.ssh/id_ed25519.pub) и не задан SSH_PUBKEY."); sys.exit(1)

    # --- рент ---
    pod_id = rent(best["id"], pub, o.hours)
    if not pod_id:
        print("Рент не удался — попробуй следующий узел из топа вручную."); sys.exit(1)
    print("POD %s" % pod_id)

    host, port, _ = wait_running(pod_id)
    if not host:
        print("Под не достиг RUNNING."); sys.exit(1)
    print("SSH root@%s -p %d" % (host, port))

    # --- ollama ---
    print("\n--- install ollama ---")
    print(ssh(host, port,
        "apt-get update -qq >/dev/null 2>&1 || true; "
        "DEBIAN_FRONTEND=noninteractive apt-get install -y -qq zstd curl >/dev/null 2>&1 || true; "
        "curl -fsSL https://ollama.com/install.sh -o /tmp/ol.sh; sh /tmp/ol.sh >/tmp/ol_inst.log 2>&1; "
        "ollama --version", timeout=300)[-400:])

    # --- pull (не прерывать!) ---
    print("\n--- ollama pull (sha256-gated, НЕ ПРЕРЫВАТЬ) ---")
    print(ssh(host, port, "nohup ollama pull %s >/tmp/pull.log 2>&1 & echo BG" % model_id,
              timeout=30)[-120:])
    for _ in range(150):
        time.sleep(20)
        out = ssh(host, port,
            "grep -q success /tmp/pull.log && echo DONE || echo RUN; "
            "tail -1 /tmp/pull.log", timeout=30)
        if "DONE" in out:
            print("PULL SUCCESS:", out.strip().splitlines()[0]); break
        if out.startswith("ERR"):
            print("ssh err:", out[-200:]); break
    else:
        print("Pull не завершился за отведённое время — проверь /tmp/pull.log вручную.")

    # --- serve ---
    print("\n--- serve 0.0.0.0:%d ---" % 11434)
    print(ssh(host, port,
        "pkill -x ollama 2>/dev/null; sleep 1; "
        "setsid env OLLAMA_HOST=0.0.0.0 OLLAMA_CONTEXT_LENGTH=%d "
        "ollama serve >>/tmp/ol.log 2>&1 </dev/null & "
        "sleep 4; curl -s http://localhost:11434/api/version" % o.context, timeout=40)[-300:])

    # --- zcode ---
    path = register_zcode(model_id, o.context)
    print("\nzcode settings: %s" % path)

    print("\n=== ТУННЕЛЬ (держать запущенным) ===")
    print("  ssh -N -L 11435:localhost:11434 -p %d root@%s" % (port, host))
    print("\nГотово. Открой zcode GUI — модель появится в списке. "
          "Туннель должен быть активен всё время работы.")


if __name__ == "__main__":
    main()