#!/usr/bin/env python3
"""Fetch exactly the selected public upstream SCI/DQ ranges, not a mirrored cube.

Python 3.9+ standard library. This downloader grants no new pixel licence.
The website distributes code, headers and observation facts, not these pixels.
No authentication, private checkout, pyKLIP or protected Lab Engine is required.
"""
import argparse
import hashlib
import json
from pathlib import Path
from urllib.request import Request, urlopen


def checked_range(url, start, length, source_bytes, expected_sha, destination):
    if destination.exists():
        data = destination.read_bytes()
        if len(data) == length and hashlib.sha256(data).hexdigest() == expected_sha:
            return data
        raise ValueError("Existing cache failed hash: " + str(destination))
    request = Request(
        url + "?ac22_exact_range=" + str(start),
        headers={
            "Range": "bytes=%d-%d" % (start, start + length - 1),
            "User-Agent": "SandboxPhysics-Public-Tutorial/1.0",
        },
    )
    with urlopen(request, timeout=60) as response:
        required = "bytes %d-%d/%d" % (start, start + length - 1, source_bytes)
        if response.status != 206 or response.headers.get("Content-Range") != required:
            raise ValueError("Server ignored exact byte range; no full-file fallback")
        data = response.read(length + 1)
    if len(data) != length or hashlib.sha256(data).hexdigest() != expected_sha:
        raise ValueError("Wrong source length or SHA256")
    destination.write_bytes(data)
    return data


def download(manifest_path, output):
    manifest = json.loads(manifest_path.read_text())
    output.mkdir(parents=True, exist_ok=True)
    total = sum(
        entry["header_bytes"]
        + entry["dq_header"]["bytes"]
        + sum(block["length"] for block in entry["ranges"] + entry["quality_ranges"])
        for entry in manifest["files"]
    )
    if total >= 25 * 1024 * 1024:
        raise ValueError("Selection exceeds the declared download budget")
    print(
        "Exact science/DQ/header selection: %d bytes; cache hits avoid repeat downloads"
        % total
    )
    base = manifest["official_upstream_url"]
    for entry in manifest["files"]:
        url = base + entry["file"]
        number = entry["file"].split("S20131210S")[1][:4]
        checked_range(
            url,
            0,
            entry["header_bytes"],
            entry["source_bytes"],
            entry["header_sha256"],
            output / (number + "-headers.bin"),
        )
        header = entry["dq_header"]
        checked_range(
            url,
            header["start"],
            header["bytes"],
            entry["source_bytes"],
            header["sha256"],
            output / (number + "-dq-header.bin"),
        )
        for channel, science, quality in zip(
            manifest["channels"], entry["ranges"], entry["quality_ranges"]
        ):
            for kind, block in [("sci", science), ("dq", quality)]:
                checked_range(
                    url,
                    block["start"],
                    block["length"],
                    entry["source_bytes"],
                    block["sha256"],
                    output / ("%s-%s-%02d.bin" % (number, kind, channel)),
                )
        print("PASS", entry["file"], "selected ranges and original headers")
    return total


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--manifest", type=Path, default=Path("data-source-manifest.json")
    )
    parser.add_argument("--output", type=Path, default=Path("gpi-cache"))
    args = parser.parse_args()
    download(args.manifest, args.output)