#!/usr/bin/env python3
"""Clone git extensions on the VM. Stdlib only. Token file optional."""
from __future__ import annotations

import json
import os
import subprocess
import sys
from pathlib import Path
from urllib.parse import urlsplit, urlunsplit

TOKEN_PATH = Path("/tmp/gpu-rent-git.token")
JOBS_PATH = Path("/tmp/gpu-rent-ext.json")
MARKER = Path("/mnt/swarm_data/.gpu-rent-extensions-seeded")


def strip_auth(url: str) -> str:
    parts = urlsplit(url)
    host = parts.hostname or ""
    if parts.port:
        host = f"{host}:{parts.port}"
    return urlunsplit((parts.scheme, host, parts.path, parts.query, parts.fragment))


def with_token(url: str, token: str) -> str:
    if not token:
        return url
    if url.startswith("https://github.com/"):
        return "https://x-access-token:" + token + "@github.com/" + url[len("https://github.com/") :]
    if url.startswith("https://gitlab.com/"):
        return "https://oauth2:" + token + "@gitlab.com/" + url[len("https://gitlab.com/") :]
    return url


def run(argv: list[str], cwd: str | None = None) -> None:
    subprocess.check_call(argv, cwd=cwd)


def is_sha(ref: str) -> bool:
    ref = ref.strip()
    return len(ref) == 40 and all(c in "0123456789abcdefABCDEF" for c in ref)


def clone_one(job: dict, token: str) -> None:
    dest = Path(job["dest"])
    url = job["url"].strip()
    ref = (job.get("ref") or "main").strip()
    dest.parent.mkdir(parents=True, exist_ok=True)
    authed = with_token(url, token)
    if dest.is_dir() and (dest / ".git").is_dir():
        origin = subprocess.check_output(["git", "-C", str(dest), "remote", "get-url", "origin"], text=True).strip()
        if strip_auth(origin) != strip_auth(url):
            print(f"FAIL origin mismatch {dest}: {origin} != {url}", file=sys.stderr)
            raise SystemExit(2)
        run(["git", "-C", str(dest), "fetch", "--recurse-submodules", "origin"])
        run(["git", "-C", str(dest), "checkout", ref])
        print(f"updated {dest}")
        return
    if dest.exists():
        print(f"FAIL {dest} exists but is not a git repo", file=sys.stderr)
        raise SystemExit(2)
    if is_sha(ref):
        run(["git", "clone", "--recurse-submodules", authed, str(dest)])
        run(["git", "-C", str(dest), "fetch", "origin", ref])
        run(["git", "-C", str(dest), "checkout", ref])
    else:
        try:
            run(["git", "clone", "--recurse-submodules", "--depth", "1", "--branch", ref, authed, str(dest)])
        except subprocess.CalledProcessError:
            run(["git", "clone", "--recurse-submodules", authed, str(dest)])
            run(["git", "-C", str(dest), "checkout", ref])
    print(f"cloned {dest}")


def main() -> int:
    token = TOKEN_PATH.read_text(encoding="utf-8").strip() if TOKEN_PATH.is_file() else ""
    jobs = json.loads(JOBS_PATH.read_text(encoding="utf-8"))
    failed = 0
    for job in jobs:
        try:
            clone_one(job, token)
        except Exception as exc:
            failed += 1
            print(f"FAIL {job.get('dest')}: {exc}", file=sys.stderr)
    if TOKEN_PATH.is_file():
        TOKEN_PATH.unlink()
    if failed:
        return 1
    MARKER.parent.mkdir(parents=True, exist_ok=True)
    MARKER.write_text("ok\n", encoding="utf-8")
    print("extensions ok")
    return 0


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