#!/usr/bin/env python3
"""Mirror VirtFusion OS template images via aria2c.

Catalog: https://security.virtfusion.net/operating_system_templates_all.json

Includes:
  - os-cdn.virtfusion.net / vf-os.b-cdn.net (needs CDN User-Agent)
  - public upstreams (Ubuntu, Alma, Rocky, os.virtfusion.net, etc.)
"""

from __future__ import annotations

import argparse
import json
import shutil
import subprocess
import sys
import urllib.request
from pathlib import Path
from urllib.parse import urlparse

CATALOG_URL = "https://security.virtfusion.net/operating_system_templates_all.json"
CATALOG_UA = "Mozilla/5.0 (compatible; vf-c6182f2d0c1011de)"
CDN_UA = "VF/RT8B4HE276103QAP3626-OG1K29N2HOO9IT0XG2VLC4U3A52C6Y"
CDN_HOSTS = ("os-cdn.virtfusion.net", "vf-os.b-cdn.net")


def fetch_catalog(timeout: int = 60) -> list[dict]:
    req = urllib.request.Request(CATALOG_URL, headers={"User-Agent": CATALOG_UA})
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        data = json.loads(resp.read().decode())
    if not isinstance(data, list):
        raise SystemExit(f"Unexpected catalog shape: {type(data).__name__}")
    return data


def strip_virtfusion_name(path: str) -> str:
    """debian/foo-virtfusion.qcow2 -> debian/foo.qcow2"""
    p = Path(path)
    name = p.name.replace("-virtfusion", "")
    return str(p.with_name(name)) if name != p.name else path


def is_cdn_url(url: str) -> bool:
    return any(host in url for host in CDN_HOSTS)


def local_out_path(url: str, distro: str | None) -> str:
    """Build a stable relative output path for a template URL."""
    parsed = urlparse(url)
    host = parsed.netloc.lower()
    rel = parsed.path.lstrip("/")

    if is_cdn_url(url) or host.endswith("virtfusion.net") or host.endswith("b-cdn.net"):
        # Keep CDN / VF layout: alma/..., images/..., cloudlinux/...
        return strip_virtfusion_name(rel)

    # Public upstreams: group under distro when available
    name = strip_virtfusion_name(Path(rel).name)
    if not name:
        return ""
    if distro:
        return f"{distro}/{name}"
    return name


def catalog_entries(
    catalog: list[dict],
    *,
    include_eol: bool = False,
    cdn_only: bool = False,
    public_only: bool = False,
) -> list[tuple[str, str]]:
    """Return (url, relative_out_path) for template images."""
    out: list[tuple[str, str]] = []
    seen: set[str] = set()
    for item in catalog:
        if not isinstance(item, dict):
            continue
        if item.get("eol") and not include_eol:
            continue
        url = (item.get("file") or {}).get("url") or ""
        if not url.startswith(("http://", "https://")):
            continue
        cdn = is_cdn_url(url)
        if cdn_only and not cdn:
            continue
        if public_only and cdn:
            continue
        path = local_out_path(url, item.get("distro") or None)
        if not path or path in seen:
            continue
        seen.add(path)
        out.append((url, path))
    return out


def rename_existing(dest: Path) -> int:
    """Rename already-downloaded *-virtfusion* files so aria2 can resume."""
    count = 0
    for src in dest.rglob("*"):
        if not src.is_file() or "-virtfusion" not in src.name:
            continue
        dst = src.with_name(src.name.replace("-virtfusion", ""))
        if dst.exists():
            print(f"skip rename (exists): {src} -> {dst}", file=sys.stderr)
            continue
        src.rename(dst)
        ctrl = Path(str(src) + ".aria2")
        if ctrl.exists():
            ctrl.rename(Path(str(dst) + ".aria2"))
        print(f"renamed: {src.relative_to(dest)} -> {dst.relative_to(dest)}", file=sys.stderr)
        count += 1
    return count


def write_aria2_input(entries: list[tuple[str, str]], path: Path) -> None:
    lines: list[str] = []
    for url, rel in entries:
        lines.append(url)
        lines.append(f"  out={rel}")
        # CDN needs the VF UA; public mirrors accept it too.
        if is_cdn_url(url):
            lines.append(f"  user-agent={CDN_UA}")
    path.write_text("\n".join(lines) + ("\n" if lines else ""), encoding="utf-8")


def run_aria2(
    input_file: Path,
    dest: Path,
    *,
    jobs: int,
    connections: int,
    dry_run: bool,
) -> int:
    aria2c = shutil.which("aria2c")
    if not aria2c:
        raise SystemExit("aria2c not found; install aria2 (e.g. apt install aria2)")

    dest.mkdir(parents=True, exist_ok=True)
    cmd = [
        aria2c,
        f"--input-file={input_file}",
        f"--dir={dest}",
        # Default UA for public images; per-entry overrides set CDN UA above.
        f"--user-agent={CATALOG_UA}",
        "--continue=true",
        "--auto-file-renaming=false",
        "--allow-overwrite=true",
        "--conditional-get=true",
        "--remote-time=true",
        "--max-tries=5",
        "--retry-wait=3",
        f"--max-concurrent-downloads={jobs}",
        f"--split={connections}",
        f"--max-connection-per-server={connections}",
        "--min-split-size=20M",
        "--file-allocation=none",
        "--check-certificate=true",
    ]
    print(" ".join(cmd), file=sys.stderr)
    if dry_run:
        return 0
    return subprocess.call(cmd)


def main() -> int:
    p = argparse.ArgumentParser(
        description="Sync VirtFusion OS template images (CDN + public) with aria2c"
    )
    p.add_argument(
        "-d",
        "--dest",
        type=Path,
        default=Path("/home/vf-data/os/mirror"),
        help="Mirror root directory (default: /home/vf-data/os/mirror)",
    )
    p.add_argument(
        "-i",
        "--input",
        type=Path,
        default=None,
        help="aria2 input file path (default: <dest>/vf-os-templates.aria2)",
    )
    p.add_argument("--include-eol", action="store_true", help="Include EOL templates")
    p.add_argument(
        "--cdn-only",
        action="store_true",
        help="Only download os-cdn / vf-os.b-cdn.net images",
    )
    p.add_argument(
        "--public-only",
        action="store_true",
        help="Only download non-CDN (public upstream) images",
    )
    p.add_argument("-j", "--jobs", type=int, default=3, help="Concurrent downloads")
    p.add_argument("-x", "--connections", type=int, default=8, help="Connections per file")
    p.add_argument(
        "--list-only",
        action="store_true",
        help="Write aria2 input file only; do not download",
    )
    p.add_argument("--dry-run", action="store_true", help="Print aria2c command and exit")
    p.add_argument(
        "--no-rename",
        action="store_true",
        help="Do not rename existing *-virtfusion* files under dest",
    )
    args = p.parse_args()

    if args.cdn_only and args.public_only:
        raise SystemExit("Use only one of --cdn-only / --public-only")

    dest: Path = args.dest
    input_file: Path = args.input or (dest / "vf-os-templates.aria2")

    catalog = fetch_catalog()
    entries = catalog_entries(
        catalog,
        include_eol=args.include_eol,
        cdn_only=args.cdn_only,
        public_only=args.public_only,
    )
    if not entries:
        raise SystemExit("No template URLs found in catalog")

    cdn_n = sum(1 for u, _ in entries if is_cdn_url(u))
    pub_n = len(entries) - cdn_n

    dest.mkdir(parents=True, exist_ok=True)
    if not args.no_rename:
        n = rename_existing(dest)
        if n:
            print(f"Renamed {n} existing file(s)", file=sys.stderr)
    write_aria2_input(entries, input_file)
    print(
        f"Wrote {len(entries)} entries (cdn={cdn_n}, public={pub_n}) -> {input_file}",
        file=sys.stderr,
    )

    if args.list_only:
        return 0
    return run_aria2(
        input_file,
        dest,
        jobs=args.jobs,
        connections=args.connections,
        dry_run=args.dry_run,
    )


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