"""QEC decoder accuracy board (SPEC_v7_decoder s3-s6, review_v7_fable D6-D9). Standard library only.

An artifact is the predicted logical-observable flips for every shot of one instance, as canonical base64 of Stim b8
rows (ceil(K/8) bytes per shot, bit j of a row = (row[j // 8] >> (j % 8)) & 1, padding bits zero):

  {"format": "b8-base64-v1", "instance": "<id>", "predictions": "<base64>",
   "estimate": {"batch_ns": int, "cores": int, "hardware": "printable ASCII"}}      (estimate optional)

Score: errors E = number of shots whose predicted row differs from the true row in any bit (a shot with several wrong
observables counts once), with shots N, rounds R, ler_shot = E/N and ler_round = E/(N R) as reduced fractions.

The true flips of ranked instances are private: they live in $QARENA_PRIVATE/decoder/<id>/key.b8 next to
private-manifest.json, which binds the key's length and SHA-256 to the SHA-256 of the public instance file. Every format
check finishes before the key is opened, so a rejection never depends on a label; key loading and scoring report only
"hidden key unavailable" or "scorer unavailable" on failure (D9), never a path, a byte or a partial count. Development
instances (-dev) carry their public labels in the instance file and are scored completely by the public bundle.
"""
import base64
import binascii
import hashlib
import json
import math
import os
import re
import time

from common import INSTANCES, Reject, exact_keys, short, want_int

FORMAT = "b8-base64-v1"
MAX_ARTIFACT_BYTES = 1048576
BLOCK_ROWS = 4096                      # process_time() is polled between blocks of at most this many rows
B64 = re.compile(r"[A-Za-z0-9+/]*={0,2}")
INST_ID = re.compile(r"[a-z0-9][a-z0-9-]{0,63}")
PRINTABLE = re.compile(r"[\x20-\x7e]{1,240}")


class KeyUnavailable(Exception):
    """No private key for a ranked instance on this machine (for example in the public bundle)."""


def frac(a, b):
    g = math.gcd(a, b)
    return {"num": format(a // g, "x"), "den": format(b // g, "x")}


def load_instance_bytes(inst):
    if not INST_ID.fullmatch(inst):
        raise ValueError("bad instance id")
    with open(os.path.join(INSTANCES, "decoder", inst + ".json"), "rb") as f:
        raw = f.read()
    return json.loads(raw.decode("utf-8")), raw


def geometry(I):
    N, K, R = I["shots"], I["observables"], I["rounds"]
    B = (K + 7) // 8
    if sum(e["shots"] for e in I["experiments"]) != N or I["bytes_per_shot"] != B or I["format"] != FORMAT:
        raise ValueError("inconsistent instance file")
    return N, K, R, B


def padding_ok(buf, B, K):
    """True iff every row's padding bits (bits K .. 8B-1) are zero."""
    r = K % 8
    if r == 0:
        return True
    return not buf[B - 1::B].translate(None, bytes(range(1 << r)))


def check_estimate(e):
    if not isinstance(e, dict):
        raise Reject("estimate must be an object with keys batch_ns, cores, hardware")
    exact_keys(e, ("batch_ns", "cores", "hardware"), where="estimate")
    want_int(e["batch_ns"], 1, 2 ** 63 - 1, "estimate.batch_ns")
    want_int(e["cores"], 1, 4096, "estimate.cores")
    h = e["hardware"]
    if not isinstance(h, str) or not PRINTABLE.fullmatch(h):
        raise Reject("estimate.hardware must be 1 to 240 printable ASCII characters")
    return {"batch_ns": e["batch_ns"], "cores": e["cores"], "hardware": h}


def decode_predictions(s, N, B, K):
    """Canonical standard base64 of exactly N*B bytes with zero padding bits in every row."""
    if not isinstance(s, str):
        raise Reject("predictions must be a base64 string")
    want = 4 * ((N * B + 2) // 3)
    if len(s) != want:
        raise Reject(f"predictions must be exactly {want} base64 characters for {N} shots of {B} byte(s), got {len(s)}")
    if not B64.fullmatch(s):
        raise Reject("predictions: only A-Z a-z 0-9 + / and final = padding are allowed (no whitespace, no URL-safe alphabet)")
    try:
        buf = base64.b64decode(s, validate=True)
    except (binascii.Error, ValueError):
        raise Reject("predictions are not valid base64")
    if len(buf) != N * B:
        raise Reject(f"predictions decode to {len(buf)} bytes, expected {N * B}")
    if base64.b64encode(buf).decode("ascii") != s:
        raise Reject("predictions are not canonical base64 (nonzero unused bits in the last character)")
    if not padding_ok(buf, B, K):
        raise Reject(f"predictions: padding bits above observable {K - 1} must be zero in every row")
    return buf


def check_format(inst, I, art):
    """Every artifact check that does not need labels. Returns (prediction bytes, validated estimate or None)."""
    exact_keys(art, ("format", "instance", "predictions"), ("estimate",))
    if art["format"] != FORMAT:
        raise Reject(f"format must be {FORMAT!r}, got {short(art['format'])}")
    if art["instance"] != inst:
        raise Reject(f"instance must be {inst!r} (the instance you submit to), got {short(art['instance'])}")
    est = check_estimate(art["estimate"]) if "estimate" in art else None
    N, K, R, B = geometry(I)
    return decode_predictions(art["predictions"], N, B, K), est


def load_key(inst, I, inst_bytes):
    """The true observable flips, after every integrity check. Ranked keys come from the private directory only,
    at a path built from the registered instance id."""
    N, K, R, B = geometry(I)
    if I["mode"] == "debug":
        key = base64.b64decode(I["labels_b64"], validate=True)
        if hashlib.sha256(key).hexdigest() != I["labels_sha256"]:
            raise ValueError("development labels do not match their digest")
    else:
        root = os.environ.get("QARENA_PRIVATE", "")
        d = os.path.join(root, "decoder", inst)
        if not root or not os.path.isfile(os.path.join(d, "key.b8")) or not os.path.isfile(os.path.join(d, "private-manifest.json")):
            raise KeyUnavailable()
        with open(os.path.join(d, "private-manifest.json"), "rb") as f:
            pm = json.loads(f.read().decode("utf-8"))
        if (pm.get("instance") != inst or pm.get("board") != "decoder" or pm.get("key_file") != "key.b8"
                or pm.get("public_manifest_sha256") != hashlib.sha256(inst_bytes).hexdigest()
                or pm.get("key_bytes") != N * B or pm.get("shots") != N or pm.get("observables") != K):
            raise ValueError("private manifest does not match the public instance")
        with open(os.path.join(d, "key.b8"), "rb") as f:
            key = f.read(N * B + 1)
        if hashlib.sha256(key).hexdigest() != pm.get("key_sha256"):
            raise ValueError("key does not match its digest")
    if len(key) != N * B or not padding_ok(key, B, K):
        raise ValueError("key has the wrong length or nonzero padding")
    return key


def count_errors(pred, key, N, B, deadline=None):
    """Number of rows (B bytes each) where pred and key differ. Returns None if process_time() passes the deadline
    while rows remain. Rows are handled in blocks of BLOCK_ROWS: XOR as big integers, OR of the B byte planes, then
    count the nonzero rows."""
    E = 0
    for r0 in range(0, N, BLOCK_ROWS):
        if deadline is not None and time.process_time() > deadline:
            return None
        r1 = min(N, r0 + BLOCK_ROWS)
        n = (r1 - r0) * B
        x = (int.from_bytes(pred[r0 * B:r1 * B], "little") ^ int.from_bytes(key[r0 * B:r1 * B], "little")).to_bytes(n, "little")
        if B == 1:
            rows = x
        else:
            acc = 0
            for j in range(B):
                acc |= int.from_bytes(x[j::B], "little")
            rows = acc.to_bytes(r1 - r0, "little")
        E += (r1 - r0) - rows.count(0)
    return E


def cpu_deadline():
    """Soft CPU deadline: 80 % of the board's budget (run_one sets RLIMIT_CPU to the full budget)."""
    try:
        cpu = float(os.environ.get("QARENA_DECODER_CPU") or os.environ.get("QARENA_CPU") or 0)
    except ValueError:
        cpu = 0
    return time.process_time() + 0.8 * cpu if cpu > 0 else None


def verify(inst, art):
    I, raw = load_instance_bytes(inst)
    N, K, R, B = geometry(I)
    pred, est = check_format(inst, I, art)          # Reject propagates: no label has been read yet
    deadline = cpu_deadline()
    try:
        key = load_key(inst, I, raw)
        E = count_errors(pred, key, N, B, deadline)
    except KeyUnavailable:
        return {"status": "error", "detail": {"reason": "hidden key unavailable",
                                              "note": "ranked labels are private to the QArena server; the format is valid"}}
    except BaseException:                            # D9: no traceback, path or byte from key loading or scoring
        return {"status": "error", "detail": {"reason": "scorer unavailable"}}
    if E is None:
        return {"status": "error", "detail": {"reason": "verifier CPU budget exhausted (not ranked)"}}
    score = {"errors": E, "shots": N, "rounds": R, "ler_shot": frac(E, N), "ler_round": frac(E, N * R)}
    detail = {"scored_shots": N}
    if est is not None:
        detail["estimate"] = est
    return {"status": "accepted", "score": score, "detail": detail}