"""Reliable, resumable downloads and streaming tar extraction (stdlib only).

Benchmark data lives on public hosts (Hugging Face and others) and ranges from a few
hundred KB to a 36 GB tar. Colab sessions drop connections, so two tools are provided:

* :func:`download` fetches one file to disk. It resumes from a ``<name>.part`` file with an
  HTTP ``Range`` request, retries transient network failures, verifies size and SHA-256 and
  only then atomically renames the ``.part`` file to its final name. A completed file is
  never overwritten and a partial file never looks complete.
* :func:`stream_extract` reads a (possibly huge, possibly compressed) tar archive as a
  stream and writes only the selected regular files to disk, without ever storing the
  archive itself.

All network access goes through an injectable :data:`Opener`, so tests run offline.
:func:`hf_resolve_url` builds revision-pinned Hugging Face download URLs.
"""

from __future__ import annotations

import hashlib
import http.client
import logging
import os
import re
import shutil
import ssl
import tarfile
import tempfile
import time
import urllib.error
import urllib.parse
import urllib.request
import zlib
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from types import TracebackType
from typing import IO, Literal, NoReturn, Protocol, cast

__all__ = [
    "Downloaded",
    "FetchError",
    "HeadersLike",
    "Opener",
    "Response",
    "default_opener",
    "download",
    "hf_resolve_url",
    "stream_extract",
]

logger = logging.getLogger(__name__)

_USER_AGENT = "nida-bench/0"
_SHA256_RE = re.compile(r"[0-9a-f]{64}")
_COMMIT_SHA_RE = re.compile(r"[0-9a-f]{40}")
_CONTENT_RANGE_RE = re.compile(r"bytes (\d+)-(\d+)/(\d+|\*)")
_DRIVE_LETTER_RE = re.compile(r"[A-Za-z]:")
_MAX_BACKOFF_S = 30.0
_COPY_CHUNK_BYTES = 1 << 20
# HTTP 4xx codes that are worth retrying; every other 4xx is a permanent client error.
_RETRYABLE_4XX = frozenset({408, 429})

# Transient network failures. ``HTTPError`` is a ``URLError`` and is classified separately.
_NETWORK_ERRORS: tuple[type[BaseException], ...] = (
    urllib.error.URLError,
    ConnectionError,
    TimeoutError,
    http.client.HTTPException,  # includes IncompleteRead
    ssl.SSLError,
)
# Everything that can go wrong while reading an archive stream (network or corrupt data).
_STREAM_ERRORS: tuple[type[BaseException], ...] = (
    *_NETWORK_ERRORS,
    tarfile.TarError,
    EOFError,
    zlib.error,
)


class FetchError(RuntimeError):
    """A download or extraction failed; the message carries URL/path and byte counts."""


class _TransientError(Exception):
    """Internal: this attempt failed but a resumed attempt may succeed."""

    def __init__(self, message: str, *, expected_bytes: int | None = None) -> None:
        super().__init__(message)
        self.expected_bytes = expected_bytes


class HeadersLike(Protocol):
    """Response headers: only ``get`` is used (``http.client.HTTPMessage`` satisfies it)."""

    def get(self, name: str, default: str | None = None, /) -> str | None:
        """Return the header value, or ``default`` when absent."""
        ...


class Response(Protocol):
    """The slice of an HTTP response this module relies on.

    ``http.client.HTTPResponse`` (what ``urllib.request.urlopen`` returns) satisfies it.
    """

    @property
    def status(self) -> int:
        """HTTP status code."""
        ...

    @property
    def headers(self) -> HeadersLike:
        """Response headers."""
        ...

    def read(self, amt: int = -1, /) -> bytes:
        """Read up to ``amt`` bytes; ``b""`` means end of stream."""
        ...

    def __enter__(self) -> Response:
        """Enter the context."""
        ...

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
        /,
    ) -> object:
        """Close the response."""
        ...


Opener = Callable[[urllib.request.Request, float], Response]
"""Callable performing a request with a timeout in seconds and returning a response."""


def default_opener(request: urllib.request.Request, timeout_s: float) -> Response:
    """Open ``request`` with :func:`urllib.request.urlopen`.

    Args:
        request: The prepared request.
        timeout_s: Socket timeout in seconds.

    Returns:
        The open HTTP response (a context manager).
    """
    response: Response = urllib.request.urlopen(request, timeout=timeout_s)
    return response


@dataclass(frozen=True)
class Downloaded:
    """A completed, verified download.

    Attributes:
        path: Final file path.
        size_bytes: File size in bytes.
        sha256: Lowercase hex SHA-256 computed over the final file.
        url: Source URL.
    """

    path: Path
    size_bytes: int
    sha256: str
    url: str


# --------------------------------------------------------------------------------------
# download
# --------------------------------------------------------------------------------------


def download(
    url: str,
    dest: Path,
    *,
    expected_sha256: str | None = None,
    expected_size: int | None = None,
    opener: Opener = default_opener,
    chunk_bytes: int = 1 << 20,
    timeout_s: float = 60.0,
    max_attempts: int = 5,
    sleep: Callable[[float], None] = time.sleep,
) -> Downloaded:
    """Download ``url`` to ``dest``, resuming and retrying as needed.

    If ``dest`` already exists it is verified (size and/or SHA-256 when given) and returned
    without any network access; it is never overwritten. Otherwise the data is written to
    ``<dest>.part`` and resumed with ``Range: bytes=<n>-`` on every retry. A server that
    answers 200 instead of 206 to a ranged request makes the download restart from zero.
    Transient failures (``URLError``, ``ConnectionError``, ``TimeoutError``, short reads,
    HTTP 5xx/408/429) are retried with backoff of 1, 2, 4, ... seconds (capped at 30). Other
    HTTP 4xx errors fail immediately. When all attempts are exhausted the ``.part`` file is
    kept so a later call can resume it.

    The SHA-256 is computed over the final file (re-read once at the end), so it also
    covers bytes written by earlier, interrupted attempts.

    Concurrent downloads of the same ``dest`` are not supported.

    Args:
        url: ``http://`` or ``https://`` URL.
        dest: Final file path; parent directories are created.
        expected_sha256: Expected lowercase/uppercase hex SHA-256 (64 chars), if known.
        expected_size: Expected size in bytes, if known.
        opener: Network access function (injectable for tests).
        chunk_bytes: Read chunk size in bytes.
        timeout_s: Socket timeout per request, seconds.
        max_attempts: Total number of attempts (first try included), at least 1.
        sleep: Sleep function used for backoff between attempts.

    Returns:
        The verified :class:`Downloaded` record.

    Raises:
        ValueError: On invalid arguments (bad URL scheme, malformed hash, ...).
        FetchError: On verification failure, permanent HTTP errors or exhausted attempts.
        OSError: On local filesystem failures (e.g. disk full).
    """
    _require_http_url(url)
    sha_expected = _normalise_sha256(expected_sha256)
    if expected_size is not None and expected_size < 0:
        raise ValueError(f"expected_size must be >= 0, got {expected_size}")
    if chunk_bytes < 1:
        raise ValueError(f"chunk_bytes must be >= 1, got {chunk_bytes}")
    if max_attempts < 1:
        raise ValueError(f"max_attempts must be >= 1, got {max_attempts}")
    if timeout_s <= 0:
        raise ValueError(f"timeout_s must be > 0, got {timeout_s}")

    if dest.exists():
        return _verify_existing(url, dest, sha_expected, expected_size)

    dest.parent.mkdir(parents=True, exist_ok=True)
    part = dest.with_name(dest.name + ".part")

    last_error: _TransientError | None = None
    known_total = expected_size  # best knowledge of the final size, for error messages
    for attempt in range(1, max_attempts + 1):
        try:
            _attempt(url, part, expected_size, opener, chunk_bytes, timeout_s)
        except _TransientError as exc:
            last_error = exc
            if exc.expected_bytes is not None:
                known_total = exc.expected_bytes
            part_bytes = part.stat().st_size if part.exists() else 0
            logger.warning(
                "download %s attempt %d/%d failed after %d bytes: %s",
                url,
                attempt,
                max_attempts,
                part_bytes,
                exc,
            )
            if attempt < max_attempts:
                sleep(min(2.0 ** (attempt - 1), _MAX_BACKOFF_S))
            continue
        last_error = None
        break

    if last_error is not None:
        part_bytes = part.stat().st_size if part.exists() else 0
        raise FetchError(
            f"download of {url} failed after {max_attempts} attempts: {part_bytes} bytes "
            f"received in {part} (expected "
            f"{known_total if known_total is not None else 'unknown'}); last error: {last_error}"
        ) from last_error

    size_bytes = part.stat().st_size
    digest = _sha256_file(part)
    if sha_expected is not None and digest != sha_expected:
        part.unlink()
        raise FetchError(
            f"sha256 mismatch for {url}: got {digest}, expected {sha_expected} "
            f"({size_bytes} bytes; partial file deleted)"
        )
    part.replace(dest)
    logger.info("downloaded %s -> %s (%d bytes, sha256 %s)", url, dest, size_bytes, digest)
    return Downloaded(path=dest, size_bytes=size_bytes, sha256=digest, url=url)


def _require_http_url(url: str) -> None:
    if not url.lower().startswith(("http://", "https://")):
        raise ValueError(f"only http(s) URLs are supported, got {url!r}")


def _normalise_sha256(value: str | None) -> str | None:
    if value is None:
        return None
    lowered = value.strip().lower()
    if _SHA256_RE.fullmatch(lowered) is None:
        raise ValueError(f"expected_sha256 must be 64 hex characters, got {value!r}")
    return lowered


def _sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while chunk := handle.read(_COPY_CHUNK_BYTES):
            digest.update(chunk)
    return digest.hexdigest()


def _verify_existing(
    url: str, dest: Path, expected_sha256: str | None, expected_size: int | None
) -> Downloaded:
    if not dest.is_file():
        raise FetchError(f"destination {dest} exists but is not a regular file (url {url})")
    size_bytes = dest.stat().st_size
    if expected_size is not None and size_bytes != expected_size:
        raise FetchError(
            f"existing file {dest} has {size_bytes} bytes, expected {expected_size} "
            f"(url {url}); refusing to overwrite, delete it manually to re-download"
        )
    digest = _sha256_file(dest)
    if expected_sha256 is not None and digest != expected_sha256:
        raise FetchError(
            f"existing file {dest} has sha256 {digest}, expected {expected_sha256} "
            f"(url {url}); refusing to overwrite, delete it manually to re-download"
        )
    logger.info("download skipped, %s already present (%d bytes)", dest, size_bytes)
    return Downloaded(path=dest, size_bytes=size_bytes, sha256=digest, url=url)


def _attempt(
    url: str,
    part: Path,
    expected_size: int | None,
    opener: Opener,
    chunk_bytes: int,
    timeout_s: float,
) -> None:
    """Run one (possibly resumed) transfer into ``part``.

    Returns normally when ``part`` holds the complete body. Raises ``_TransientError`` when
    a resumed attempt may help and ``FetchError`` for permanent problems.
    """
    offset = part.stat().st_size if part.exists() else 0
    if expected_size is not None:
        if offset > expected_size:
            logger.warning("%s is larger than expected size %d, restarting", part, expected_size)
            part.unlink()
            offset = 0
        elif offset == expected_size:
            return  # everything is already on disk; only verification is left

    headers = {"User-Agent": _USER_AGENT}
    if offset > 0:
        headers["Range"] = f"bytes={offset}-"
    request = urllib.request.Request(url, headers=headers)

    try:
        response = opener(request, timeout_s)
    except urllib.error.HTTPError as exc:
        code = exc.code
        exc.close()
        _raise_for_http_status(url, part, code, offset)
    except _NETWORK_ERRORS as exc:
        raise _TransientError(f"connecting to {url} failed: {exc!r}") from exc

    with response:
        status = response.status
        if status not in (200, 206):
            raise FetchError(f"unexpected HTTP status {status} from {url}")
        if status == 200 and offset > 0:
            logger.warning("server ignored Range for %s, restarting from byte 0", url)
            offset = 0
        total = _expected_total(url, response, status, offset)
        if total is not None and expected_size is not None and total != expected_size:
            part.unlink(missing_ok=True)
            raise FetchError(
                f"{url} is {total} bytes according to the server, expected {expected_size}"
            )

        position = offset
        mode = "ab" if (status == 206 and offset > 0) else "wb"
        with part.open(mode) as handle:
            try:
                while chunk := response.read(chunk_bytes):
                    handle.write(chunk)
                    position += len(chunk)
                    if total is not None and position > total:
                        break
            except _NETWORK_ERRORS as exc:
                raise _TransientError(
                    f"reading {url} failed at byte {position}: {exc!r}",
                    expected_bytes=total,
                ) from exc
            handle.flush()
            os.fsync(handle.fileno())

    if total is not None and position != total:
        if position > total:
            part.unlink()
            raise _TransientError(
                f"{url} sent more than its Content-Length: {position} bytes, expected {total}",
                expected_bytes=total,
            )
        raise _TransientError(
            f"incomplete body from {url}: {position} bytes, expected {total}",
            expected_bytes=total,
        )
    if expected_size is not None and position != expected_size:
        if position > expected_size:
            part.unlink()
            raise FetchError(
                f"{url} delivered {position} bytes, expected {expected_size} (partial file deleted)"
            )
        raise _TransientError(
            f"stream from {url} ended early: {position} bytes, expected {expected_size}",
            expected_bytes=expected_size,
        )


def _raise_for_http_status(url: str, part: Path, code: int, offset: int) -> NoReturn:
    if code == 416 and offset > 0:
        # The stale .part is at or beyond the end of the remote file: start over.
        part.unlink(missing_ok=True)
        raise _TransientError(f"HTTP 416 for ranged request at byte {offset}; restarting")
    if 400 <= code < 500 and code not in _RETRYABLE_4XX:
        raise FetchError(f"HTTP {code} for {url}")
    raise _TransientError(f"HTTP {code} for {url}")


def _parse_length(url: str, name: str, value: str) -> int:
    try:
        parsed = int(value)
    except ValueError:
        raise FetchError(f"malformed {name} header from {url}: {value!r}") from None
    if parsed < 0:
        raise FetchError(f"malformed {name} header from {url}: {value!r}")
    return parsed


def _expected_total(url: str, response: Response, status: int, offset: int) -> int | None:
    """Return the total file size implied by the response headers, if the server says."""
    length_header = response.headers.get("Content-Length")
    length = _parse_length(url, "Content-Length", length_header) if length_header else None
    if status == 200:
        return length
    range_header = response.headers.get("Content-Range")
    if range_header is None:
        raise FetchError(f"206 response from {url} has no Content-Range header")
    match = _CONTENT_RANGE_RE.fullmatch(range_header.strip())
    if match is None:
        raise FetchError(f"malformed Content-Range from {url}: {range_header!r}")
    start = int(match.group(1))
    if start != offset:
        raise FetchError(
            f"server returned bytes starting at {start} from {url}, requested {offset}"
        )
    if match.group(3) != "*":
        return int(match.group(3))
    return offset + length if length is not None else None


# --------------------------------------------------------------------------------------
# stream_extract
# --------------------------------------------------------------------------------------


def stream_extract(
    url_or_path: str | Path,
    dest_dir: Path,
    select: Callable[[str], bool],
    *,
    opener: Opener = default_opener,
    timeout_s: float = 60.0,
) -> list[Path]:
    """Extract selected regular files from a tar archive while streaming it.

    The archive (plain, gzip, bz2 or xz; auto-detected) is read sequentially from a URL or a
    local path and is never stored. A stream cannot be resumed mid-archive: if reading fails
    the call raises :class:`FetchError` and must simply be run again. Files extracted before
    the failure are kept and skipped quickly on the rerun (the archive is read again from
    the start, so the whole transfer is repeated).

    For every regular-file member the name is normalised (leading ``./`` and empty or ``.``
    components removed) and validated. Unsafe names (empty, absolute, containing ``..``
    components, backslashes, NUL bytes or a Windows drive prefix) raise
    :class:`FetchError` even when ``select`` would reject them: an archive with such members
    is treated as hostile. Selected members are written to ``dest_dir / name`` through a
    temporary file in the same directory and an atomic rename. A target that already exists
    with the member's size is left untouched. Non-regular members (directories, symlinks,
    hard links, devices) are never written; those that ``select`` accepts are counted in the
    log message. If a name occurs twice in the archive the last occurrence wins, as with
    ``tar``.

    Args:
        url_or_path: ``http(s)://`` URL (fetched with ``opener``) or a local file path.
        dest_dir: Output directory; created if missing.
        select: Predicate on the normalised member name (``/``-separated, relative).
        opener: Network access function (injectable for tests); unused for local paths.
        timeout_s: Socket timeout per request, seconds.

    Returns:
        Paths of the extracted or already present files, in archive order, without
        duplicates.

    Raises:
        FetchError: On unsafe member names, unreadable/corrupt/truncated archives and network
            failures.
        OSError: On local filesystem failures (e.g. disk full).
    """
    if timeout_s <= 0:
        raise ValueError(f"timeout_s must be > 0, got {timeout_s}")
    dest_dir.mkdir(parents=True, exist_ok=True)
    root = dest_dir.resolve()
    label = str(url_or_path)

    result: dict[str, Path] = {}  # normalised name -> path; insertion order == archive order
    written_this_run: set[str] = set()
    scanned = 0
    extracted = 0
    already_present = 0
    ignored_non_regular = 0

    with _open_source(url_or_path, opener, timeout_s) as stream:
        try:
            # tarfile in "r|*" mode only ever calls ``stream.read``; urllib responses are not
            # nominally IO[bytes], hence the cast.
            with tarfile.open(fileobj=cast(IO[bytes], stream), mode="r|*") as archive:
                for member in archive:
                    scanned += 1
                    if not member.isreg():
                        if select(_collapse_name(member.name)):
                            ignored_non_regular += 1
                        _forget_members(archive)
                        continue
                    name = _safe_member_name(member.name)
                    if select(name):
                        target = _resolve_target(root, dest_dir, name)
                        if (
                            name not in written_this_run
                            and target.is_file()
                            and target.stat().st_size == member.size
                        ):
                            already_present += 1
                        else:
                            source = archive.extractfile(member)
                            if source is None:
                                raise FetchError(f"cannot read member {name!r} of archive {label}")
                            _write_member_atomically(source, target, member.size, name, label)
                            written_this_run.add(name)
                            extracted += 1
                        result.setdefault(name, target)
                    # Streaming mode would otherwise keep every TarInfo alive.
                    _forget_members(archive)
        except _STREAM_ERRORS as exc:
            raise FetchError(
                f"reading archive {label} failed after {scanned} members "
                f"({extracted} extracted, {already_present} already present): {exc!r}; "
                "re-run to continue (archives cannot be resumed mid-stream)"
            ) from exc

    logger.info(
        "stream_extract %s: scanned %d members, extracted %d, already present %d, "
        "ignored %d selected non-regular members (directories/links/devices)",
        label,
        scanned,
        extracted,
        already_present,
        ignored_non_regular,
    )
    return list(result.values())


def _forget_members(archive: tarfile.TarFile) -> None:
    """Drop ``TarFile``'s per-member cache; streaming a huge archive would otherwise grow it."""
    # ``members`` is a runtime attribute that typeshed does not declare.
    archive.members = []  # type: ignore[attr-defined]


@contextmanager
def _open_source(
    url_or_path: str | Path, opener: Opener, timeout_s: float
) -> Iterator[Response | IO[bytes]]:
    if isinstance(url_or_path, str) and url_or_path.lower().startswith(("http://", "https://")):
        request = urllib.request.Request(url_or_path, headers={"User-Agent": _USER_AGENT})
        try:
            response = opener(request, timeout_s)
        except _NETWORK_ERRORS as exc:  # includes HTTPError
            raise FetchError(f"opening {url_or_path} failed: {exc!r}") from exc
        with response:
            if response.status != 200:
                raise FetchError(f"unexpected HTTP status {response.status} from {url_or_path}")
            yield response
    else:
        path = Path(url_or_path)
        if not path.is_file():
            raise FetchError(f"archive {path} does not exist or is not a file")
        with path.open("rb") as handle:
            yield handle


def _collapse_name(name: str) -> str:
    """Drop leading ``./`` and empty/``.`` components; performs no safety checks."""
    return "/".join(part for part in name.split("/") if part not in ("", "."))


def _safe_member_name(raw_name: str) -> str:
    """Normalise a regular-file member name or raise ``FetchError`` if it is unsafe."""
    reason: str | None = None
    if "\x00" in raw_name:
        reason = "contains a NUL byte"
    elif "\\" in raw_name:
        reason = "contains a backslash"
    elif raw_name.startswith("/"):
        reason = "is an absolute path"
    elif _DRIVE_LETTER_RE.match(raw_name) is not None:
        reason = "starts with a drive letter"
    elif ".." in raw_name.split("/"):
        reason = "contains a '..' component"
    else:
        name = _collapse_name(raw_name)
        if name:
            return name
        reason = "is empty"
    raise FetchError(f"unsafe tar member name {raw_name!r}: {reason}")


def _resolve_target(root: Path, dest_dir: Path, name: str) -> Path:
    """Return ``dest_dir / name`` after checking that it stays inside ``dest_dir``."""
    target = dest_dir.joinpath(*name.split("/"))
    try:
        target.resolve().relative_to(root)
    except ValueError:
        raise FetchError(f"tar member {name!r} would be written outside {dest_dir}") from None
    return target


def _write_member_atomically(
    source: IO[bytes], target: Path, expected_bytes: int, name: str, label: str
) -> None:
    """Copy ``source`` to ``target`` via a temp file in the same directory and a rename."""
    target.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp_name = tempfile.mkstemp(prefix=f".{target.name}.", suffix=".tmp", dir=target.parent)
    tmp_path = Path(tmp_name)
    try:
        with os.fdopen(fd, "wb") as handle:
            shutil.copyfileobj(source, handle, _COPY_CHUNK_BYTES)
            written = handle.tell()
            if written != expected_bytes:
                raise FetchError(
                    f"member {name!r} of {label}: wrote {written} bytes, expected {expected_bytes}"
                )
            handle.flush()
            os.fsync(handle.fileno())
        tmp_path.replace(target)
    except BaseException:
        tmp_path.unlink(missing_ok=True)
        raise


# --------------------------------------------------------------------------------------
# Hugging Face URLs
# --------------------------------------------------------------------------------------


def hf_resolve_url(
    repo_id: str,
    revision: str,
    path_in_repo: str,
    repo_type: Literal["model", "dataset"],
) -> str:
    """Build a revision-pinned Hugging Face ``resolve`` URL.

    Args:
        repo_id: Repository id such as ``"google/fleurs"``.
        revision: Full 40-character lowercase hex commit SHA (branches and tags are refused
            so the downloaded data is pinned).
        path_in_repo: File path inside the repository (``/``-separated, relative).
        repo_type: ``"dataset"`` or ``"model"``.

    Returns:
        ``https://huggingface.co/[datasets/]<repo_id>/resolve/<revision>/<path>`` with every
        path component percent-encoded (``/`` kept).

    Raises:
        ValueError: On a malformed revision, repo id or path, or an unknown repo type.
    """
    if repo_type not in ("model", "dataset"):
        raise ValueError(f"repo_type must be 'model' or 'dataset', got {repo_type!r}")
    if _COMMIT_SHA_RE.fullmatch(revision) is None:
        raise ValueError(f"revision must be a 40-character lowercase hex commit SHA: {revision!r}")
    _require_clean_path("repo_id", repo_id)
    _require_clean_path("path_in_repo", path_in_repo)
    prefix = "datasets/" if repo_type == "dataset" else ""
    quoted_repo = urllib.parse.quote(repo_id, safe="/")
    quoted_path = urllib.parse.quote(path_in_repo, safe="/")
    return f"https://huggingface.co/{prefix}{quoted_repo}/resolve/{revision}/{quoted_path}"


def _require_clean_path(what: str, value: str) -> None:
    parts = value.split("/")
    if not value or any(part in ("", ".", "..") for part in parts) or "\\" in value:
        raise ValueError(
            f"{what} must be a relative '/'-separated path without empty, '.' or '..' "
            f"components: {value!r}"
        )