#!/usr/bin/env python3
"""Convert a PrimeFiles 23# coprime-reduction segment to prime text."""

from __future__ import annotations

import argparse
import array
import re
import zlib
from pathlib import Path
from typing import Iterator

BLOCK_WIDTH = 223_092_870
BITMAP_WIDTH = BLOCK_WIDTH * 2
BYTES_PER_BLOCK = 4_561_920
BITS_PER_BLOCK = BYTES_PER_BLOCK * 8
BASE_WHEEL = 9_699_690  # 19#
BASE_RESIDUE_COUNT = 1_658_880
IMPLICIT_PRIMES = (2, 3, 5, 7, 11, 13, 17, 19, 23)
MAX_RAW_BYTES = 9_123_840
MAX_INPUT_BYTES = 17_000_000
PUBLIC_NAME = re.compile(
    r"primefiles-(?P<number>\d+)-(?P<start>\d+)-(?P<end>\d+)-coprime23-v4\.bin(?:\.gz)?$"
)


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Decode a PrimeFiles 23# coprime-reduction bitmap into text."
    )
    parser.add_argument("input", type=Path, help="23# coprime-reduction .bin or .bin.gz segment")
    parser.add_argument("output", type=Path, nargs="?", help="output text file (default: input name with .txt)")
    parser.add_argument("--start", type=int, help="bitmap start; inferred from an original PrimeFiles filename")
    parser.add_argument("--end", type=int, help="exclusive logical range end; defaults to represented range")
    parser.add_argument(
        "--include-implicit-small-primes",
        action="store_true",
        help="write 2, 3, 5, 7, 11, 13, 17, 19 and 23 when in the requested range",
    )
    parser.add_argument("--format-version", type=int, choices=(4,), help="format version; defaults to 4")
    return parser.parse_args()


def default_output_path(source: Path) -> Path:
    name = source.name
    if name.endswith(".gz"):
        name = name[:-3]
    if name.endswith(".bin"):
        name = name[:-4]
    return source.with_name(name + ".txt")


def infer_range(source: Path) -> tuple[int, int, int]:
    match = PUBLIC_NAME.fullmatch(source.name)
    if match is None:
        raise ValueError("cannot infer the range from this filename; provide --start")
    segment_number = int(match.group("number"))
    logical_start = int(match.group("start"))
    logical_end = int(match.group("end"))
    bitmap_start = 0 if segment_number == 0 and logical_start == 24 else logical_start
    return bitmap_start, logical_start, logical_end


def read_payload(source: Path) -> bytes:
    if source.stat().st_size > MAX_INPUT_BYTES:
        raise ValueError(f"input exceeds the {MAX_INPUT_BYTES:,}-byte safety limit")
    compressed = source.read_bytes()
    if not compressed.startswith(b"\x1f\x8b"):
        if len(compressed) > MAX_RAW_BYTES:
            raise ValueError("uncompressed bitmap exceeds its format size")
        return compressed
    decoder = zlib.decompressobj(16 + zlib.MAX_WBITS)
    payload = decoder.decompress(compressed, MAX_RAW_BYTES + 1)
    if len(payload) > MAX_RAW_BYTES or decoder.unconsumed_tail:
        raise ValueError("gzip payload expands beyond the bitmap safety limit")
    if not decoder.eof:
        raise ValueError("truncated gzip stream")
    if decoder.unused_data:
        raise ValueError("gzip stream contains trailing or concatenated data")
    return payload


def iter_coprime_residues() -> Iterator[int]:
    """Yield the sorted residues of 23# using a small, bounded 19# table."""
    base_residues = array.array("I")
    small_primes = (2, 3, 5, 7, 11, 13, 17, 19)
    for value in range(1, BASE_WHEEL):
        if all(value % prime for prime in small_primes):
            base_residues.append(value)
    if len(base_residues) != BASE_RESIDUE_COUNT:
        raise ValueError("invalid 19# residue table")
    for _ in range(2):
        for block in range(23):
            offset = block * BASE_WHEEL
            for residue in base_residues:
                value = offset + residue
                if value % 23:
                    yield value


def convert(
    payload: bytes,
    destination: Path,
    bitmap_start: int,
    logical_start: int,
    logical_end: int,
    include_small_primes: bool,
) -> int:
    count = 0
    residues = iter_coprime_residues()
    with destination.open("w", encoding="ascii", newline="\n") as output:
        if include_small_primes:
            for prime in IMPLICIT_PRIMES:
                # The public filename starts at 24 for segment zero, but the
                # optional implicit primes belong to its physical [0, end) range.
                if bitmap_start <= prime < logical_end:
                    output.write(f"{prime}\n")
                    count += 1

        candidate_index = 0
        for byte in payload:
            for bit in range(8):
                residue = next(residues)
                if byte & (1 << bit):
                    value = bitmap_start + candidate_index // BITS_PER_BLOCK * BLOCK_WIDTH + residue
                    if logical_start <= value < logical_end:
                        output.write(f"{value}\n")
                        count += 1
                candidate_index += 1
    return count


def main() -> None:
    args = parse_arguments()
    if args.format_version not in (None, 4):
        raise SystemExit("--format-version must be 4")
    inferred_logical_start = None
    inferred_end = None
    if args.start is None:
        try:
            bitmap_start, inferred_logical_start, inferred_end = infer_range(args.input)
        except ValueError as error:
            raise SystemExit(str(error)) from error
    else:
        bitmap_start = args.start
    if bitmap_start < 0 or bitmap_start % BLOCK_WIDTH:
        raise SystemExit(f"--start must be a non-negative multiple of {BLOCK_WIDTH:,}")

    try:
        payload = read_payload(args.input)
    except (OSError, ValueError, zlib.error) as error:
        raise SystemExit(f"invalid PrimeFiles bitmap: {error}") from error
    if len(payload) != MAX_RAW_BYTES:
        raise SystemExit(f"invalid 23# coprime-reduction bitmap length: {len(payload):,} bytes")

    represented_end = bitmap_start + BITMAP_WIDTH
    logical_start = inferred_logical_start if inferred_logical_start is not None else bitmap_start
    logical_end = args.end if args.end is not None else (inferred_end or represented_end)
    if logical_end <= logical_start or logical_end > represented_end:
        raise SystemExit(f"--end must be above {logical_start:,} and at most {represented_end:,}")

    destination = args.output or default_output_path(args.input)
    count = convert(payload, destination, bitmap_start, logical_start, logical_end, args.include_implicit_small_primes)
    print(f"Wrote {count:,} primes to {destination} (PrimeFiles v4, 23# coprime reduction)")


if __name__ == "__main__":
    main()