#!/usr/bin/env python3

# This ChatGPT-generated sciprt helps us check the downloads specified
# in our JSON file.  In this directory, run:
#   ./check_package_files.py ../package_pololu_stm32_index.json
#
# This will use the files already downloaded to this directory if possible,
# which saves a lot of time but does not detect broken downloads.
#
# To detect broken downloads, run this script in a new, empty
# directory and afterwards delete the files it downaloaded.

import argparse
import concurrent.futures
import hashlib
import json
import os
import sys
import tempfile
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import Any

@dataclass(frozen=True)
class ArchiveRef:
    kind: str
    package: str
    name: str
    version: str
    url: str
    archive_file_name: str
    checksum: str
    size: int


def parse_checksum(checksum: str) -> tuple[str, str]:
    """
    Arduino package indexes usually use checksums like:
        SHA-256:012345abcdef...
    """
    if ":" not in checksum:
        raise ValueError(f"unsupported checksum format: {checksum!r}")

    algorithm, digest = checksum.split(":", 1)
    algorithm = algorithm.strip().lower().replace("-", "")
    digest = digest.strip().lower()

    if algorithm != "sha256":
        raise ValueError(f"unsupported checksum algorithm: {algorithm!r}")

    if len(digest) != 64:
        raise ValueError(f"bad SHA-256 digest length: {checksum!r}")

    return algorithm, digest


def require_field(obj: dict[str, Any], field: str, context: str) -> Any:
    if field not in obj:
        raise ValueError(f"{context} is missing required field {field!r}")
    return obj[field]


def iter_archive_refs(index: dict[str, Any]) -> list[ArchiveRef]:
    refs: list[ArchiveRef] = []

    for package in index.get("packages", []):
        package_name = package.get("name", "<unknown-package>")

        for platform in package.get("platforms", []):
            context = f"platform {package_name}:{platform.get('architecture', '?')}:{platform.get('version', '?')}"
            refs.append(
                ArchiveRef(
                    kind="platform",
                    package=package_name,
                    name=str(platform.get("architecture", "<unknown-arch>")),
                    version=str(platform.get("version", "<unknown-version>")),
                    url=str(require_field(platform, "url", context)),
                    archive_file_name=str(require_field(platform, "archiveFileName", context)),
                    checksum=str(require_field(platform, "checksum", context)),
                    size=int(require_field(platform, "size", context)),
                )
            )

        for tool in package.get("tools", []):
            tool_name = tool.get("name", "<unknown-tool>")
            tool_version = tool.get("version", "<unknown-version>")

            for system in tool.get("systems", []):
                host = system.get("host", "<unknown-host>")
                context = f"tool {package_name}:{tool_name}:{tool_version}:{host}"
                refs.append(
                    ArchiveRef(
                        kind="tool",
                        package=package_name,
                        name=str(tool_name),
                        version=str(tool_version),
                        url=str(require_field(system, "url", context)),
                        archive_file_name=str(require_field(system, "archiveFileName", context)),
                        checksum=str(require_field(system, "checksum", context)),
                        size=int(require_field(system, "size", context)),
                    )
                )

    return refs


def check_conflicts(refs: list[ArchiveRef]) -> None:
    """
    The same archiveFileName may appear multiple times. That's OK only if
    url/checksum/size all agree.
    """
    by_name: dict[str, ArchiveRef] = {}

    for ref in refs:
        old = by_name.get(ref.archive_file_name)
        if old is None:
            by_name[ref.archive_file_name] = ref
            continue

        if (
            old.url != ref.url
            or old.checksum != ref.checksum
            or old.size != ref.size
        ):
            raise ValueError(
                "conflicting entries for archiveFileName "
                f"{ref.archive_file_name!r}:\n"
                f"  old: url={old.url!r}, checksum={old.checksum!r}, size={old.size}\n"
                f"  new: url={ref.url!r}, checksum={ref.checksum!r}, size={ref.size}"
            )


def dedupe_by_archive_name(refs: list[ArchiveRef]) -> list[ArchiveRef]:
    result: dict[str, ArchiveRef] = {}
    for ref in refs:
        result.setdefault(ref.archive_file_name, ref)
    return list(result.values())


def download_one(ref: ArchiveRef, dest_dir: Path) -> None:
    dest = dest_dir / ref.archive_file_name

    if dest.exists():
        return

    dest.parent.mkdir(parents=True, exist_ok=True)

    fd, tmp_name = tempfile.mkstemp(
        prefix=dest.name + ".",
        suffix=".tmp",
        dir=str(dest.parent),
    )
    os.close(fd)

    tmp_path = Path(tmp_name)

    try:
        print(f"Downloading {ref.archive_file_name}")
        print(f"  {ref.url}")

        request = urllib.request.Request(
            ref.url,
            headers={
                "User-Agent": "arduino-package-index-checker/1.0",
            },
        )

        with urllib.request.urlopen(request, timeout=60) as response:
            with tmp_path.open("wb") as out:
                while True:
                    chunk = response.read(1024 * 1024)
                    if not chunk:
                        break
                    out.write(chunk)

        # Atomic on the same filesystem.
        tmp_path.replace(dest)

    except Exception:
        tmp_path.unlink(missing_ok=True)
        raise


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()

    with path.open("rb") as f:
        while True:
            chunk = f.read(1024 * 1024)
            if not chunk:
                break
            h.update(chunk)

    return h.hexdigest()


def verify_one(ref: ArchiveRef, dest_dir: Path) -> list[str]:
    errors: list[str] = []
    path = dest_dir / ref.archive_file_name

    if not path.exists():
        return [f"{ref.archive_file_name}: file is missing"]

    actual_size = path.stat().st_size
    if actual_size != ref.size:
        errors.append(
            f"{ref.archive_file_name}: size mismatch: "
            f"expected {ref.size}, got {actual_size}"
        )

    try:
        _, expected_sha256 = parse_checksum(ref.checksum)
    except ValueError as e:
        errors.append(f"{ref.archive_file_name}: {e}")
        return errors

    actual_sha256 = sha256_file(path)
    if actual_sha256.lower() != expected_sha256.lower():
        errors.append(
            f"{ref.archive_file_name}: sha256 mismatch:\n"
            f"  expected {expected_sha256}\n"
            f"  got      {actual_sha256}"
        )

    return errors


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Download and verify files referenced by an Arduino package index JSON file."
    )
    parser.add_argument("package_index_json", help="Path to Arduino package index JSON")
    parser.add_argument(
        "-j",
        "--jobs",
        type=int,
        default=4,
        help="Number of parallel downloads/verifications, default: 4",
    )
    args = parser.parse_args()

    index_path = Path(args.package_index_json)
    dest_dir = Path.cwd()

    with index_path.open("r", encoding="utf-8") as f:
        index = json.load(f)

    refs = iter_archive_refs(index)
    check_conflicts(refs)

    unique_refs = dedupe_by_archive_name(refs)

    print(f"Found {len(refs)} archive references.")
    print(f"Found {len(unique_refs)} unique archive files.")

    missing_refs = [
        ref for ref in unique_refs
        if not (dest_dir / ref.archive_file_name).exists()
    ]

    if missing_refs:
        print(f"{len(missing_refs)} files are missing; downloading...")
    else:
        print("No downloads needed; all files already exist.")

    download_errors: list[str] = []

    with concurrent.futures.ThreadPoolExecutor(max_workers=args.jobs) as executor:
        future_to_ref = {
            executor.submit(download_one, ref, dest_dir): ref
            for ref in missing_refs
        }

        for future in concurrent.futures.as_completed(future_to_ref):
            ref = future_to_ref[future]
            try:
                future.result()
            except Exception as e:
                download_errors.append(
                    f"{ref.archive_file_name}: download failed: {e}"
                )

    if download_errors:
        print()
        print("Download errors:")
        for error in download_errors:
            print(f"  {error}")

    print()
    print("Verifying size and SHA-256...")

    verify_errors: list[str] = []

    with concurrent.futures.ThreadPoolExecutor(max_workers=args.jobs) as executor:
        future_to_ref = {
            executor.submit(verify_one, ref, dest_dir): ref
            for ref in unique_refs
        }

        for future in concurrent.futures.as_completed(future_to_ref):
            verify_errors.extend(future.result())

    if verify_errors:
        print()
        print("Verification errors:")
        for error in verify_errors:
            print(f"  {error}")
    else:
        print("All files verified successfully.")

    if download_errors or verify_errors:
        return 1

    return 0


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