"""Resumable, checksummed, node-shardable downloader for the HF dataset.

Design (CLAUDE.md §4.1): dry-run, month-range filters, retry with exponential
backoff, resumability (hf_hub native), sha256 verification against LFS metadata,
per-file logs, disk-space checks, idempotence, node-aware sharding
(--shard k/n downloads files where index % n == k, so cluster nodes split work
without coordination).
"""

from __future__ import annotations

import hashlib
import os
import shutil
import time
from datetime import datetime, timezone
from pathlib import Path

from huggingface_hub import HfApi, hf_hub_download

from . import paths
from .logging_utils import get_logger
from .manifest import FileRecord, Manifest

CHUNK = 1 << 20
MIN_FREE_BYTES = 20 * (1 << 30)  # refuse to start a file with <20 GB free


def _sha256(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(CHUNK), b""):
            h.update(chunk)
    return h.hexdigest()


def month_range(start: str, end: str) -> list[tuple[int, int]]:
    y0, m0 = map(int, start.split("-"))
    y1, m1 = map(int, end.split("-"))
    out = []
    y, m = y0, m0
    while (y, m) <= (y1, m1):
        out.append((y, m))
        m += 1
        if m == 13:
            y, m = y + 1, 1
    return out


def build_manifest(start: str, end: str) -> Manifest:
    """Populate the manifest from pinned repo metadata (idempotent)."""
    cfg = paths.base_config()["dataset"]
    api = HfApi()
    info = api.dataset_info(cfg["id"], revision=cfg["revision"], files_metadata=True)
    meta = {s.rfilename: s for s in info.siblings}
    man = Manifest(paths.get("manifests") / "raw_files.json")
    for y, m in month_range(start, end):
        remote = cfg["file_pattern"].format(year=y, month=m)
        if remote not in meta:
            raise FileNotFoundError(f"{remote} not in repo at pinned revision")
        s = meta[remote]
        existing = man.records.get(remote)
        if existing and existing.status == "complete":
            continue
        man.upsert(FileRecord(
            dataset_id=cfg["id"], dataset_revision=cfg["revision"],
            remote_path=remote,
            local_path=str(Path("data/raw") / remote),
            year=y, month=m,
            expected_bytes=s.size,
            lfs_sha256=(s.lfs.sha256 if s.lfs else None),
        ))
    man.save()
    return man


def download(start: str, end: str, *, dry_run: bool = False,
             shard: str | None = None, max_retries: int = 5) -> None:
    cfg = paths.base_config()["dataset"]
    log = get_logger("download", paths.get("logs") / "download" / "download.jsonl")
    man = build_manifest(start, end)
    todo = sorted(man.pending(), key=lambda r: (r.year, r.month))
    if shard:
        k, n = map(int, shard.split("/"))
        todo = [r for i, r in enumerate(todo) if i % n == k]
    total = sum(r.expected_bytes or 0 for r in todo)
    log.info(f"{len(todo)} files to download, {total/1e9:.1f} GB"
             + (f" (shard {shard})" if shard else ""))
    if dry_run:
        for r in todo:
            log.info(f"DRY RUN would fetch {r.remote_path} ({(r.expected_bytes or 0)/1e6:.0f} MB)")
        return

    raw_root = paths.get("raw")
    for rec in todo:
        if shutil.disk_usage(raw_root).free < MIN_FREE_BYTES:
            log.error("disk free below safety threshold; stopping")
            break
        rec.status = "downloading"
        rec.download_started_at = datetime.now(timezone.utc).isoformat()
        man.upsert(rec); man.save()
        for attempt in range(max_retries):
            try:
                local = hf_hub_download(
                    cfg["id"], rec.remote_path, repo_type="dataset",
                    revision=cfg["revision"], local_dir=raw_root,
                    token=os.environ.get("HF_TOKEN"),
                )
                digest = _sha256(Path(local))
                if rec.lfs_sha256 and digest != rec.lfs_sha256:
                    raise IOError(f"checksum mismatch for {rec.remote_path}")
                rec.sha256 = digest
                rec.downloaded_bytes = Path(local).stat().st_size
                rec.status = "complete"
                rec.download_completed_at = datetime.now(timezone.utc).isoformat()
                log.info(f"complete {rec.remote_path} ({rec.downloaded_bytes/1e6:.0f} MB)")
                break
            except Exception as e:  # noqa: BLE001 — retry any transient failure
                wait = 2 ** attempt
                log.warning(f"attempt {attempt+1}/{max_retries} failed for "
                            f"{rec.remote_path}: {e}; retrying in {wait}s")
                time.sleep(wait)
        else:
            rec.status = "failed"
            log.error(f"FAILED {rec.remote_path} after {max_retries} attempts")
        man.upsert(rec); man.save()
    man.export_tabular()


def verify(start: str, end: str) -> bool:
    """Re-hash every complete file against LFS sha256. Returns True if all pass."""
    log = get_logger("verify", paths.get("logs") / "download" / "verify.jsonl")
    man = Manifest(paths.get("manifests") / "raw_files.json")
    ok = True
    for rec in sorted(man.records.values(), key=lambda r: (r.year, r.month)):
        p = paths.repo_root() / rec.local_path
        if rec.status != "complete":
            log.warning(f"not complete: {rec.remote_path} ({rec.status})")
            ok = False
            continue
        if not p.exists():
            log.error(f"missing local file: {p}")
            rec.status = "failed"; ok = False; continue
        digest = _sha256(p)
        if rec.lfs_sha256 and digest != rec.lfs_sha256:
            log.error(f"checksum mismatch: {rec.remote_path}")
            rec.status = "failed"; ok = False
    man.save()
    log.info("verification " + ("PASSED" if ok else "FAILED"))
    return ok