#!/usr/bin/env python3
"""Bounded native bzip2 block-order experiment. Authored code: CC0-1.0."""
import argparse
import bz2
import ctypes
import ctypes.util
import datetime
import hashlib
import json
import os
from pathlib import Path
import platform
import shutil
import subprocess

BLOCK_MAGIC = f"{0x314159265359:048b}"
END_MAGIC = f"{0x177245385090:048b}"
MASK = (1 << 32) - 1
PAYLOAD_BYTES = 3_250_000
SEED = b"hesela-bzip2-block-order-v1"
EXPECTED_IDS = ("original", "swap-1-2", "swap-1-32", "swap-1-33",
                "block-crc-bit", "stream-crc-bit")


def require(condition, message):
    if not condition:
        raise ValueError(message)


def digest(data):
    return hashlib.sha256(data).hexdigest()


def rol(value, count=1):
    count %= 32
    value &= MASK
    return ((value << count) | (value >> ((32 - count) % 32))) & MASK


def combined(crcs):
    value = 0
    for crc in crcs:
        value = rol(value) ^ crc
    return value


def closed_form(crcs):
    value = 0
    for index, crc in enumerate(crcs):
        value ^= rol(crc, len(crcs) - 1 - index)
    return value


def bits(data):
    return "".join(f"{byte:08b}" for byte in data)


def packed(value):
    require(set(value) <= {"0", "1"}, "non-bit input")
    value += "0" * (-len(value) % 8)
    return bytes(int(value[i:i + 8], 2) for i in range(0, len(value), 8))


def positions(value, marker):
    result, start = [], 0
    while True:
        position = value.find(marker, start)
        if position < 0:
            return result
        result.append(position)
        start = position + 1


def parse_fixture(data):
    # Strict scanner for our generated fixture, not a general-purpose decoder.
    value = bits(data)
    require(data[:4] == b"BZh1", "expected native -1 header")
    starts = positions(value, BLOCK_MAGIC)
    endings = positions(value, END_MAGIC)
    require(len(starts) == 33 and starts[0] == 32, "expected 33 unique blocks")
    require(len(endings) == 1 and starts[-1] < endings[0], "ambiguous end marker")
    end = endings[0]
    require(len(value) - end in range(80, 88), "unexpected trailer length")
    require(set(value[end + 80:]) <= {"0"}, "nonzero padding")
    blocks = [value[a:b] for a, b in zip(starts, starts[1:] + [end])]
    crcs = [int(block[48:80], 2) for block in blocks]
    stored = int(value[end + 48:end + 80], 2)
    require(combined(crcs) == stored == closed_form(crcs), "CRC formula disagreement")
    require(packed(value[:32] + "".join(blocks) + value[end:]) == data,
            "scanner did not reconstruct original bytes")
    return value[:32], blocks, value[end:], crcs, starts


def payload_blocks(header, blocks):
    return [bz2.decompress(packed(header + block + END_MAGIC + block[48:80]))
            for block in blocks]


def mutate(header, blocks, trailer, order=None, block_crc=False, stream_crc=False):
    order = list(range(len(blocks))) if order is None else order
    require(sorted(order) == list(range(len(blocks))), "not a permutation")
    selected = [blocks[index] for index in order]
    if block_crc:
        block = selected[0]
        selected[0] = block[:48] + str(1 - int(block[48])) + block[49:]
    if stream_crc:
        trailer = trailer[:48] + str(1 - int(trailer[48])) + trailer[49:]
    return packed(header + "".join(selected) + trailer)


def clean_environment():
    env = dict(os.environ)
    for name in ("BZIP", "BZIP2"):
        env.pop(name, None)
    env["LC_ALL"] = "C"
    return env


def command(binary, flags, data=b""):
    return subprocess.run([str(binary), *flags], input=data, capture_output=True,
                          timeout=90, env=clean_environment(), check=False)


def python_decode(data, expected):
    decoder = bz2.BZ2Decompressor()
    try:
        output = decoder.decompress(data)
    except OSError as error:
        return {"accepted": False, "error_type": type(error).__name__}
    return {
        "accepted": decoder.eof and not decoder.unused_data,
        "eof": decoder.eof, "unused_bytes": len(decoder.unused_data),
        "output_bytes": len(output), "output_sha256": digest(output),
        "matches_expected": output == expected,
    }


def run(binary):
    binary = Path(shutil.which(str(binary)) or binary).resolve(strict=True)
    expected = hashlib.shake_256(SEED).digest(PAYLOAD_BYTES)
    encoded = command(binary, ["-1c"], expected)
    require(encoded.returncode == 0 and not encoded.stderr, "native compression failed")
    original = encoded.stdout
    header, blocks, trailer, crcs, starts = parse_fixture(original)
    parts = payload_blocks(header, blocks)
    require(b"".join(parts) == expected, "independent input oracle mismatch")
    rows = []
    orders = {"original": list(range(33))}
    for distance in (1, 31, 32):
        order = list(range(33))
        order[0], order[distance] = order[distance], order[0]
        orders[f"swap-1-{distance + 1}"] = order
    for name in EXPECTED_IDS:
        order = orders.get(name, list(range(33)))
        fixture = mutate(header, blocks, trailer, order,
                         name == "block-crc-bit", name == "stream-crc-bit")
        tested = command(binary, ["-t"], fixture)
        decoded = command(binary, ["-dc"], fixture)
        small = command(binary, ["-dsc"], fixture)
        intended_order = b"".join(parts[i] for i in order)
        accepted = name in ("original", "swap-1-33")
        expected_code = 0 if accepted else 2
        require(tested.returncode == decoded.returncode == small.returncode == expected_code,
                f"unexpected native outcomes for {name}")
        if accepted:
            require(decoded.stdout == small.stdout == intended_order,
                    f"accepted output differs from block oracle for {name}")
            require(not tested.stderr and not decoded.stderr and not small.stderr,
                    f"unexpected warning for {name}")
        python = python_decode(fixture, expected)
        require(python["accepted"] == accepted, f"Python path disagrees for {name}")
        rows.append({
            "id": name, "compressed_bytes": len(fixture),
            "compressed_sha256": digest(fixture),
            "test_exit": tested.returncode, "decode_exit": decoded.returncode,
            "small_decode_exit": small.returncode,
            "decoded_bytes": len(decoded.stdout), "decoded_sha256": digest(decoded.stdout),
            "matches_original": decoded.stdout == expected,
            "accepted_exact_original": decoded.returncode == 0 and decoded.stdout == expected,
            "same_output_size": len(decoded.stdout) == len(expected),
            "stderr_test": tested.stderr.decode("utf-8", errors="replace"),
            "stderr_decode": decoded.stderr.decode("utf-8", errors="replace"),
            "python": python, "block_order_one_based": [i + 1 for i in order],
            "same_compressed_size": len(fixture) == len(original),
            "same_stored_stream_crc": name != "stream-crc-bit",
            "all_original_block_bits_retained": name != "block-crc-bit",
            "recomputed_combined_crc_hex": f"{combined([crcs[i] for i in order]):08x}",
        })
    lib = ctypes.CDLL(ctypes.util.find_library("bz2"))
    lib.BZ2_bzlibVersion.restype = ctypes.c_char_p
    return {
        "schema_version": 1, "kind": "designed-native-experiment",
        "generated_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        "environment": {
            "python": platform.python_version(), "platform": platform.platform(),
            "binary": str(binary), "binary_sha256": digest(binary.read_bytes()),
            "binary_version": command(binary, ["--version"]).stderr.decode().splitlines()[0],
            "libbz2_version": lib.BZ2_bzlibVersion().decode(),
        },
        "fixture": {
            "generator": "SHAKE-256", "seed_ascii": SEED.decode(),
            "payload_bytes": len(expected), "payload_sha256": digest(expected),
            "compressed_bytes": len(original), "compressed_sha256": digest(original),
            "block_count": len(blocks), "block_start_bits_msb_first": starts,
            "block_decoded_bytes": [len(part) for part in parts],
            "block_sha256": [digest(part) for part in parts],
            "block_crc_hex": [f"{crc:08x}" for crc in crcs],
            "stored_combined_crc_hex": f"{combined(crcs):08x}",
        },
        "observations": rows,
        "scope": {
            "hardware_faults": False, "random_fault_rate_estimated": False,
            "independent_decoder_implementations": False,
            "one_based_block_positions": True,
        },
    }


def strict_json(path):
    def pairs(items):
        obj = {}
        for key, value in items:
            require(key not in obj, f"duplicate JSON key: {key}")
            obj[key] = value
        return obj

    def invalid(value):
        raise ValueError(f"non-finite JSON constant: {value}")

    return json.loads(Path(path).read_text(), object_pairs_hook=pairs,
                      parse_constant=invalid)


def canonical(value):
    return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False)


def semantics(result):
    return {key: result[key] for key in ("schema_version", "kind", "fixture",
                                         "observations", "scope")}


def compare(retained, fresh):
    require(canonical(semantics(retained)) == canonical(semantics(fresh)),
            "retained evidence differs from fresh execution")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--bzip2", default="bzip2")
    target = parser.add_mutually_exclusive_group(required=True)
    target.add_argument("--output", type=Path)
    target.add_argument("--check", type=Path)
    args = parser.parse_args()
    result = run(args.bzip2)
    if args.check:
        compare(strict_json(args.check), result)
        print("Exact semantic replay passed: 6 cases, 33 native blocks.")
    else:
        with args.output.open("x") as stream:
            json.dump(result, stream, indent=2, allow_nan=False)
            stream.write("\n")
        print(f"Recorded 6 native cases in {args.output}")


if __name__ == "__main__":
    main()