"""One-time download of the TinyStoriesV2 corpus.

This is the only network access in the entire project. Files land in
data/raw/ (gitignored). SHA-256 hashes are recorded in data/raw/SHA256SUMS
so the corpus can be verified later without re-downloading.

Usage: python scripts/download_data.py
"""

import hashlib
import urllib.request
from pathlib import Path

BASE_URL = "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main"
FILES = [
    "TinyStoriesV2-GPT4-train.txt",  # ~2.2 GB
    "TinyStoriesV2-GPT4-valid.txt",  # ~22 MB
]
RAW_DIR = Path(__file__).resolve().parent.parent / "data" / "raw"
CHUNK = 1 << 20  # 1 MB


def download(name: str) -> Path:
    dest = RAW_DIR / name
    if dest.exists():
        print(f"{name}: already present, skipping download")
        return dest
    tmp = dest.with_suffix(dest.suffix + ".part")
    url = f"{BASE_URL}/{name}"
    print(f"{name}: downloading")
    with urllib.request.urlopen(url) as response, open(tmp, "wb") as f:
        total = int(response.headers.get("Content-Length") or 0)
        done = 0
        while chunk := response.read(CHUNK):
            f.write(chunk)
            done += len(chunk)
            if total:
                print(
                    f"\r  {done / 1e6:,.0f} / {total / 1e6:,.0f} MB"
                    f" ({done * 100 // total}%)",
                    end="",
                    flush=True,
                )
        print()
    tmp.rename(dest)  # only a fully written file ever gets the real name
    return dest


def sha256(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        while chunk := f.read(CHUNK):
            h.update(chunk)
    return h.hexdigest()


def main() -> None:
    RAW_DIR.mkdir(parents=True, exist_ok=True)
    lines = []
    for name in FILES:
        path = download(name)
        digest = sha256(path)
        print(f"{name}: {path.stat().st_size / 1e6:,.0f} MB  sha256={digest}")
        lines.append(f"{digest}  {name}\n")
    sums = RAW_DIR / "SHA256SUMS"
    sums.write_text("".join(lines))
    print(f"hashes recorded in {sums}")


if __name__ == "__main__":
    main()