#!/usr/bin/env python3
"""Download GGUF files for llama.cpp from a JSON job list. Stdlib only."""
from __future__ import annotations

import json
import os
import sys
import urllib.error
import urllib.request
from pathlib import Path

JOBS = Path("/tmp/gpu-rent-llamacpp-models.json")
MODELS_DIR = Path("/mnt/swarm_data/llamacpp/models")
TOKEN_FILE = Path("/tmp/gpu-rent-hf.token")


def main() -> int:
    if TOKEN_FILE.is_file():
        try:
            os.environ["HF_TOKEN"] = TOKEN_FILE.read_text(encoding="utf-8").strip()
        finally:
            try:
                TOKEN_FILE.unlink(missing_ok=True)
            except OSError:
                pass
    if not JOBS.is_file():
        print("no jobs file", file=sys.stderr)
        return 1
    jobs = json.loads(JOBS.read_text(encoding="utf-8"))
    if not isinstance(jobs, list) or not jobs:
        print("llamacpp fetch: пустой список — skip")
        return 0
    MODELS_DIR.mkdir(parents=True, exist_ok=True)
    token = (os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or "").strip()
    failed = 0
    for i, job in enumerate(jobs, 1):
        if not isinstance(job, dict):
            continue
        url = str(job.get("url") or "").strip()
        name = str(job.get("filename") or "").strip()
        if not url:
            continue
        if not name:
            name = url.rstrip("/").rsplit("/", 1)[-1] or "model.gguf"
        dest = MODELS_DIR / name
        if dest.is_file() and dest.stat().st_size > 1_000_000:
            print(f"[{i}/{len(jobs)}] уже есть {name} ({dest.stat().st_size} bytes)")
            continue
        print(f"[{i}/{len(jobs)}] download {name}")
        partial = dest.with_suffix(dest.suffix + ".partial")
        headers = {"User-Agent": "gpu-rent/1"}
        if token:
            headers["Authorization"] = f"Bearer {token}"
        req = urllib.request.Request(url, headers=headers)
        try:
            with urllib.request.urlopen(req, timeout=600) as resp, partial.open("wb") as out:
                while True:
                    chunk = resp.read(1024 * 1024)
                    if not chunk:
                        break
                    out.write(chunk)
            partial.replace(dest)
            print(f"ok {name} ({dest.stat().st_size} bytes)")
        except (urllib.error.URLError, urllib.error.HTTPError, OSError, TimeoutError) as exc:
            failed += 1
            print(f"FAIL {name}: {exc}", file=sys.stderr)
            try:
                partial.unlink(missing_ok=True)
            except OSError:
                pass
    if failed:
        return 1
    print("llamacpp fetch ok")
    return 0


if __name__ == "__main__":
    sys.exit(main())