#!/usr/bin/env python3
"""Bounded synthetic Zstandard dictionary and frame-completeness experiment."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import platform
import re
import shutil
import subprocess
import tempfile
from datetime import datetime, timezone


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


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


def run(args, payload):
    env = {k: v for k, v in os.environ.items() if not k.startswith("ZSTD")}
    env["LC_ALL"] = "C"
    return subprocess.run(["zstd", *args], input=payload, capture_output=True,
                          timeout=15, check=False, env=env)


def compress(payload, dictionary=None, checksum=True):
    args = ["-q", "-c", "-3", "--single-thread",
            "--check" if checksum else "--no-check"]
    if dictionary is not None:
        args += ["-D", str(dictionary), "--no-dictID"]
    result = run(args, payload)
    require(result.returncode == 0, result.stderr.decode("utf-8", "replace"))
    require(result.stdout[:4] == bytes.fromhex("28b52ffd"), "Expected zstd frame")
    return result.stdout


def observe(case_id, frame, expected, dictionary=None, verify_checksum=True):
    args = ["-q", "-d", "-c", "--check" if verify_checksum else "--no-check"]
    if dictionary is not None:
        args += ["-D", str(dictionary)]
    result = run(args, frame)
    return {
        "id": case_id, "exitCode": result.returncode,
        "checksumPolicy": "verify-if-present" if verify_checksum else "ignore",
        "outputBytes": len(result.stdout), "outputSha256": digest(result.stdout),
        "expectedBytes": len(expected), "expectedSha256": digest(expected),
        "bytesMatch": result.stdout == expected,
        "accepted": result.returncode == 0 and result.stdout == expected,
        "stderr": result.stderr.decode("utf-8", "replace").strip(),
        "compressedBytes": len(frame), "compressedSha256": digest(frame),
    }


CASE_EXPECTATIONS = {
    "plain-control": (True, True),
    "dictionary-correct": (True, True),
    "dictionary-missing": (False, False),
    "dictionary-wrong-checked": (False, False),
    "dictionary-wrong-unchecked": (True, False),
    "dictionary-correct-unchecked": (True, True),
    "two-frames-complete": (True, True),
    "second-frame-removed": (True, False),
    "last-byte-truncated": (False, None),
    "dictionary-check-ignored": (True, False),
    "checksum-damaged-verified": (False, None),
    "checksum-damaged-ignored": (True, True),
    "skippable-metadata-original": (True, True),
    "skippable-metadata-changed": (True, True),
}


def validate(rows):
    require(len(rows) == len(CASE_EXPECTATIONS), "Unexpected number of cases")
    require({row["id"] for row in rows} == set(CASE_EXPECTATIONS), "Invalid case IDs")
    for row in rows:
        success, match = CASE_EXPECTATIONS[row["id"]]
        require(type(row["exitCode"]) is int, "Invalid exit status")
        require(type(row["outputBytes"]) is int and row["outputBytes"] >= 0,
                "Invalid output length")
        for key in ("expectedBytes", "compressedBytes"):
            require(type(row[key]) is int and row[key] >= 0, "Invalid " + key)
        for key in ("outputSha256", "expectedSha256", "compressedSha256"):
            require(isinstance(row[key], str) and re.fullmatch("[0-9a-f]{64}", row[key]),
                    "Invalid " + key)
        require(type(row["bytesMatch"]) is bool, "Invalid match flag")
        require(type(row["accepted"]) is bool, "Invalid acceptance flag")
        ignored = row["id"] in {"dictionary-check-ignored", "checksum-damaged-ignored"}
        require(row["checksumPolicy"] == ("ignore" if ignored else "verify-if-present"),
                row["id"] + ": changed checksum policy")
        require((row["exitCode"] == 0) == success, row["id"] + ": changed exit result")
        if match is not None:
            require(row["bytesMatch"] == match, row["id"] + ": changed byte result")
        require(row["accepted"] == (success and row["bytesMatch"]),
                row["id"] + ": invalid acceptance")
        require(row["bytesMatch"] == (
            row["outputBytes"] == row["expectedBytes"]
            and row["outputSha256"] == row["expectedSha256"]
        ), row["id"] + ": inconsistent byte oracle")

    by_id = {row["id"]: row for row in rows}
    for left, right in [("dictionary-wrong-checked", "dictionary-check-ignored"),
                        ("checksum-damaged-verified", "checksum-damaged-ignored")]:
        require(by_id[left]["compressedSha256"] == by_id[right]["compressedSha256"],
                "Checksum-policy comparison must use identical compressed input")
    a, b = (by_id["skippable-metadata-" + suffix] for suffix in ("original", "changed"))
    require(a["compressedBytes"] == b["compressedBytes"], "Metadata size changed")
    require(a["compressedSha256"] != b["compressedSha256"], "Metadata did not change")
    require(a["outputSha256"] == b["outputSha256"] == by_id["plain-control"]["outputSha256"],
            "Skippable metadata unexpectedly changed decoded payload")


def skippable(metadata):
    require(len(metadata) <= 4096, "Synthetic metadata limit exceeded")
    return bytes.fromhex("502a4d18") + len(metadata).to_bytes(4, "little") + metadata


def experiment():
    require(shutil.which("zstd") is not None, "zstd CLI is required")
    # Raw dictionary, not a trained/formatted dictionary.
    dictionary = b"".join(hashlib.sha256(("hesela-" + str(i)).encode()).digest()
                          for i in range(256))
    payload = dictionary[2048:4096]
    wrong = bytes(value ^ 0xFF for value in dictionary)
    with tempfile.TemporaryDirectory(prefix="hesela-zstd-") as temp:
        good_path, wrong_path = Path(temp) / "correct.raw", Path(temp) / "wrong.raw"
        good_path.write_bytes(dictionary)
        wrong_path.write_bytes(wrong)
        plain = compress(payload)
        checked = compress(payload, good_path)
        unchecked = compress(payload, good_path, checksum=False)
        # RFC 8878 descriptor byte: low two bits encode Dictionary_ID field size.
        require(checked[4] & 3 == 0 and unchecked[4] & 3 == 0,
                "Dictionary ID unexpectedly present")
        require(checked[4] & 4 != 0 and unchecked[4] & 4 == 0,
                "Unexpected checksum flags")
        first, second = compress(b"record-one\n"), compress(b"record-two\n")
        whole = b"record-one\nrecord-two\n"
        require(plain[4] & 4 != 0, "Control must include a checksum")
        damaged_checksum = plain[:-1] + bytes([plain[-1] ^ 1])
        metadata = b"generation=0001"
        changed_metadata = b"generation=0002"
        rows = [
            observe("plain-control", plain, payload),
            observe("dictionary-correct", checked, payload, good_path),
            observe("dictionary-missing", checked, payload),
            observe("dictionary-wrong-checked", checked, payload, wrong_path),
            observe("dictionary-wrong-unchecked", unchecked, payload, wrong_path),
            observe("dictionary-correct-unchecked", unchecked, payload, good_path),
            observe("two-frames-complete", first + second, whole),
            observe("second-frame-removed", first, whole),
            observe("last-byte-truncated", (first + second)[:-1], whole),
            observe("dictionary-check-ignored", checked, payload, wrong_path,
                    verify_checksum=False),
            observe("checksum-damaged-verified", damaged_checksum, payload),
            observe("checksum-damaged-ignored", damaged_checksum, payload,
                    verify_checksum=False),
            observe("skippable-metadata-original", skippable(metadata) + plain, payload),
            observe("skippable-metadata-changed", skippable(changed_metadata) + plain, payload),
        ]
        validate(rows)
    return {
        "schemaVersion": 1,
        "generatedAt": datetime.now(timezone.utc).isoformat(),
        "sourceSha256": digest(Path(__file__).read_bytes()),
        "environment": {"zstd": run(["--version"], b"").stdout.decode().strip(),
                        "python": platform.python_version(),
                        "system": platform.system(), "machine": platform.machine()},
        "method": {
            "type": "synthetic-codec-contract-demonstration",
            "dictionaryType": "raw-content", "dictionaryBytes": len(dictionary),
            "dictionarySha256": digest(dictionary),
            "wrongDictionarySha256": digest(wrong),
            "payloadBytes": len(payload), "payloadSha256": digest(payload),
            "dictionaryIdPresent": False, "compressionLevel": 3,
            "commandTimeoutSeconds": 15,
            "acceptance": "zero exit status AND exact expected decoded bytes; payload only, not metadata",
            "skippableMetadata": metadata.decode("ascii"),
            "changedSkippableMetadata": changed_metadata.decode("ascii"),
            "checksumMutation": "XOR final checksum byte with 1; compressed data unchanged",
            "scope": "One CLI build; no disk corruption rate, latency or power-loss claim.",
        },
        "observations": rows,
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, help="JSON output path; stdout by default")
    args = parser.parse_args()
    result = json.dumps(experiment(), indent=2) + "\n"
    if args.output:
        args.output.write_text(result, encoding="utf-8")
    else:
        print(result, end="")


if __name__ == "__main__":
    main()
