from pathlib import Path
import zipfile


ROOT = Path(__file__).resolve().parent
OUTPUT = ROOT / "lite_runtime.zip"
INPUTS = (ROOT / "eval_exp.py", ROOT / "chinatravel", ROOT / "leaderboard_data")


def include(path: Path) -> bool:
    relative = path.relative_to(ROOT)
    if "__pycache__" in relative.parts or path.suffix in {".pyc", ".pyo"}:
        return False
    return not (path.name == "human1000.jsonl" and "oracle" in relative.parts)


def archive_files():
    for source in INPUTS:
        if source.is_file():
            yield source
            continue
        yield from (path for path in sorted(source.rglob("*")) if path.is_file())


temporary = OUTPUT.with_suffix(".zip.tmp")
with zipfile.ZipFile(
    temporary,
    mode="w",
    compression=zipfile.ZIP_DEFLATED,
    compresslevel=9,
) as archive:
    for path in archive_files():
        if include(path):
            archive.write(path, path.relative_to(ROOT).as_posix())

with zipfile.ZipFile(temporary) as archive:
    names = archive.namelist()
    private_oracles = [name for name in names if name.endswith("/human1000.jsonl")]
    if private_oracles:
        raise RuntimeError(f"Private Oracle files entered the Lite archive: {private_oracles}")
    if "chinatravel/environment/database/intercity_transport/train/from_北京_to_上海.json" not in names:
        raise RuntimeError("UTF-8 database filenames were not preserved in the Lite archive.")

temporary.replace(OUTPUT)
print(f"Wrote {OUTPUT} with {len(names)} files ({OUTPUT.stat().st_size} bytes).")