src/data/io_utils.py

from __future__ import annotations

import hashlib
import json
import tarfile
import time
import zipfile
from pathlib import Path
from typing import Any, Dict, Optional
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

import yaml


def load_yaml(path: Path) -> Dict[str, Any]:
    with path.open("r", encoding="utf-8") as handle:
        data = yaml.safe_load(handle)
    return data or {}


def ensure_dir(path: Path) -> None:
    path.mkdir(parents=True, exist_ok=True)


def sha256_file(path: Path, chunk_size: int = 1024 * 1024) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while True:
            block = handle.read(chunk_size)
            if not block:
                break
            digest.update(block)
    return digest.hexdigest()


def write_json(path: Path, payload: Dict[str, Any]) -> None:
    ensure_dir(path.parent)
    with path.open("w", encoding="utf-8") as handle:
        json.dump(payload, handle, indent=2, ensure_ascii=True)


def read_json(path: Path) -> Any:
    with path.open("r", encoding="utf-8") as handle:
        return json.load(handle)


def _download_once(
    url: str,
    destination: Path,
    timeout_seconds: int,
    chunk_size: int,
) -> None:
    temp_file = destination.with_suffix(destination.suffix + ".part")
    resume_pos = temp_file.stat().st_size if temp_file.exists() else 0

    headers = {}
    if resume_pos > 0:
        headers["Range"] = f"bytes={resume_pos}-"

    request = Request(url=url, headers=headers)
    with urlopen(request, timeout=timeout_seconds) as response:
        status = getattr(response, "status", 200)
        supports_resume = status == 206 and resume_pos > 0
        mode = "ab" if supports_resume else "wb"
        with temp_file.open(mode) as out:
            while True:
                chunk = response.read(chunk_size)
                if not chunk:
                    break
                out.write(chunk)
    temp_file.replace(destination)


def download_file(
    url: str,
    destination: Path,
    timeout_seconds: int = 120,
    retries: int = 3,
    chunk_size: int = 1024 * 1024,
) -> None:
    ensure_dir(destination.parent)
    attempts = max(1, retries)
    for attempt in range(1, attempts + 1):
        try:
            _download_once(
                url=url,
                destination=destination,
                timeout_seconds=timeout_seconds,
                chunk_size=chunk_size,
            )
            return
        except (HTTPError, URLError, TimeoutError, OSError) as exc:
            if attempt == attempts:
                raise RuntimeError(
                    f"Download failed after {attempts} attempts for {url}: {exc}"
                ) from exc
            sleep_s = min(10, attempt * 2)
            time.sleep(sleep_s)


def verify_checksum(path: Path, expected_sha256: Optional[str]) -> bool:
    if not expected_sha256:
        return True
    actual = sha256_file(path)
    return actual.lower() == expected_sha256.lower()


def _is_within_directory(base_dir: Path, target_path: Path) -> bool:
    try:
        base = base_dir.resolve(strict=False)
        target = target_path.resolve(strict=False)
        target.relative_to(base)
        return True
    except ValueError:
        return False


def extract_archive(archive_path: Path, output_dir: Path, safe_extract: bool = True) -> None:
    ensure_dir(output_dir)
    lower_name = archive_path.name.lower()

    if zipfile.is_zipfile(archive_path):
        with zipfile.ZipFile(archive_path, "r") as zf:
            members = zf.infolist()
            if safe_extract:
                for member in members:
                    target = output_dir / member.filename
                    if not _is_within_directory(output_dir, target):
                        raise RuntimeError(
                            f"Unsafe zip member path detected: {member.filename}"
                        )
            zf.extractall(output_dir)
        return

    if tarfile.is_tarfile(archive_path):
        with tarfile.open(archive_path, "r:*") as tf:
            members = tf.getmembers()
            if safe_extract:
                for member in members:
                    target = output_dir / member.name
                    if not _is_within_directory(output_dir, target):
                        raise RuntimeError(
                            f"Unsafe tar member path detected: {member.name}"
                        )
            tf.extractall(output_dir)
        return

    raise RuntimeError(f"Unsupported archive format: {lower_name}")