#!/usr/bin/env python3
"""Pull Ollama models from a JSON list. Stdlib only. Runs on the VM.

Uses POST /api/pull with stream JSON for completed/total + speed lines.
"""
from __future__ import annotations

import json
import subprocess
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

JOBS = Path("/tmp/gpu-rent-ollama-models.json")
MARKER = Path("/mnt/swarm_data/.gpu-rent-ollama-pulling")
OLLAMA = "http://127.0.0.1:11434"


def fmt_bytes(n: float) -> str:
    n = float(n)
    for unit, div in (("GB", 1024**3), ("MB", 1024**2), ("KB", 1024), ("B", 1)):
        if n >= div or unit == "B":
            if unit == "B":
                return f"{int(n)}B"
            return f"{n / div:.1f}{unit}"
    return f"{n:.0f}B"


def progress_line(label: str, done: int, total: int | None, speed: float, *, width: int = 22) -> str:
    if total and total > 0:
        pct = min(100.0, 100.0 * done / total)
        filled = min(width, max(0, int(width * done / total)))
        bar = "#" * filled + "-" * (width - filled)
        return (
            f"{label} [{bar}] {pct:5.1f}%  "
            f"{fmt_bytes(done)}/{fmt_bytes(total)}  {fmt_bytes(speed)}/s"
        )
    return f"{label} {fmt_bytes(done)}  {fmt_bytes(speed)}/s"


def listed() -> set[str]:
    """Exact tags from `ollama list` (NAME column), e.g. qwen2.5:7b."""
    try:
        out = subprocess.check_output(["ollama", "list"], text=True, stderr=subprocess.DEVNULL)
    except (subprocess.CalledProcessError, FileNotFoundError):
        return set()
    names: set[str] = set()
    for i, line in enumerate(out.splitlines()):
        if i == 0 and line.lower().startswith("name"):
            continue
        parts = line.split()
        if parts:
            names.add(parts[0])
    return names


def already_have(have: set[str], wanted: str) -> bool:
    """Exact tag match only — qwen2.5:3b must not satisfy qwen2.5:7b."""
    if wanted in have:
        return True
    if ":" not in wanted and f"{wanted}:latest" in have:
        return True
    if wanted.endswith(":latest") and wanted.rsplit(":", 1)[0] in have:
        return True
    return False


def pull_stream(name: str, label: str) -> None:
    body = json.dumps({"name": name, "stream": True}).encode("utf-8")
    req = urllib.request.Request(
        f"{OLLAMA}/api/pull",
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    t0 = time.monotonic()
    last_print = 0.0
    with urllib.request.urlopen(req, timeout=7200) as resp:
        while True:
            raw = resp.readline()
            if not raw:
                break
            try:
                ev = json.loads(raw.decode("utf-8"))
            except json.JSONDecodeError:
                continue
            if ev.get("error"):
                raise RuntimeError(str(ev["error"]))
            status = str(ev.get("status") or "")
            total = int(ev.get("total") or 0) or None
            done = int(ev.get("completed") or 0)
            now = time.monotonic()
            if total and done:
                if now - last_print >= 1.0 or done >= total:
                    elapsed = max(now - t0, 0.001)
                    speed = done / elapsed
                    line = progress_line(label, done, total, speed)
                    end = "\n" if done >= total else "\r"
                    print(line, end=end, flush=True)
                    last_print = now
            elif status and status not in {"success"} and now - last_print >= 2.0:
                print(f"{label} {status}", flush=True)
                last_print = now
            if status == "success":
                break
    print(f"{label} ok", flush=True)


def main() -> int:
    if not JOBS.is_file():
        print("no jobs file")
        return 1
    models = json.loads(JOBS.read_text(encoding="utf-8"))
    if not isinstance(models, list) or not models:
        print("ollama pull: пустой список — skip")
        return 0
    have = listed()
    MARKER.parent.mkdir(parents=True, exist_ok=True)
    MARKER.write_text(f"{int(time.time())}\n", encoding="utf-8")
    failed = 0
    try:
        for i, name in enumerate(models, 1):
            name = str(name).strip()
            if not name:
                continue
            prefix = f"[{i}/{len(models)}]"
            if already_have(have, name):
                print(f"{prefix} уже есть {name}")
                continue
            print(f"{prefix} ollama pull {name}", flush=True)
            try:
                pull_stream(name, label=f"{prefix} {name}")
                have.add(name)
            except (urllib.error.URLError, urllib.error.HTTPError, OSError, TimeoutError, RuntimeError) as exc:
                failed += 1
                print(f"FAIL pull {name}: {exc}")
    finally:
        MARKER.unlink(missing_ok=True)
    if failed:
        return 1
    print("ollama pull ok")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())