#!/usr/bin/env python3
"""Copy competition notebook sources out of the Meta Kaggle Code zip.

Files in that dataset are named by KernelVersions.Id. Only versions referenced
by the local kernel-version map are extracted.
"""

from __future__ import annotations

import shutil
import zipfile
from pathlib import Path

from download_competitions import CODE_SOURCE_DIR, META_DIR, ensure_kernel_versions


def main() -> None:
    versions = ensure_kernel_versions()
    needed = set(versions.values())
    print(f"need {len(needed)} version files", flush=True)
    root = META_DIR / "meta-kaggle-code"
    zips = sorted(root.glob("*.zip"), key=lambda path: path.stat().st_size)
    if not zips:
        raise SystemExit(f"no zip in {root}")
    zip_path = zips[-1]
    print(f"zip {zip_path} bytes {zip_path.stat().st_size}", flush=True)
    CODE_SOURCE_DIR.mkdir(parents=True, exist_ok=True)
    matched = 0
    shown = 0
    with zipfile.ZipFile(zip_path) as archive:
        for info in archive.infolist():
            if info.is_dir():
                continue
            name = Path(info.filename)
            if shown < 8:
                print(f"zip entry {info.filename}", flush=True)
                shown += 1
            if name.stem not in needed:
                continue
            suffix = name.suffix or ".src"
            dest = CODE_SOURCE_DIR / f"{name.stem}{suffix}"
            if not (dest.exists() and dest.stat().st_size > 0):
                with archive.open(info, "r") as src, dest.open("wb") as out:
                    shutil.copyfileobj(src, out, length=1024 * 1024)
            matched += 1
            if matched % 5000 == 0:
                print(f"extracted {matched}", flush=True)
    print(f"extracted {matched} of {len(needed)}", flush=True)
    if len(needed) > 10000 and matched < 1000:
        raise SystemExit(2)


if __name__ == "__main__":
    main()