#!/usr/bin/env python3
"""Validate SigNoz ClickHouse native-backup inventories without data access.

The shell restore-drill driver uses this helper for two narrow jobs:

* prove that a local artifact is a completed, hashed native BACKUP with the
  exact six-database source inventory expected by P0-OBS-002; and
* compare the restored database/table/engine set with that source inventory.

Row and byte values are recorded as bounded aggregate evidence, not exact
parity gates, because the source inventory is captured while online ingestion
can still advance before the native BACKUP snapshot is committed.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
import sys
from dataclasses import dataclass
from pathlib import Path


EXPECTED_DATABASES = (
    "signoz_analytics",
    "signoz_logs",
    "signoz_metadata",
    "signoz_meter",
    "signoz_metrics",
    "signoz_traces",
)
CRITICAL_TABLES = (
    ("signoz_logs", "logs_v2"),
    ("signoz_metrics", "time_series_v4"),
    ("signoz_metrics", "samples_v4"),
    ("signoz_traces", "signoz_index_v3"),
)
SAFE_IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
SAFE_OPERATION_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$")
SHA256 = re.compile(r"^[0-9a-f]{64}$")
SCHEMA_SHA256 = re.compile(r"^[0-9A-F]{64}$")


class ContractError(ValueError):
    """A bounded, non-secret backup or inventory contract failure."""


@dataclass(frozen=True)
class TableReceipt:
    database: str
    table: str
    engine: str
    rows: int
    bytes: int
    schema_hash: str


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


def parse_inventory(path: Path) -> dict[tuple[str, str], TableReceipt]:
    receipts: dict[tuple[str, str], TableReceipt] = {}
    try:
        lines = path.read_text(encoding="utf-8").splitlines()
    except (OSError, UnicodeError) as exc:
        raise ContractError("inventory_read_failed") from exc

    if not lines:
        raise ContractError("inventory_empty")

    for line_number, line in enumerate(lines, start=1):
        fields = line.split("\t")
        if len(fields) != 6:
            raise ContractError(f"inventory_field_count_invalid_line_{line_number}")
        database, table, engine, rows_text, bytes_text, schema_hash = fields
        if not SAFE_IDENTIFIER.fullmatch(database):
            raise ContractError(
                f"inventory_database_identifier_unsafe_line_{line_number}"
            )
        if not SAFE_IDENTIFIER.fullmatch(table):
            raise ContractError(f"inventory_table_identifier_unsafe_line_{line_number}")
        if not SAFE_IDENTIFIER.fullmatch(engine):
            raise ContractError(
                f"inventory_engine_identifier_unsafe_line_{line_number}"
            )
        if not rows_text.isdecimal() or not bytes_text.isdecimal():
            raise ContractError(f"inventory_counter_invalid_line_{line_number}")
        if not SCHEMA_SHA256.fullmatch(schema_hash):
            raise ContractError(f"inventory_schema_hash_invalid_line_{line_number}")
        key = (database, table)
        if key in receipts:
            raise ContractError(f"inventory_duplicate_table_line_{line_number}")
        receipts[key] = TableReceipt(
            database=database,
            table=table,
            engine=engine,
            rows=int(rows_text),
            bytes=int(bytes_text),
            schema_hash=schema_hash,
        )
    return receipts


def require_database_set(receipts: dict[tuple[str, str], TableReceipt]) -> None:
    observed = tuple(sorted({receipt.database for receipt in receipts.values()}))
    if observed != EXPECTED_DATABASES:
        raise ContractError("signoz_database_set_mismatch")


def require_critical_tables(
    receipts: dict[tuple[str, str], TableReceipt], *, require_nonzero: bool
) -> None:
    for key in CRITICAL_TABLES:
        receipt = receipts.get(key)
        if receipt is None:
            raise ContractError(f"critical_table_missing_{key[0]}_{key[1]}")
        if require_nonzero and receipt.rows <= 0:
            raise ContractError(f"critical_table_empty_{key[0]}_{key[1]}")


def emit(values: list[tuple[str, object]]) -> None:
    for key, value in values:
        print(f"{key}={value}")


def preflight(args: argparse.Namespace) -> None:
    manifest_path = Path(args.manifest)
    artifact_path = Path(args.artifact)
    inventory_path = Path(args.source_inventory)
    try:
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    except (OSError, UnicodeError, json.JSONDecodeError) as exc:
        raise ContractError("backup_manifest_read_failed") from exc
    if not isinstance(manifest, dict):
        raise ContractError("backup_manifest_not_object")
    if manifest.get("status") != "BACKUP_CREATED":
        raise ContractError("backup_terminal_not_BACKUP_CREATED")

    artifact_sha256 = sha256_file(artifact_path)
    inventory_sha256 = sha256_file(inventory_path)
    artifact_size = artifact_path.stat().st_size
    manifest_sha256 = manifest.get("sha256")
    manifest_inventory_sha256 = manifest.get("source_inventory_sha256")
    manifest_size = manifest.get("size_bytes")
    if not isinstance(manifest_sha256, str) or not SHA256.fullmatch(manifest_sha256):
        raise ContractError("backup_manifest_sha256_invalid")
    if manifest_sha256 != artifact_sha256:
        raise ContractError("backup_artifact_sha256_mismatch")
    if manifest_inventory_sha256 != inventory_sha256:
        raise ContractError("source_inventory_sha256_mismatch")
    if not isinstance(manifest_size, int) or manifest_size <= 0:
        raise ContractError("backup_manifest_size_invalid")
    if manifest_size != artifact_size:
        raise ContractError("backup_artifact_size_mismatch")
    if tuple(manifest.get("databases", ())) != EXPECTED_DATABASES:
        raise ContractError("backup_manifest_database_allowlist_mismatch")

    operation_id = manifest.get("operation_id")
    if not isinstance(operation_id, str) or not SAFE_OPERATION_ID.fullmatch(
        operation_id
    ):
        raise ContractError("backup_operation_id_invalid")

    receipts = parse_inventory(inventory_path)
    require_database_set(receipts)
    require_critical_tables(receipts, require_nonzero=True)
    replicated_count = sum(
        receipt.engine.startswith("Replicated") for receipt in receipts.values()
    )
    distributed_count = sum(
        receipt.engine == "Distributed" for receipt in receipts.values()
    )
    physical_count = sum("MergeTree" in receipt.engine for receipt in receipts.values())
    if replicated_count <= 0:
        raise ContractError("replicated_table_precondition_missing")
    if distributed_count <= 0:
        raise ContractError("distributed_table_precondition_missing")
    if physical_count <= 0:
        raise ContractError("physical_table_precondition_missing")

    emit(
        [
            ("artifact_sha256", artifact_sha256),
            ("artifact_size_bytes", artifact_size),
            ("source_inventory_sha256", inventory_sha256),
            ("source_database_count", len(EXPECTED_DATABASES)),
            ("source_table_count", len(receipts)),
            ("replicated_table_count", replicated_count),
            ("distributed_table_count", distributed_count),
            ("physical_table_count", physical_count),
            ("backup_operation_id", operation_id),
        ]
    )


def compare(args: argparse.Namespace) -> None:
    source = parse_inventory(Path(args.source_inventory))
    restored = parse_inventory(Path(args.restored_inventory))
    require_database_set(source)
    require_database_set(restored)
    require_critical_tables(source, require_nonzero=True)
    require_critical_tables(restored, require_nonzero=True)

    source_keys = set(source)
    restored_keys = set(restored)
    if source_keys != restored_keys:
        raise ContractError("database_table_manifest_set_mismatch")
    for key in sorted(source_keys):
        if source[key].engine != restored[key].engine:
            raise ContractError(f"table_engine_mismatch_{key[0]}_{key[1]}")
        if source[key].schema_hash != restored[key].schema_hash:
            raise ContractError(f"table_schema_mismatch_{key[0]}_{key[1]}")

    source_rows = sum(receipt.rows for receipt in source.values())
    restored_rows = sum(receipt.rows for receipt in restored.values())
    source_bytes = sum(receipt.bytes for receipt in source.values())
    restored_bytes = sum(receipt.bytes for receipt in restored.values())
    physical_count = sum("MergeTree" in receipt.engine for receipt in restored.values())
    emit(
        [
            ("manifest_parity", 1),
            ("restored_database_count", len(EXPECTED_DATABASES)),
            ("restored_table_count", len(restored)),
            ("physical_table_count", physical_count),
            ("critical_nonzero_count", len(CRITICAL_TABLES)),
            ("source_total_rows", source_rows),
            ("restored_total_rows", restored_rows),
            ("row_delta", restored_rows - source_rows),
            ("source_total_bytes", source_bytes),
            ("restored_total_bytes", restored_bytes),
            ("byte_delta", restored_bytes - source_bytes),
        ]
    )


def check_list(args: argparse.Namespace) -> None:
    receipts = parse_inventory(Path(args.inventory))
    require_database_set(receipts)
    for key in sorted(receipts):
        receipt = receipts[key]
        if "MergeTree" in receipt.engine:
            print(f"{receipt.database}\t{receipt.table}")


def parser() -> argparse.ArgumentParser:
    root = argparse.ArgumentParser()
    commands = root.add_subparsers(dest="command", required=True)

    preflight_parser = commands.add_parser("preflight")
    preflight_parser.add_argument("--manifest", required=True)
    preflight_parser.add_argument("--artifact", required=True)
    preflight_parser.add_argument("--source-inventory", required=True)
    preflight_parser.set_defaults(handler=preflight)

    compare_parser = commands.add_parser("compare")
    compare_parser.add_argument("--source-inventory", required=True)
    compare_parser.add_argument("--restored-inventory", required=True)
    compare_parser.set_defaults(handler=compare)

    check_parser = commands.add_parser("check-list")
    check_parser.add_argument("--inventory", required=True)
    check_parser.set_defaults(handler=check_list)
    return root


def main() -> int:
    args = parser().parse_args()
    try:
        args.handler(args)
    except (ContractError, OSError) as exc:
        print(f"ERROR={exc}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())