scripts/download_sources.py

#!/usr/bin/env python3
"""List or download approved source artifacts from the source manifest."""

from __future__ import annotations

import argparse
import hashlib
import json
import shutil
import ssl
import subprocess
import sys
import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path

import certifi
import pymupdf


PROJECT_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_MANIFEST = PROJECT_ROOT / "data" / "sources" / "download_manifest.json"
DEFAULT_OUTPUT = PROJECT_ROOT / "data" / "raw" / "sources"
TLS_CONTEXT = ssl.create_default_context(cafile=certifi.where())


def human_size(byte_count: int | None) -> str:
    if byte_count is None:
        return "unknown"
    value = float(byte_count)
    for unit in ("B", "KiB", "MiB", "GiB"):
        if value < 1024 or unit == "GiB":
            return f"{value:.1f} {unit}"
        value /= 1024
    raise AssertionError("unreachable")


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def load_manifest(path: Path) -> list[dict]:
    data = json.loads(path.read_text(encoding="utf-8"))
    if data.get("schema_version") != 1 or not isinstance(data.get("sources"), list):
        raise ValueError(f"Unsupported manifest: {path}")
    defaults = data.get("defaults", {})
    if not isinstance(defaults, dict):
        raise ValueError(f"Manifest defaults must be an object: {path}")
    return [{**defaults, **source} for source in data["sources"]]


def select_sources(sources: list[dict], requested: list[str]) -> list[dict]:
    if not requested:
        return sources
    by_id = {source["source_id"]: source for source in sources}
    unknown = sorted(set(requested) - set(by_id))
    if unknown:
        raise ValueError(f"Unknown source_id: {', '.join(unknown)}")
    return [by_id[source_id] for source_id in requested]


def validate_artifact(path: Path, filename: str) -> None:
    """Reject an HTML redirect or error page saved under a PDF filename."""
    if Path(filename).suffix.casefold() == ".pdf":
        with path.open("rb") as stream:
            signature = stream.read(5)
        if signature != b"%PDF-":
            raise ValueError(
                f"Expected a PDF for {filename}, received signature {signature!r}"
            )
        try:
            with pymupdf.open(path) as document:
                if document.page_count < 1:
                    raise ValueError(f"PDF has no readable pages: {filename}")
        except (pymupdf.FileDataError, RuntimeError) as error:
            raise ValueError(f"Unreadable PDF {filename}: {error}") from error


def curl_download(url: str, destination: Path) -> None:
    command = [
        "curl",
        "--fail",
        "--location",
        "--silent",
        "--show-error",
        "--proto",
        "=https",
        "--retry",
        "3",
        "--retry-delay",
        "2",
    ]
    if destination.exists() and destination.stat().st_size > 0:
        try:
            subprocess.run(
                command
                + ["--continue-at", "-", "--output", str(destination), url],
                check=True,
            )
            return
        except subprocess.CalledProcessError:
            print(f"[restart] server cannot resume {destination.name}")
            destination.unlink()
    subprocess.run(command + ["--output", str(destination), url], check=True)


def download(source: dict, output_dir: Path) -> dict:
    output_dir.mkdir(parents=True, exist_ok=True)
    destination = output_dir / source["filename"]
    partial = destination.with_suffix(destination.suffix + ".part")

    if destination.exists():
        print(f"[skip] {source['source_id']}: {destination.name} already exists")
    else:
        print(f"[download] {source['source_id']} -> {destination}")
        request = urllib.request.Request(
            source["url"], headers={"User-Agent": "controlai-source-fetcher/0.1"}
        )
        try:
            with urllib.request.urlopen(request, context=TLS_CONTEXT) as response, partial.open("wb") as stream:
                shutil.copyfileobj(response, stream, length=1024 * 1024)
            validate_artifact(partial, source["filename"])
            partial.replace(destination)
        except urllib.error.HTTPError as error:
            if error.code not in {403, 429}:
                partial.unlink(missing_ok=True)
                raise
            curl_download(source["url"], partial)
            validate_artifact(partial, source["filename"])
            partial.replace(destination)
        except urllib.error.URLError as error:
            partial.unlink(missing_ok=True)
            if not isinstance(error.reason, ssl.SSLCertVerificationError):
                raise
            curl_download(source["url"], partial)
            validate_artifact(partial, source["filename"])
            partial.replace(destination)
        except Exception:
            partial.unlink(missing_ok=True)
            raise

    validate_artifact(destination, source["filename"])
    actual_bytes = destination.stat().st_size
    expected_bytes = source.get("expected_bytes")
    if expected_bytes is not None and actual_bytes != expected_bytes:
        raise ValueError(
            f"Size mismatch for {source['source_id']}: "
            f"expected {expected_bytes}, got {actual_bytes}"
        )

    return {
        "source_id": source["source_id"],
        "title": source.get("title"),
        "authors": source.get("authors", []),
        "filename": source["filename"],
        "url": source["url"],
        "license": source["license"],
        "corpus_tier": source.get("corpus_tier"),
        "release": source.get("release"),
        "intended_use": source.get("intended_use"),
        "coverage": source.get("coverage", []),
        "bytes": actual_bytes,
        "sha256": sha256(destination),
        "downloaded_at": datetime.now(timezone.utc).isoformat(),
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--manifest", type=Path, default=DEFAULT_MANIFEST)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--source", action="append", default=[], help="Download only this source_id; repeatable")
    parser.add_argument(
        "--exclude-source",
        action="append",
        default=[],
        help="Skip this source_id; repeatable",
    )
    parser.add_argument("--download", action="store_true", help="Actually download; otherwise only list")
    parser.add_argument(
        "--continue-on-error",
        action="store_true",
        help="Record successful files and continue after a failed source",
    )
    args = parser.parse_args()

    all_sources = load_manifest(args.manifest)
    sources = select_sources(all_sources, args.source)
    known_ids = {source["source_id"] for source in all_sources}
    unknown_exclusions = sorted(set(args.exclude_source) - known_ids)
    if unknown_exclusions:
        raise ValueError(f"Unknown excluded source_id: {', '.join(unknown_exclusions)}")
    excluded = set(args.exclude_source)
    sources = [source for source in sources if source["source_id"] not in excluded]
    known_total = sum(source.get("expected_bytes") or 0 for source in sources)

    print(f"Manifest: {args.manifest}")
    print(f"Output:   {args.output}")
    for source in sources:
        print(
            f"- {source['source_id']}: {human_size(source.get('expected_bytes'))} "
            f"[{source['license']}]"
        )
    print(f"Known total: {human_size(known_total)} + unknown sizes")

    if not args.download:
        print("Dry run only. Add --download to fetch these files.")
        return 0

    lock_path = args.output / "source_lock.json"
    existing_records = []
    if lock_path.exists():
        existing_records = json.loads(lock_path.read_text(encoding="utf-8"))
    by_source_id = {record["source_id"]: record for record in existing_records}
    failures: list[tuple[str, str]] = []
    for source in sources:
        try:
            record = download(source, args.output)
        except Exception as error:
            if not args.continue_on_error:
                raise
            failures.append((source["source_id"], str(error)))
            print(f"[error] {source['source_id']}: {error}", file=sys.stderr)
            continue
        by_source_id[record["source_id"]] = record
        lock_records = [by_source_id[source_id] for source_id in sorted(by_source_id)]
        lock_path.write_text(
            json.dumps(lock_records, indent=2) + "\n", encoding="utf-8"
        )

    print(f"Wrote checksums and provenance: {lock_path}")
    if failures:
        print(f"Completed with {len(failures)} failed source(s).", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (OSError, ValueError, json.JSONDecodeError) as error:
        print(f"error: {error}", file=sys.stderr)
        raise SystemExit(1)