"""Dataset downloader — registry + SHA256-verified fetch."""

from __future__ import annotations

import hashlib
import urllib.request
from dataclasses import dataclass
from pathlib import Path


@dataclass(frozen=True)
class DatasetSpec:
    url: str
    sha256: str
    target_subdir: str
    unpack: bool = True


# REGISTRY NOTES:
# Entries with url="" are manual-download-only. The ETH Research Collection
# (DOI 10.3929/ethz-b-000690084) gates downloads behind a DSpace UI that
# doesn't expose a stable direct URL, so the bundle has to be fetched by
# hand and placed beside the project. SHA256 remains authoritative: the
# download helper refuses anything that doesn't match.
DATASET_REGISTRY: dict[str, DatasetSpec] = {
    "euroc_machine_hall": DatasetSpec(
        # 12.6 GB bundle containing MH_01 … MH_05 (each as inner .zip + .bag).
        # Download from ETH Research Collection:
        #   https://doi.org/10.3929/ethz-b-000690084
        # After fetching, unpack the inner MH_0N_easy.zip of interest into
        # datasets/euroc/MH_0N/ so the adapter finds mav0/.
        url="",  # manual download — see comment above
        sha256="5ed7d07903f8d19b6c8808e2ae8a0872b281f6e34ef5497023b8ac58c3de0f6f",
        target_subdir="euroc",
        unpack=False,  # the bundle itself is not unpacked end-to-end; see README
    ),
    "vpair_sample": DatasetSpec(
        url="",  # manual download only — see Zenodo link on
                 # https://github.com/AerVisLoc/vpair
        sha256="c83f41fa7c9f39debe4898abceb009bff8a2b17b4840defb1831ecf127a9bf6b",
        target_subdir="vpair/sample",
        unpack=True,
    ),
}


def verify_sha256(path: Path, expected_hex: str) -> bool:
    """Return True iff SHA256 of `path` hex-matches `expected_hex`."""
    h = hashlib.sha256()
    with path.open("rb") as fh:
        for chunk in iter(lambda: fh.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest() == expected_hex


def download_dataset(spec: DatasetSpec, dest: Path) -> Path:
    """Download `spec.url` to `dest` if not already present and valid.

    If `dest` already exists and its SHA256 matches `spec.sha256`, return
    immediately without touching the network. Otherwise download, verify,
    and raise RuntimeError on hash mismatch (deleting the bad file).
    """
    if dest.exists() and verify_sha256(dest, spec.sha256):
        return dest
    dest.parent.mkdir(parents=True, exist_ok=True)
    tmp = dest.with_suffix(dest.suffix + ".part")
    urllib.request.urlretrieve(spec.url, tmp)
    tmp.rename(dest)
    if not verify_sha256(dest, spec.sha256):
        dest.unlink()
        raise RuntimeError(f"SHA256 mismatch after download of {spec.url}")
    return dest
