#!/usr/bin/env python3
"""Korvos AUSTRAC export bundle integrity verifier.

Reproduces the integrity checks documented at:
    https://korvos.com.au/trust/hash-chain

Uses only the Python standard library (>= 3.10). No external dependencies.
Runs offline. No network calls. No Korvos infrastructure required.

Usage
-----
    python3 verify.py path/to/korvos-austrac-export.zip

Exit codes
----------
    0  every check passed
    1  one or more file SHA-256s drifted from the manifest
    2  one or more rows failed the chain recomputation
    3  bundle is malformed or missing required files
    4  chain head in chain_verifications.json disagrees with the JSONL tail
"""

from __future__ import annotations

import hashlib
import json
import sys
import zipfile
from typing import Any, Iterable


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def sha256_hex(data: bytes) -> str:
    """Return the lowercase hex SHA-256 of the given bytes."""
    return hashlib.sha256(data).hexdigest()


def read_zip_text(z: zipfile.ZipFile, name: str) -> str:
    """Read a UTF-8 text file out of the bundle, raising with a clear message."""
    try:
        return z.read(name).decode("utf-8")
    except KeyError as exc:
        raise SystemExit(f"[3] bundle is missing {name!r}: {exc}")


def read_zip_bytes(z: zipfile.ZipFile, name: str) -> bytes:
    try:
        return z.read(name)
    except KeyError as exc:
        raise SystemExit(f"[3] bundle is missing {name!r}: {exc}")


def iter_jsonl(text: str) -> Iterable[dict[str, Any]]:
    for i, line in enumerate(text.splitlines(), start=1):
        line = line.strip()
        if not line:
            continue
        try:
            yield json.loads(line)
        except json.JSONDecodeError as exc:
            raise SystemExit(f"[3] JSONL parse failure at line {i}: {exc}")


# ---------------------------------------------------------------------------
# Checks
# ---------------------------------------------------------------------------


def check_file_hashes(z: zipfile.ZipFile, manifest: dict[str, Any]) -> list[str]:
    """Verify chains/*.jsonl SHA-256s match what manifest.json records."""
    failures: list[str] = []
    for entry in manifest.get("files", []):
        filename = entry["filename"]
        recorded = entry["sha256"]
        actual_bytes = read_zip_bytes(z, filename)
        actual = sha256_hex(actual_bytes)
        if actual != recorded:
            failures.append(
                f"{filename}: manifest sha256 {recorded} != actual {actual}"
            )
    return failures


def check_chain_rows(z: zipfile.ZipFile, manifest: dict[str, Any]) -> list[str]:
    """For each chain, recompute every row's record_hash and compare."""
    failures: list[str] = []
    genesis = bytes(32)  # 32 zero bytes — the chain seed

    for entry in manifest.get("files", []):
        if entry.get("read_error"):
            # The manifest is honest about a chain it could not read at
            # export time. Surface but don't fail — auditor decides.
            print(
                f"NOTE: {entry['filename']} was not exported "
                f"(read_error: {entry['read_error']})"
            )
            continue

        filename = entry["filename"]
        table = entry["table"]
        text = read_zip_text(z, filename)

        expected_prev = genesis
        expected_seq = 0

        for row in iter_jsonl(text):
            expected_seq += 1
            seq = int(row["sequence_no"])
            if seq != expected_seq:
                failures.append(
                    f"{table}: sequence gap at row id={row.get('id')!r}: "
                    f"expected seq={expected_seq}, got seq={seq}"
                )
                break

            row_prev = bytes.fromhex(row["prev_hash"])
            if row_prev != expected_prev:
                failures.append(
                    f"{table}@seq={seq}: prev_hash chain break "
                    f"(expected {expected_prev.hex()}, got {row['prev_hash']})"
                )
                break

            payload_hex = row.get("__chain_payload_hex")
            if not isinstance(payload_hex, str):
                failures.append(
                    f"{table}@seq={seq}: bundle does not carry "
                    "__chain_payload_hex — bundle predates the offline-verifier "
                    "RPC (migration 20260608000010) or was generated by a "
                    "non-Korvos exporter. Cannot verify per-row hash offline."
                )
                break

            payload = bytes.fromhex(payload_hex)
            recomputed = hashlib.sha256(row_prev + payload).hexdigest()
            if recomputed != row["record_hash"]:
                failures.append(
                    f"{table}@seq={seq} (id={row.get('id')!r}): "
                    f"record_hash mismatch — recomputed {recomputed} != "
                    f"recorded {row['record_hash']}"
                )
                break

            expected_prev = bytes.fromhex(row["record_hash"])

    return failures


def check_chain_heads(z: zipfile.ZipFile, manifest: dict[str, Any]) -> list[str]:
    """The chain head in manifest.files[].head must match the JSONL tail."""
    failures: list[str] = []
    for entry in manifest.get("files", []):
        if entry.get("read_error"):
            continue
        recorded_head = entry.get("head")
        text = read_zip_text(z, entry["filename"])
        rows = list(iter_jsonl(text))
        if not rows:
            if recorded_head is not None:
                failures.append(
                    f"{entry['table']}: manifest records a head "
                    f"({recorded_head}) but the JSONL is empty"
                )
            continue
        tail = rows[-1]
        if recorded_head is None:
            failures.append(
                f"{entry['table']}: JSONL has rows but manifest head is null"
            )
            continue
        if int(recorded_head["sequence_no"]) != int(tail["sequence_no"]):
            failures.append(
                f"{entry['table']}: head sequence_no "
                f"{recorded_head['sequence_no']} != tail {tail['sequence_no']}"
            )
        if recorded_head["record_hash"] != tail["record_hash"]:
            failures.append(
                f"{entry['table']}: head record_hash "
                f"{recorded_head['record_hash']} != tail {tail['record_hash']}"
            )
    return failures


# ---------------------------------------------------------------------------
# Entry
# ---------------------------------------------------------------------------


def main(argv: list[str]) -> int:
    if len(argv) != 2:
        print("Usage: python3 verify.py <bundle.zip>", file=sys.stderr)
        return 3

    bundle_path = argv[1]
    try:
        z = zipfile.ZipFile(bundle_path, "r")
    except (FileNotFoundError, zipfile.BadZipFile) as exc:
        print(f"[3] cannot open bundle: {exc}", file=sys.stderr)
        return 3

    with z:
        manifest_text = read_zip_text(z, "manifest.json")
        try:
            manifest = json.loads(manifest_text)
        except json.JSONDecodeError as exc:
            print(f"[3] manifest.json is not valid JSON: {exc}", file=sys.stderr)
            return 3

        if manifest.get("format_version") != 1:
            print(
                f"[3] unsupported manifest.format_version: "
                f"{manifest.get('format_version')!r} (this script expects 1)",
                file=sys.stderr,
            )
            return 3

        print(
            f"Bundle: {bundle_path}\n"
            f"Tenant: {manifest.get('tenant', {}).get('legal_name')} "
            f"({manifest.get('tenant', {}).get('id')})\n"
            f"Range:  {manifest.get('range', {}).get('from')} -> "
            f"{manifest.get('range', {}).get('to')}\n"
            f"Generated: {manifest.get('generated_at')}\n"
            f"Chains:    {len(manifest.get('files', []))}\n"
        )

        file_failures = check_file_hashes(z, manifest)
        if file_failures:
            print("FILE-HASH FAILURES:", file=sys.stderr)
            for f in file_failures:
                print(f"  - {f}", file=sys.stderr)
            return 1
        print("[OK] all chain files match their manifest SHA-256")

        row_failures = check_chain_rows(z, manifest)
        if row_failures:
            print("CHAIN-ROW FAILURES:", file=sys.stderr)
            for f in row_failures:
                print(f"  - {f}", file=sys.stderr)
            return 2
        print("[OK] every row's record_hash matches sha256(prev_hash || payload)")

        head_failures = check_chain_heads(z, manifest)
        if head_failures:
            print("CHAIN-HEAD FAILURES:", file=sys.stderr)
            for f in head_failures:
                print(f"  - {f}", file=sys.stderr)
            return 4
        print("[OK] every chain head matches its JSONL tail")

        print("\nIntegrity verified. All checks passed.")
        return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))
