#!/usr/bin/env python3
"""Generate the synthetic files used by the Parquet opening-failure article.

The lab separates wrong bytes, incomplete files, a valid footer with a damaged
checksummed page, a valid encrypted file that needs keys, and a valid file whose
missing extension ViewParquet rejects before DuckDB sees it. These are tiny
teaching fixtures, not a corruption-recovery suite or compatibility benchmark.
"""

from __future__ import annotations

import base64
import hashlib
import json
import shutil
import struct
import tempfile
from datetime import datetime, timedelta, timezone
from decimal import Decimal
from pathlib import Path
from typing import Any

import duckdb
import pyarrow as pa
import pyarrow.parquet as pq
import pyarrow.parquet.encryption as pe


SCRIPT_PATH = Path(__file__).resolve()
REPO_ROOT = (
    SCRIPT_PATH.parents[1]
    if SCRIPT_PATH.parent.name == "scripts" and (SCRIPT_PATH.parents[1] / "public").exists()
    else Path.cwd()
)
OUTPUT_DIR = REPO_ROOT / "public" / "data" / "seo" / "parquet-open-failure-lab"
GENERATED_ON = "2026-08-25"

VALID_CONTROL = OUTPUT_DIR / "valid-control.parquet"
VALID_NO_EXTENSION = OUTPUT_DIR / "valid-control-no-extension"
HTML_RESPONSE = OUTPUT_DIR / "html-error-response.parquet"
GIT_LFS_POINTER = OUTPUT_DIR / "git-lfs-pointer.parquet"
ZERO_BYTE = OUTPUT_DIR / "zero-byte.parquet"
TRUNCATED_FOOTER = OUTPUT_DIR / "truncated-footer.parquet"
UNCLOSED_WRITER = OUTPUT_DIR / "unclosed-writer.parquet"
BODY_CORRUPTION = OUTPUT_DIR / "footer-valid-body-corrupt.parquet"
ENCRYPTED_FOOTER = OUTPUT_DIR / "encrypted-footer.parquet"
MANIFEST = OUTPUT_DIR / "manifest.json"
PUBLIC_GENERATOR = OUTPUT_DIR / "generate_parquet_open_failure_fixtures.py"

CONTROL_ROWS = 5000
CONTROL_EVENT_ID_SUM = 12_497_500
CONTROL_PAYLOAD_LENGTH_SUM = 2_135_000

FIXTURE_ORDER = (
    VALID_CONTROL,
    VALID_NO_EXTENSION,
    HTML_RESPONSE,
    GIT_LFS_POINTER,
    ZERO_BYTE,
    TRUNCATED_FOOTER,
    UNCLOSED_WRITER,
    BODY_CORRUPTION,
    ENCRYPTED_FOOTER,
)


class InMemoryKmsClient(pe.KmsClient):
    """Deliberately insecure key wrapper for synthetic teaching data only."""

    def __init__(self, config: pe.KmsConnectionConfig):
        super().__init__()
        self.keys = config.custom_kms_conf

    def wrap_key(self, key_bytes: bytes, master_key_identifier: str) -> bytes:
        master_key = self.keys[master_key_identifier].encode("utf-8")
        return base64.b64encode(master_key + key_bytes)

    def unwrap_key(self, wrapped_key: bytes, master_key_identifier: str) -> bytes:
        decoded = base64.b64decode(wrapped_key)
        expected_key = self.keys[master_key_identifier].encode("utf-8")
        if decoded[: len(expected_key)] != expected_key:
            raise ValueError("Incorrect synthetic master key")
        return decoded[len(expected_key) :]


def make_crypto_environment() -> tuple[pe.KmsConnectionConfig, pe.CryptoFactory]:
    kms_config = pe.KmsConnectionConfig(
        custom_kms_conf={
            "footer_key": "0123456789112345",
            "column_key": "1234567890123450",
        }
    )
    crypto_factory = pe.CryptoFactory(lambda config: InMemoryKmsClient(config))
    return kms_config, crypto_factory


def sql_path(path: Path) -> str:
    return path.as_posix().replace("'", "''")


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def visible_hex(data: bytes) -> str:
    return data.hex(" ")


def portable_error(error: Exception, path: Path) -> str:
    first_line = str(error).splitlines()[0]
    return first_line.replace(str(path), path.name)


def run_duckdb_check(
    connection: duckdb.DuckDBPyConnection,
    path: Path,
    sql: str,
) -> dict[str, Any]:
    try:
        cursor = connection.execute(sql)
        rows = cursor.fetchall()
        return {
            "status": "ok",
            "columns": [column[0] for column in cursor.description],
            "rows": [list(row) for row in rows],
        }
    except Exception as error:  # DuckDB exposes multiple exception subclasses.
        return {"status": "error", "error": portable_error(error, path)}


def run_pyarrow_checks(path: Path) -> dict[str, Any]:
    try:
        parquet_file = pq.ParquetFile(path)
        footer = {
            "status": "ok",
            "rows": parquet_file.metadata.num_rows,
            "row_groups": parquet_file.metadata.num_row_groups,
            "columns": parquet_file.schema.names,
        }
    except Exception as error:
        footer = {"status": "error", "error": portable_error(error, path)}

    try:
        table = pq.read_table(path, page_checksum_verification=True)
        full_decode = {
            "status": "ok",
            "rows": table.num_rows,
            "columns": table.column_names,
        }
    except Exception as error:
        full_decode = {"status": "error", "error": portable_error(error, path)}

    return {"footer_metadata": footer, "full_decode_with_checksum_verification": full_decode}


def inspect_fixture(
    connection: duckdb.DuckDBPyConnection,
    path: Path,
    role: str,
    expected_viewparquet_entry: str,
) -> dict[str, Any]:
    data = path.read_bytes()
    escaped_path = sql_path(path)
    full_decode_expression = "count(*)" if path == ENCRYPTED_FOOTER else "sum(length(payload))"
    return {
        "name": path.name,
        "role": role,
        "bytes": len(data),
        "sha256": sha256(path),
        "first_16_bytes_hex": visible_hex(data[:16]),
        "last_16_bytes_hex": visible_hex(data[-16:]),
        "boundary_magic": {
            "first_four_ascii": data[:4].decode("ascii", errors="replace"),
            "last_four_ascii": data[-4:].decode("ascii", errors="replace"),
        },
        "native_duckdb": {
            "footer_schema": run_duckdb_check(
                connection,
                path,
                f"DESCRIBE SELECT * FROM read_parquet('{escaped_path}')",
            ),
            "metadata_count": run_duckdb_check(
                connection,
                path,
                f"SELECT count(*) AS rows FROM read_parquet('{escaped_path}')",
            ),
            "full_decode": run_duckdb_check(
                connection,
                path,
                f"SELECT {full_decode_expression} AS decoded_check "
                f"FROM read_parquet('{escaped_path}')",
            ),
        },
        "pyarrow": run_pyarrow_checks(path),
        "expected_viewparquet_entry": expected_viewparquet_entry,
    }


def write_control_and_corruption() -> dict[str, Any]:
    values = [
        f"row-{index:06d}-" + ("abcdefghijklmnopqrstuvwxyz" * 16)
        for index in range(CONTROL_ROWS)
    ]
    table = pa.table(
        {
            "event_id": pa.array(range(CONTROL_ROWS), type=pa.int32()),
            "payload": pa.array(values),
        }
    )
    pq.write_table(
        table,
        VALID_CONTROL,
        compression="gzip",
        use_dictionary=False,
        data_page_version="2.0",
        data_page_size=32 * 1024,
        write_page_checksum=True,
        write_batch_size=128,
    )

    valid_bytes = bytearray(VALID_CONTROL.read_bytes())
    if valid_bytes[:4] != b"PAR1" or valid_bytes[-4:] != b"PAR1":
        raise RuntimeError("Control file does not have ordinary Parquet boundaries")

    payload_column = pq.ParquetFile(VALID_CONTROL).metadata.row_group(0).column(1)
    data_page_offset = payload_column.data_page_offset
    gzip_offset = valid_bytes.find(b"\x1f\x8b\x08", data_page_offset)
    footer_length = struct.unpack("<I", valid_bytes[-8:-4])[0]
    footer_start = len(valid_bytes) - footer_length - 8
    if gzip_offset < data_page_offset or gzip_offset >= footer_start:
        raise RuntimeError("Could not locate a GZIP page inside the control body")

    original_byte = valid_bytes[gzip_offset]
    valid_bytes[gzip_offset] ^= 0x01
    BODY_CORRUPTION.write_bytes(valid_bytes)
    if BODY_CORRUPTION.read_bytes()[:4] != b"PAR1" or BODY_CORRUPTION.read_bytes()[-4:] != b"PAR1":
        raise RuntimeError("Corruption recipe unexpectedly changed the boundary markers")

    return {
        "fixture": BODY_CORRUPTION.name,
        "page_checksum_written": True,
        "payload_data_page_offset": data_page_offset,
        "flipped_byte_offset": gzip_offset,
        "original_byte_hex": f"{original_byte:02x}",
        "corrupted_byte_hex": f"{valid_bytes[gzip_offset]:02x}",
        "footer_start_offset": footer_start,
        "footer_copied_intact": True,
    }


def write_unclosed_writer_fixture(path: Path) -> None:
    schema = pa.schema([("event_id", pa.int32()), ("payload", pa.string())])
    table = pa.table(
        {
            "event_id": pa.array(range(64), type=pa.int32()),
            "payload": [f"row-{index}" for index in range(64)],
        }
    )

    with tempfile.TemporaryDirectory(prefix="viewparquet-unclosed-writer-") as temp_directory:
        in_progress = Path(temp_directory) / "in-progress.parquet"
        sink = pa.OSFile(str(in_progress), "wb")
        writer = pq.ParquetWriter(sink, schema, compression="NONE")
        try:
            writer.write_table(table)
            sink.flush()
            # Copy the bytes before close() writes the footer. Closing afterwards
            # only finalizes the temporary source, not the published fixture.
            shutil.copyfile(in_progress, path)
        finally:
            writer.close()
            sink.close()


def write_encrypted_fixture(path: Path) -> dict[str, Any]:
    table = pa.table(
        {
            "id": pa.array([1, 2, 3, 4], type=pa.int64()),
            "name": pa.array(["Ada", "Linus", None, "Grace"], type=pa.string()),
            "tags": pa.array(
                [["parquet", "duckdb"], [], None, ["arrow"]],
                type=pa.list_(pa.string()),
            ),
            "event_time": pa.array(
                [
                    datetime(2026, 8, 25, 8, 30, tzinfo=timezone.utc),
                    datetime(2026, 8, 25, 9, 0, tzinfo=timezone.utc),
                    None,
                    datetime(2026, 8, 25, 10, 15, tzinfo=timezone.utc),
                ],
                type=pa.timestamp("ms", tz="UTC"),
            ),
            "amount": pa.array(
                [Decimal("12.34"), Decimal("0.00"), None, Decimal("-5.50")],
                type=pa.decimal128(8, 2),
            ),
        }
    )
    kms_config, crypto_factory = make_crypto_environment()
    encryption_config = pe.EncryptionConfiguration(
        footer_key="footer_key",
        column_keys={"column_key": ["id", "name"]},
        encryption_algorithm="AES_GCM_V1",
        plaintext_footer=False,
        internal_key_material=True,
        cache_lifetime=timedelta(minutes=5),
        data_key_length_bits=256,
    )
    encryption_properties = crypto_factory.file_encryption_properties(
        kms_config,
        encryption_config,
    )
    pq.write_table(table, path, encryption_properties=encryption_properties)
    encrypted_bytes = path.read_bytes()
    if encrypted_bytes[:4] != b"PARE" or encrypted_bytes[-4:] != b"PARE":
        raise RuntimeError("Encrypted-footer fixture does not have PARE boundaries")

    # Recreate the crypto objects so validation does not depend on writer state.
    kms_config, crypto_factory = make_crypto_environment()
    decryption_properties = crypto_factory.file_decryption_properties(
        kms_config,
        pe.DecryptionConfiguration(cache_lifetime=timedelta(minutes=5)),
    )
    decoded = pq.read_table(path, decryption_properties=decryption_properties)
    if not decoded.equals(table):
        raise RuntimeError("Encrypted fixture did not round-trip with its synthetic test keys")
    return {
        "fixture": path.name,
        "rows_with_test_keys": decoded.num_rows,
        "round_trip_with_test_keys": True,
        "warning": (
            "The generator's key wrapper and embedded test keys are deliberately insecure and "
            "must never be copied into production."
        ),
        "byte_stability": (
            "Encrypted bytes are nondeterministic because PyArrow generates random data keys and nonces; "
            "use the manifest hash for this generated copy."
        ),
    }


def main() -> None:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    connection = duckdb.connect()

    corruption_recipe = write_control_and_corruption()
    valid_bytes = VALID_CONTROL.read_bytes()
    footer_length = struct.unpack("<I", valid_bytes[-8:-4])[0]
    footer_start = len(valid_bytes) - footer_length - 8
    shutil.copyfile(VALID_CONTROL, VALID_NO_EXTENSION)

    HTML_RESPONSE.write_bytes(
        b"<!doctype html><html><title>AccessDenied</title>"
        b"<body>403 AccessDenied</body></html>\n"
    )
    GIT_LFS_POINTER.write_bytes(
        b"version https://git-lfs.github.com/spec/v1\n"
        + b"oid sha256:"
        + (b"0" * 64)
        + b"\nsize 24576\n"
    )
    ZERO_BYTE.write_bytes(b"")
    # Remove the footer-length word and final PAR1 marker. The body and most of
    # the serialized footer remain, but a reader cannot locate the metadata.
    TRUNCATED_FOOTER.write_bytes(valid_bytes[:-8])
    write_unclosed_writer_fixture(UNCLOSED_WRITER)
    encrypted_validation = write_encrypted_fixture(ENCRYPTED_FOOTER)

    PUBLIC_GENERATOR.write_text(SCRIPT_PATH.read_text(encoding="utf-8"), encoding="utf-8")

    role_by_name = {
        VALID_CONTROL.name: "valid_checksummed_gzip_control",
        VALID_NO_EXTENSION.name: "valid_control_without_extension",
        HTML_RESPONSE.name: "html_response_renamed_parquet",
        GIT_LFS_POINTER.name: "git_lfs_pointer_renamed_parquet",
        ZERO_BYTE.name: "zero_byte_file",
        TRUNCATED_FOOTER.name: "truncated_before_footer_length_and_magic",
        UNCLOSED_WRITER.name: "pyarrow_writer_bytes_copied_before_close",
        BODY_CORRUPTION.name: "parseable_footer_with_one_damaged_checksummed_page",
        ENCRYPTED_FOOTER.name: "valid_encrypted_footer_requiring_keys",
    }
    expected_entry_by_name = {
        VALID_CONTROL.name: "accepted_then_opened",
        VALID_NO_EXTENSION.name: "rejected_before_duckdb_due_to_filename_extension",
        HTML_RESPONSE.name: "accepted_by_extension_then_duckdb_error",
        GIT_LFS_POINTER.name: "accepted_by_extension_then_duckdb_error",
        ZERO_BYTE.name: "accepted_by_extension_then_duckdb_error",
        TRUNCATED_FOOTER.name: "accepted_by_extension_then_duckdb_error",
        UNCLOSED_WRITER.name: "accepted_by_extension_then_duckdb_error",
        BODY_CORRUPTION.name: "accepted_by_extension_then_decode_error",
        ENCRYPTED_FOOTER.name: "accepted_by_extension_then_reader_requires_encryption_config",
    }
    fixtures = [
        inspect_fixture(
            connection,
            path,
            role_by_name[path.name],
            expected_entry_by_name[path.name],
        )
        for path in FIXTURE_ORDER
    ]

    manifest = {
        "purpose": "Reproducible teaching fixtures for diagnosing why a Parquet file will not open",
        "generated_on": GENERATED_ON,
        "generator": "scripts/generate_parquet_open_failure_fixtures.py",
        "native_duckdb_version": connection.execute("SELECT version()").fetchone()[0],
        "pyarrow_version": pa.__version__,
        "control": {
            "rows": CONTROL_ROWS,
            "event_id_sum": CONTROL_EVENT_ID_SUM,
            "payload_length_sum": CONTROL_PAYLOAD_LENGTH_SUM,
            "row_groups": pq.ParquetFile(VALID_CONTROL).metadata.num_row_groups,
            "footer_length_bytes": footer_length,
            "footer_start_offset": footer_start,
            "compression": "GZIP",
            "data_page_version": "2.0",
            "page_checksums": True,
        },
        "corruption_recipe": corruption_recipe,
        "encrypted_validation": encrypted_validation,
        "evidence_artifacts": [
            {
                "name": PUBLIC_GENERATOR.name,
                "role": "fixture_generator",
                "bytes": PUBLIC_GENERATOR.stat().st_size,
                "sha256": sha256(PUBLIC_GENERATOR),
            }
        ],
        "limitations": [
            "Tiny synthetic teaching files; not a representative corruption or compatibility corpus",
            "Exact exception wording is specific to the recorded reader versions and can change",
            "The internal-corruption fixture flips one GZIP header byte and does not model every damaged page",
            "Boundary magic and a parseable footer do not prove that every data page decodes",
            "Encrypted output is semantically reproducible but not byte-stable because its data keys and nonces are random",
            "Remote authentication, CORS, range handling, and browser memory are documented separately and are not represented by these local files",
        ],
        "fixtures": fixtures,
    }
    MANIFEST.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(manifest, indent=2))


if __name__ == "__main__":
    main()