#!/usr/bin/env python3
"""
pfxor - windowed-palindrome prime field XOR cipher.

Keystream construction (validated 2026-08-04):
  - Field of midpoint M: radii d with M-d and M+d both prime.
  - Take a window of W consecutive gaps starting at shell distance D.
  - Bit i = sign(g_i - g_{W-1-i}), for i in the first half only.
  - Ties (g_i == g_mirror) are dropped.

Properties:
  eps = 0 by exchangeability (distribution-free, no debiasing needed)
  sieve-proof: bits are comparisons, expose no residues mod p
  random access: key is (M, D); window may sit at any depth in the flesh

HARD RULES:
  1. Never reuse a radius range. Not "a different window" - a DISJOINT one.
     Windows sharing a centre reuse bits verbatim; windows merely sharing
     gaps correlate at rho = 1/3. Both are fatal.
  2. Only the first half of the palindrome is usable. The second half is
     its bitwise complement reversed.
  3. Use M >= 10^30. Below 10^15 brute force over M is reachable.
"""
import argparse, hashlib, json, os, sys, time
import numpy as np
from gmpy2 import mpz, is_prime
from sympy import primerange

SIEVE_LIMIT = 300_000          # legacy default

def sieve_limit_for(span):
    """Presieve cost is O(pi(limit)) per segment; benefit scales with span.
    Measured optimum tracks span/250 (see bench)."""
    return max(1000, min(2_000_000, int(span // 250)))
_SMALL = None

def small_primes():
    global _SMALL
    if _SMALL is None:
        _SMALL = list(primerange(3, SIEVE_LIMIT))
    return _SMALL


def enumerate_field(M, d_start, n_hits, seg=80_000_000, progress=False):
    """Radii d >= d_start with M-d and M+d both prime. Returns the first n_hits."""
    M = mpz(M)
    sp = list(primerange(3, sieve_limit_for(seg)))
    parity = 1 if M % 2 == 0 else 0        # M-d must be odd
    base = d_start + ((parity - d_start) % 2)
    hits = []
    t0 = time.time()
    while len(hits) < n_hits:
        n = seg // 2
        alive = np.ones(n, dtype=bool)
        for p in sp:
            r = int(M % p)
            for t in (r, p - r):           # d == +-M mod p can never be a hit
                s = (t - base) % (2 * p)
                if s % 2:
                    s = (s + p) % (2 * p)
                start = s // 2
                if start < n:
                    alive[start::p] = False
        for off in np.nonzero(alive)[0]:
            d = base + 2 * int(off)
            if is_prime(M - d) and is_prime(M + d):
                hits.append(d)
                if len(hits) >= n_hits:
                    break
        if progress:
            el = time.time() - t0
            print(f"  ... {len(hits):,}/{n_hits:,} hits, d up to {base+seg:,}, "
                  f"{el:.0f}s", file=sys.stderr)
        base += seg
    return np.array(hits, dtype=object)


def window_size(nbits):
    """Deterministic from the payload alone, so decryption reproduces it.
    Tie rate is ~0.2-0.3%; 3% headroom is generous."""
    return 2 * int(np.ceil(nbits * 1.03)) + 2


def keystream(M, D, nbits, progress=False):
    W = window_size(nbits)
    hits = enumerate_field(M, D, W + 1, progress=progress)
    g = np.array([int(x) for x in np.diff(hits)], dtype=np.int64)[:W]
    d = g - g[::-1]
    half = d[:W // 2]
    keep = half != 0
    bits = (half[keep] > 0).astype(np.uint8)
    ties = int((~keep).sum())
    if len(bits) < nbits:
        raise RuntimeError(
            f"tie rate unexpectedly high: {len(bits)} bits from W={W}, need {nbits}")
    return bits[:nbits], {"W": W, "ties": ties,
                          "tie_rate": ties / len(half),
                          "d_first": int(hits[0]), "d_last": int(hits[-1]),
                          "ones": float(bits[:nbits].mean())}


def xor_file(data: bytes, bits: np.ndarray) -> bytes:
    ks = np.packbits(bits)
    return bytes(np.frombuffer(data, dtype=np.uint8) ^ ks[:len(data)])


def run(args):
    data = open(args.infile, "rb").read()
    nbits = len(data) * 8
    M, D = int(args.midpoint), int(args.radius)

    if M < 10**30:
        print(f"WARNING: M has {len(str(M))} digits. Below 10^30 the midpoint "
              f"is reachable by brute force.", file=sys.stderr)
    if M % 2 == 0:
        print("NOTE: even M - field radii will be odd.", file=sys.stderr)

    print(f"payload   {len(data):,} bytes ({nbits:,} bits)", file=sys.stderr)
    print(f"midpoint  {len(str(M))} digits", file=sys.stderr)
    print(f"radius    d = {D:,}", file=sys.stderr)
    print(f"window    W = {window_size(nbits):,} gaps -> "
          f"{window_size(nbits)+1:,} field objects", file=sys.stderr)
    t0 = time.time()
    bits, meta = keystream(M, D, nbits, progress=args.progress)
    el = time.time() - t0

    out = xor_file(data, bits)
    open(args.outfile, "wb").write(out)

    meta.update({
        "midpoint_digits": len(str(M)), "radius_start": D,
        "bytes": len(data), "seconds": round(el, 1),
        "radius_span": [meta["d_first"], meta["d_last"]],
        "sha256_in": hashlib.sha256(data).hexdigest()[:16],
        "sha256_out": hashlib.sha256(out).hexdigest()[:16],
    })
    print(f"\nenumerated to d = {meta['d_last']:,} in {el:.1f}s "
          f"({len(data)/max(el,1e-9):.0f} B/s)", file=sys.stderr)
    print(f"ties dropped {meta['ties']:,} ({100*meta['tie_rate']:.3f}%)", file=sys.stderr)
    print(f"keystream ones = {meta['ones']:.5f} "
          f"(z = {(meta['ones']-0.5)*2*np.sqrt(nbits):+.2f})", file=sys.stderr)
    print(f"\nwrote {args.outfile}", file=sys.stderr)
    print(f"\nRADIUS RANGE USED: [{meta['d_first']:,} .. {meta['d_last']:,}]", file=sys.stderr)
    print(f"Do not reuse this range, or any overlapping range, with this "
          f"midpoint.", file=sys.stderr)
    if args.meta:
        json.dump(meta, open(args.meta, "w"), indent=2)


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("infile")
    ap.add_argument("outfile")
    ap.add_argument("-m", "--midpoint", required=True,
                    help="midpoint M (decimal, or 10**30+61 style expression)")
    ap.add_argument("-r", "--radius", required=True, help="start shell distance d")
    ap.add_argument("--meta", help="write run metadata to this JSON file")
    ap.add_argument("--progress", action="store_true")
    a = ap.parse_args()
    a.midpoint = eval(a.midpoint) if "*" in str(a.midpoint) or "+" in str(a.midpoint) else a.midpoint
    a.radius = eval(a.radius) if "*" in str(a.radius) or "+" in str(a.radius) else a.radius
    run(a)
