#!/usr/bin/env python3
"""CC0-1.0. Controlled Parquet footer experiment; creates temporary fixtures only."""
import argparse
import copy
import hashlib
import importlib.metadata
import json
import platform
import struct
import sys
import tempfile
import zlib
from datetime import datetime, timezone
from pathlib import Path

import pyarrow as pa
import pyarrow.compute as pc
import pyarrow.dataset as ds
import pyarrow.parquet as pq
import thriftpy2
from thriftpy2.protocol.compact import TCompactProtocol
from thriftpy2.transport.memory import TMemoryBuffer

SCHEMA_URL = "https://raw.githubusercontent.com/apache/parquet-format/apache-parquet-format-2.10.0/src/main/thrift/parquet.thrift"
SCHEMA_HASH = "f3152ab8a2574c96f4b0ea3df7694f9170e6a6c18eb9888505115f9acd5e54ce"
VERSIONS = {"pyarrow": "21.0.0", "thriftpy2": "0.5.3", "ply": "3.11", "six": "1.17.0"}


def require(ok, message):
    if not ok:
        raise RuntimeError(message)


def sha(raw):
    return hashlib.sha256(raw).hexdigest()


class CountingBuffer(TMemoryBuffer):
    def __init__(self, raw):
        super().__init__(raw)
        self.consumed = 0

    def read(self, size):
        value = super().read(size)
        self.consumed += len(value)
        return value


def decode(cls, raw):
    transport = CountingBuffer(raw)
    value = cls()
    value.read(TCompactProtocol(transport))
    return value, transport.consumed


def encode(value):
    transport = TMemoryBuffer()
    value.write(TCompactProtocol(transport))
    return transport.getvalue()


def footer(raw, schema):
    require(raw[:4] == raw[-4:] == b"PAR1", "Only plaintext Parquet fixtures supported")
    length = struct.unpack("<I", raw[-8:-4])[0]
    start = len(raw) - 8 - length
    require(start >= 4, "Invalid footer extent")
    value, consumed = decode(schema.FileMetaData, raw[start:-8])
    require(consumed == length, "Footer parser did not consume exact footer")
    require(encode(value) == raw[start:-8], "Unmodified footer must roundtrip byte-for-byte")
    return start, value


def page_receipts(raw, schema, metadata):
    receipts = []
    for rg_index, group in enumerate(metadata.row_groups):
        for col_index, chunk in enumerate(group.columns):
            col = chunk.meta_data
            require(col.codec == schema.CompressionCodec.UNCOMPRESSED, "Unexpected compression")
            require(col.dictionary_page_offset is None, "Unexpected dictionary")
            offset = col.data_page_offset
            end = offset + col.total_compressed_size
            count = 0
            while offset < end:
                header, header_bytes = decode(schema.PageHeader, raw[offset:end])
                body_offset = offset + header_bytes
                body_end = body_offset + header.compressed_page_size
                require(body_end <= end and header_bytes > 0, "Page outside column extent")
                require(header.crc is not None, "Missing CRC")
                crc = zlib.crc32(raw[body_offset:body_end]) & 0xffffffff
                require(crc == (header.crc & 0xffffffff), "Independent page CRC mismatch")
                receipts.append({"rowGroup": rg_index, "column": col_index,
                                 "bodyOffset": body_offset, "bodyBytes": body_end-body_offset,
                                 "crc32": crc})
                offset = body_end
                count += 1
            require(offset == end and count == 1, "Expected one page per column chunk")
    require(len(receipts) == 4, "Expected four checked pages")
    return receipts


def variant(raw, start, metadata, mode):
    changed = copy.deepcopy(metadata)
    stats = changed.row_groups[0].columns[1].meta_data.statistics
    if mode == "narrow":
        require(struct.unpack("<q", stats.min_value)[0] == 100000, "Wrong original minimum")
        for key in ("min", "min_value"):
            require(getattr(stats, key) == struct.pack("<q", 100000), "Missing mirrored minimum")
            setattr(stats, key, struct.pack("<q", 100032))
    elif mode == "absent":
        changed.row_groups[0].columns[1].meta_data.statistics = None
    elif mode == "wide":
        for key, value in (("min", 99000), ("min_value", 99000),
                           ("max", 101063), ("max_value", 101063)):
            setattr(stats, key, struct.pack("<q", value))
        stats.is_min_value_exact = False
        stats.is_max_value_exact = False
    else:
        raise ValueError(mode)
    serialized = encode(changed)
    output = raw[:start] + serialized + struct.pack("<I", len(serialized)) + b"PAR1"
    require(output[:start] == raw[:start], "Page region changed")
    return output


def read(path, filtered=False):
    options = {"page_checksum_verification": True, "use_threads": False}
    if filtered:
        options["filters"] = [("measurement", "=", 100017)]
    return pq.read_table(path, **options)


def eligible(path):
    fragments = list(ds.dataset(path, format="parquet").get_fragments())
    require(len(fragments) == 1, "Expected one file fragment")
    parts = fragments[0].split_by_row_group(filter=ds.field("measurement") == 100017)
    return [rg.id for part in parts for rg in part.row_groups]


def run_once(root, schema):
    table = pa.Table.from_arrays(
        [pa.array(range(128), type=pa.int64()), pa.array(range(100000, 100128), type=pa.int64())],
        schema=pa.schema([pa.field("id", pa.int64(), nullable=False),
                          pa.field("measurement", pa.int64(), nullable=False)]))
    clean = root / "clean.parquet"
    pq.write_table(table, clean, row_group_size=64, compression="NONE", use_dictionary=False,
                   write_statistics=True, write_page_checksum=True, data_page_version="1.0",
                   version="2.6", write_page_index=False)
    raw = clean.read_bytes()
    start, meta = footer(raw, schema)
    pages = page_receipts(raw, schema, meta)
    require(read(clean).equals(table), "Clean full-read oracle mismatch")
    paths = {"clean": clean}
    artifacts = {}
    for mode in ("narrow", "absent", "wide"):
        changed = variant(raw, start, meta, mode)
        if mode == "narrow":
            require(len(changed) == len(raw), "Narrowed footer changed file size")
            byte_changes = [{"offset": i, "before": a, "after": b}
                            for i, (a, b) in enumerate(zip(raw, changed)) if a != b]
            require(len(byte_changes) == 2 and all(c["offset"] >= start for c in byte_changes),
                    "Only two mirrored minimum bytes may change")
        path = root / (mode + ".parquet")
        path.write_bytes(changed)
        paths[mode] = path
        _, parsed = footer(changed, schema)
        require(page_receipts(changed, schema, parsed) == pages, "Page checks changed")
        require(read(path).equals(table), "Footer mutation changed decoded values")
        artifacts[mode] = {"sha256": sha(changed), "bytes": len(changed),
                           "pageRegionSha256": sha(changed[:start]),
                           "eligibleRowGroups": eligible(path)}
        if mode == "narrow":
            artifacts[mode]["byteChanges"] = byte_changes
            artifacts[mode]["changedFields"] = ["row_groups[0].columns[1].meta_data.statistics.min",
                                               "row_groups[0].columns[1].meta_data.statistics.min_value"]
    altered_meta = pq.read_metadata(paths["narrow"])
    minimum = altered_meta.row_group(0).column(1).statistics.min
    maximum = altered_meta.row_group(0).column(1).statistics.max
    require((minimum, maximum) == (100032, 100063), "Unexpected narrowed statistics")
    oracle = table.slice(17, 1)
    observations = []
    for mode in ("clean", "narrow", "absent", "wide"):
        result = read(paths[mode], filtered=True)
        wanted = table.slice(0, 0) if mode == "narrow" else oracle
        require(result.equals(wanted), "Unexpected filtered result in " + mode)
        observations.append({"case": mode + "-filtered", "rows": result.num_rows,
                             "ids": result["id"].to_pylist(), "expectedRows": 1,
                             "matchesOracle": result.equals(oracle), "checksumVerification": True})
    full = read(paths["narrow"])
    observations.append({"case": "narrow-full-scan", "rows": full.num_rows,
                         "expectedRows": 128, "matchesOracle": full.equals(table),
                         "checksumVerification": True})
    post = full.filter(pc.equal(full["measurement"], 100017))
    require(post.equals(oracle), "Post-scan filter lost target")
    observations.append({"case": "narrow-post-scan-filter", "rows": post.num_rows,
                         "ids": post["id"].to_pylist(), "expectedRows": 1,
                         "matchesOracle": post.equals(oracle), "checksumVerification": True})
    target_page = next(p for p in pages if p["rowGroup"] == 0 and p["column"] == 1)
    require(target_page["bodyBytes"] == 64 * 8, "Expected required PLAIN INT64 page body")
    body_fault = bytearray(raw)
    body_fault[target_page["bodyOffset"] + 17 * 8] ^= 1
    bad = root / "body-fault.parquet"
    bad.write_bytes(body_fault)
    try:
        read(bad)
    except OSError as error:
        require("CRC checksum verification failed" in str(error), "Unexpected failure")
        observations.append({"case": "body-fault-positive-control", "rows": None,
                             "expectedRows": 128, "matchesOracle": False,
                             "checksumVerification": True, "outcome": "checksum-error"})
    else:
        raise RuntimeError("Payload control did not detect CRC error")
    require(eligible(clean) == [0], "Clean pruning control failed")
    require(artifacts["narrow"]["eligibleRowGroups"] == [], "Narrowed group was not pruned")
    require(all(artifacts[m]["eligibleRowGroups"] == [0] for m in ("absent", "wide")),
            "Pruning controls failed")
    actual_min = pc.min(full["measurement"].slice(0, 64)).as_py()
    actual_max = pc.max(full["measurement"].slice(0, 64)).as_py()
    require(not minimum <= actual_min <= actual_max <= maximum, "Bounds audit missed fault")
    return {"originalSha256": sha(raw), "originalBytes": len(raw), "footerOffset": start,
            "pageRegionSha256": sha(raw[:start]), "pages": pages, "variants": artifacts,
            "statsAudit": {"storedMin": minimum, "storedMax": maximum,
                           "decodedMin": actual_min, "decodedMax": actual_max,
                           "boundsContainValues": False},
            "observations": observations}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--schema", required=True, type=Path)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--repetitions", default=3, type=int)
    args = parser.parse_args()
    require(args.output is None or not (args.output.exists() or args.output.is_symlink()),
            "Output already exists; choose a new path")
    require(sys.flags.optimize < 2, "Use normal Python or -O; the pinned Thrift parser requires docstrings and does not support -OO")
    require(1 <= args.repetitions <= 20, "Repetitions must be 1..20")
    for package, version in VERSIONS.items():
        require(importlib.metadata.version(package) == version, "Version mismatch: " + package)
    require(sha(args.schema.read_bytes()) == SCHEMA_HASH, "Schema hash mismatch")
    schema = thriftpy2.load(str(args.schema), module_name="hesela_parquet_stats_thrift")
    runs = []
    for _ in range(args.repetitions):
        with tempfile.TemporaryDirectory(prefix="hesela-parquet-stats-") as directory:
            runs.append(run_once(Path(directory), schema))
    require(all(run == runs[0] for run in runs), "Fresh-fixture repetitions disagree")
    result = {"schemaVersion": 1, "experiment": "parquet-statistics-pruning",
              "evidenceType": "measured-api-behavior", "synthetic": True,
              "generatedAt": datetime.now(timezone.utc).isoformat(),
              "sourceSha256": sha(Path(__file__).read_bytes()),
              "requirementsSha256": sha(Path(__file__).with_name("parquet-statistics-requirements.txt").read_bytes()),
              "schemaUrl": SCHEMA_URL, "schemaSha256": SCHEMA_HASH,
              "environment": {"python": platform.python_version(), "packages": VERSIONS,
                              "arrowCpp": pa.cpp_version, "system": platform.system(),
                              "kernel": platform.release(), "machine": platform.machine()},
              "repetitions": args.repetitions, "casesPerRepetition": 7,
              "method": "Generate 128 synthetic rows; structurally change first-group measurement minimum in plaintext footer from 100000 to 100032; keep all page bytes unchanged. Compare CRC-enabled scans and filtered reads, missing/wide-statistics controls, independent page CRCs and an altered-page control.",
              "limits": "One implementation and host; deliberately selected, structurally valid but semantically invalid metadata. Not a spontaneous bit-error model, prevalence estimate, security exploit, benchmark or test of encrypted, compressed, nested or dictionary files. Repetitions are not independent reliability trials.",
              "runs": runs}
    text = json.dumps(result, indent=2) + "\n"
    if args.output:
        with args.output.open("x", encoding="utf-8") as output:
            output.write(text)
    else:
        print(text, end="")


if __name__ == "__main__":
    main()