"""Shared helpers for QArena verifiers: strict JSON parsing and exact GF(2) linear algebra.

GF(2) vectors are Python ints used as bitsets (bit i = coordinate i). All routines are exact.
"""
import json
import os

HERE = os.path.dirname(os.path.abspath(__file__))
INSTANCES = os.path.join(os.path.dirname(HERE), "instances")


class Reject(Exception):
    """The artifact is invalid. The message is shown to the submitter."""


def _no_dupes(pairs):
    d = {}
    for k, v in pairs:
        if k in d:
            raise Reject(f"duplicate JSON key {k!r}")
        d[k] = v
    return d


def _bad_const(x):
    raise Reject(f"non-finite number {x} not allowed")


def _finite_float(s):
    v = float(s)
    if v != v or v in (float("inf"), float("-inf")):  # 1e999 parses to inf without passing through parse_constant
        raise Reject(f"non-finite number {s[:20]}")
    return v


MAX_CONTAINERS = 2_500_000   # the largest legal artifact (qrom, 10^6 gates) needs about 10^6


def load_artifact(path, max_bytes=20 * 1024 * 1024):
    if os.path.getsize(path) > max_bytes:
        raise Reject("artifact too large")
    with open(path, "rb") as f:
        raw = f.read()
    try:
        txt = raw.decode("utf-8")
    except UnicodeDecodeError:
        raise Reject("artifact is not UTF-8")
    if raw.count(b"[") + raw.count(b"{") > MAX_CONTAINERS:  # parsing deep junk costs ~50 bytes per input byte; macOS has no RLIMIT_AS
        raise Reject(f"artifact has more than {MAX_CONTAINERS:,} JSON arrays and objects")
    try:
        obj = json.loads(txt, object_pairs_hook=_no_dupes, parse_constant=_bad_const, parse_float=_finite_float)
    except Reject:
        raise
    except ValueError as e:
        raise Reject(f"invalid JSON: {e}")
    if not isinstance(obj, dict):
        raise Reject("artifact must be a JSON object")
    return obj


def load_instance(board, inst):
    with open(os.path.join(INSTANCES, board, inst + ".json"), encoding="utf-8") as f:
        return json.load(f)


def exact_keys(obj, required, optional=(), where="artifact"):
    keys = set(obj)
    missing = set(required) - keys
    extra = keys - set(required) - set(optional)
    if missing:
        raise Reject(f"{where}: missing keys {sorted(missing)}")
    if extra:
        raise Reject(f"{where}: unknown keys {sorted(extra)}")


def short(x, n=60):
    """repr of x, cut to n characters, so a rejection never echoes a whole artifact field."""
    r = repr(x)
    return r if len(r) <= n else r[:n] + "...(truncated)"


def is_int(x):
    return isinstance(x, int) and not isinstance(x, bool)


def want_int(x, lo, hi, what):
    if not is_int(x) or not lo <= x <= hi:
        raise Reject(f"{what} must be an integer in [{lo}, {hi}], got {short(x)}")
    return x


def want_list(x, what, max_len=None, min_len=0):
    if not isinstance(x, list):
        raise Reject(f"{what} must be a list")
    if len(x) < min_len or (max_len is not None and len(x) > max_len):
        raise Reject(f"{what} must have length in [{min_len}, {max_len}], got {len(x)}")
    return x


def bitstring(s, n, what):
    """'0110' (character i = coordinate i) -> int bitset. Exact length n, only '0'/'1'."""
    if not isinstance(s, str) or len(s) != n or s.strip("01") != "":
        raise Reject(f"{what} must be a string of exactly {n} characters '0'/'1'")
    v = 0
    for i, ch in enumerate(s):
        if ch == "1":
            v |= 1 << i
    return v


def popcount(v):
    return bin(v).count("1")


def rref(rows):
    """Reduced row echelon basis over GF(2): dict pivot_bit -> row (each pivot appears in exactly one row)."""
    basis = {}
    for r in rows:
        for p, b in basis.items():
            if r >> p & 1:
                r ^= b
        if r:
            p = r.bit_length() - 1
            for q in list(basis):
                if basis[q] >> p & 1:
                    basis[q] ^= r
            basis[p] = r
    return basis


def rank(rows):
    return len(rref(rows))


def reduce(v, basis):
    for p, b in basis.items():
        if v >> p & 1:
            v ^= b
    return v


def in_span(v, basis):
    return reduce(v, basis) == 0


def nullspace(rows, n):
    """Basis (list of int bitsets over n coordinates) of {x : r.x = 0 for every row r}."""
    basis = rref(rows)
    pivots = set(basis)
    free = [i for i in range(n) if i not in pivots]
    out = []
    for f in free:
        x = 1 << f
        for p, b in basis.items():
            if b >> f & 1:
                x |= 1 << p
        out.append(x)
    return out


def dot(a, b):
    return popcount(a & b) & 1