#!/usr/bin/env python3
"""Mirror PyTorch wheel files and generate a pip simple index."""

from __future__ import annotations

import argparse
import concurrent.futures
import hashlib
import html
import os
import posixpath
import re
import sys
import time
from html.parser import HTMLParser
from pathlib import Path
from urllib.error import HTTPError, URLError
from urllib.parse import quote, unquote, urljoin, urlparse, urlunparse
from urllib.request import Request, urlopen


DEFAULT_SOURCE = "https://mirrors.tencent.com/pytorch-wheels/whl/cu129/"
DEFAULT_DEST = Path("/data_hdd/wheels/mirror")
DEFAULT_WORKERS = 6
USER_AGENT = "pytorch-wheel-mirror/1.0"
CHUNK_SIZE = 1024 * 1024


class LinkParser(HTMLParser):
    def __init__(self) -> None:
        super().__init__()
        self.hrefs: list[str] = []

    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
        if tag.lower() != "a":
            return
        for name, value in attrs:
            if name.lower() == "href" and value:
                self.hrefs.append(value)


def canonical_url(url: str, *, directory: bool = False) -> str:
    parsed = urlparse(url)
    path = posixpath.normpath(parsed.path)
    if parsed.path.endswith("/") or directory:
        path = path.rstrip("/") + "/"
    if not path.startswith("/"):
        path = "/" + path
    return urlunparse((parsed.scheme, parsed.netloc, path, "", "", ""))


def request_url(url: str, *, method: str = "GET", headers: dict[str, str] | None = None, timeout: int = 30):
    request_headers = {"User-Agent": USER_AGENT}
    if headers:
        request_headers.update(headers)
    request = Request(url, headers=request_headers, method=method)
    return urlopen(request, timeout=timeout)


def fetch_text(url: str, *, timeout: int, retries: int) -> str:
    last_error: Exception | None = None
    for attempt in range(retries + 1):
        try:
            with request_url(url, timeout=timeout) as response:
                body = response.read()
                charset = response.headers.get_content_charset() or "utf-8"
                return body.decode(charset, errors="replace")
        except (HTTPError, URLError, TimeoutError, OSError) as exc:
            last_error = exc
            if attempt < retries:
                time.sleep(min(2**attempt, 8))
    raise RuntimeError(f"failed to fetch {url}: {last_error}")


def extract_links(index_html: str, base_url: str) -> list[str]:
    parser = LinkParser()
    parser.feed(index_html)
    links: list[str] = []
    for href in parser.hrefs:
        href = html.unescape(href).strip()
        if not href or href.startswith("#"):
            continue
        absolute = urljoin(base_url, href)
        parsed = urlparse(absolute)
        if parsed.scheme not in {"http", "https"}:
            continue
        links.append(urlunparse((parsed.scheme, parsed.netloc, parsed.path, "", "", "")))
    return links


def crawl_wheels(source_url: str, *, timeout: int, retries: int) -> list[str]:
    source_root = canonical_url(source_url, directory=True)
    pending = [source_root]
    visited: set[str] = set()
    wheels: set[str] = set()

    while pending:
        current = pending.pop()
        current = canonical_url(current, directory=True)
        if current in visited:
            continue
        visited.add(current)
        print(f"crawl {current}")

        index_html = fetch_text(current, timeout=timeout, retries=retries)
        for link in extract_links(index_html, current):
            parsed = urlparse(link)
            decoded_path = unquote(parsed.path)
            if decoded_path.endswith("/"):
                directory_url = canonical_url(link, directory=True)
                if directory_url.startswith(source_root) and directory_url not in visited:
                    pending.append(directory_url)
                continue
            if decoded_path.lower().endswith(".whl"):
                file_url = canonical_url(link)
                if file_url.startswith(source_root):
                    wheels.add(file_url)

    return sorted(wheels)


def source_local_root(source_url: str) -> Path:
    parts = [unquote(part) for part in urlparse(source_url).path.split("/") if part]
    if "pytorch-wheels" in parts:
        parts = parts[parts.index("pytorch-wheels") + 1 :]
    return Path(*parts) if parts else Path(".")


def url_to_local_path(url: str, source_url: str, dest_root: Path) -> Path:
    source_root = canonical_url(source_url, directory=True)
    relative_url = canonical_url(url)[len(source_root) :]
    relative_path = Path(*[unquote(part) for part in relative_url.split("/") if part])
    return dest_root / source_local_root(source_url) / relative_path


def remote_size(url: str, *, timeout: int) -> int | None:
    try:
        with request_url(url, method="HEAD", timeout=timeout) as response:
            size = response.headers.get("Content-Length")
            return int(size) if size is not None else None
    except (HTTPError, URLError, TimeoutError, OSError, ValueError):
        return None


def download_file(url: str, local_path: Path, *, timeout: int, retries: int) -> str:
    local_path.parent.mkdir(parents=True, exist_ok=True)
    part_path = local_path.with_name(local_path.name + ".part")
    expected_size = remote_size(url, timeout=timeout)

    if local_path.exists() and expected_size is not None and local_path.stat().st_size == expected_size:
        return f"skip {local_path.name}"
    if local_path.exists() and expected_size is None:
        return f"skip {local_path.name}"

    last_error: Exception | None = None
    for attempt in range(retries + 1):
        resume_from = part_path.stat().st_size if part_path.exists() else 0
        headers: dict[str, str] = {}
        mode = "ab" if resume_from else "wb"
        if resume_from:
            headers["Range"] = f"bytes={resume_from}-"

        try:
            with request_url(url, headers=headers, timeout=timeout) as response:
                if resume_from and response.status != 206:
                    resume_from = 0
                    mode = "wb"
                with part_path.open(mode + "") as output:
                    while True:
                        chunk = response.read(CHUNK_SIZE)
                        if not chunk:
                            break
                        output.write(chunk)

            downloaded_size = part_path.stat().st_size
            if expected_size is not None and downloaded_size != expected_size:
                raise RuntimeError(f"size mismatch: got {downloaded_size}, expected {expected_size}")

            part_path.replace(local_path)
            return f"download {local_path.name}"
        except (HTTPError, URLError, TimeoutError, OSError, RuntimeError) as exc:
            last_error = exc
            if attempt < retries:
                time.sleep(min(2**attempt, 8))

    raise RuntimeError(f"failed to download {url}: {last_error}")


def normalize_package_name(name: str) -> str:
    return re.sub(r"[-_.]+", "-", name).lower()


def wheel_package_name(filename: str) -> str:
    return normalize_package_name(filename.split("-", 1)[0])


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as wheel_file:
        for chunk in iter(lambda: wheel_file.read(CHUNK_SIZE), b""):
            digest.update(chunk)
    return digest.hexdigest()


def write_html(path: Path, body: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(body, encoding="utf-8")


def generate_simple_index(dest_root: Path, wheel_root: Path, *, hashes: bool) -> None:
    wheels = sorted(wheel_root.rglob("*.whl"))
    packages: dict[str, list[Path]] = {}
    for wheel in wheels:
        packages.setdefault(wheel_package_name(wheel.name), []).append(wheel)

    simple_root = dest_root / "simple"
    project_links = []
    for package in sorted(packages):
        project_links.append(f'<a href="{quote(package)}/">{html.escape(package)}</a>')

        wheel_links = []
        package_dir = simple_root / package
        for wheel in packages[package]:
            relative = os.path.relpath(wheel, package_dir).replace(os.sep, "/")
            href = quote(relative, safe="/._-+~")
            if hashes:
                href += f"#sha256={sha256_file(wheel)}"
            wheel_links.append(f'<a href="{href}">{html.escape(wheel.name)}</a>')

        write_html(package_dir / "index.html", html_page("\n".join(wheel_links)))

    write_html(simple_root / "index.html", html_page("\n".join(project_links)))
    print(f"indexed {len(wheels)} wheels in {len(packages)} packages")


def html_page(links: str) -> str:
    return f"<!doctype html>\n<html><body>\n{links}\n</body></html>\n"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Mirror PyTorch wheels and generate a pip simple index.")
    parser.add_argument("--source", default=DEFAULT_SOURCE, help=f"source URL, default: {DEFAULT_SOURCE}")
    parser.add_argument("--dest", type=Path, default=DEFAULT_DEST, help=f"mirror root, default: {DEFAULT_DEST}")
    parser.add_argument("--workers", type=int, default=DEFAULT_WORKERS, help=f"download workers, default: {DEFAULT_WORKERS}")
    parser.add_argument("--timeout", type=int, default=30, help="HTTP timeout in seconds")
    parser.add_argument("--retries", type=int, default=3, help="network retry count")
    parser.add_argument("--dry-run", action="store_true", help="crawl and list wheels without downloading")
    parser.add_argument("--no-hashes", action="store_true", help="do not add sha256 fragments to simple index links")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if args.workers < 1:
        print("--workers must be at least 1", file=sys.stderr)
        return 2

    source = canonical_url(args.source, directory=True)
    dest_root = args.dest.resolve()
    wheel_root = dest_root / source_local_root(source)

    wheels = crawl_wheels(source, timeout=args.timeout, retries=args.retries)
    print(f"found {len(wheels)} wheels")

    if args.dry_run:
        for url in wheels:
            print(url)
        return 0

    failures = 0
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as executor:
        future_to_url = {
            executor.submit(download_file, url, url_to_local_path(url, source, dest_root), timeout=args.timeout, retries=args.retries): url
            for url in wheels
        }
        for future in concurrent.futures.as_completed(future_to_url):
            url = future_to_url[future]
            try:
                print(future.result())
            except Exception as exc:
                failures += 1
                print(f"error {url}: {exc}", file=sys.stderr)

    generate_simple_index(dest_root, wheel_root, hashes=not args.no_hashes)
    if failures:
        print(f"completed with {failures} download failures", file=sys.stderr)
        return 1
    return 0


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