#!/usr/bin/env python3
"""
Security policy (RAP / Data Mask) benchmark.

Measures performance overhead of row access policies and data masking policies
compared to equivalent manual WHERE clauses and function calls.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path
from typing import Any

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from scripts.bench_common.bendsql import BendSQL, QueryResult, csv_records, kv_map


# ---------------------------------------------------------------------------
# Policy and query definitions
# ---------------------------------------------------------------------------

RAP_POLICIES: dict[str, dict[str, str]] = {
    "simple": {
        "body": "(tenant_id INT, status VARCHAR) RETURNS BOOLEAN -> tenant_id = 1",
        "columns": "(tenant_id, status)",
    },
    "range": {
        "body": "(tenant_id INT, status VARCHAR) RETURNS BOOLEAN -> tenant_id BETWEEN 10 AND 30",
        "columns": "(tenant_id, status)",
    },
    "complex": {
        "body": "(tenant_id INT, status VARCHAR) RETURNS BOOLEAN -> tenant_id % 10 = 0 AND status = 'active'",
        "columns": "(tenant_id, status)",
    },
}

MASK_POLICIES: dict[str, dict[str, str]] = {
    "simple": {
        "body": "(val STRING) RETURNS STRING -> '***'",
        "column": "email",
    },
    "range": {
        "body": "(val STRING) RETURNS STRING -> CONCAT(LEFT(val, 3), '***@***.com')",
        "column": "email",
    },
    "complex": {
        "body": "(val STRING) RETURNS STRING -> CASE WHEN LENGTH(val) > 5 THEN CONCAT(LEFT(val, 2), REPEAT('*', LENGTH(val) - 4), RIGHT(val, 2)) ELSE '***' END",
        "column": "email",
    },
}

def rap_where_equivalent(complexity: str) -> str:
    mapping = {
        "simple": "tenant_id = 1",
        "range": "tenant_id BETWEEN 10 AND 30",
        "complex": "tenant_id % 10 = 0 AND status = 'active'",
    }
    return mapping[complexity]


def mask_func_equivalent(complexity: str) -> str:
    mapping = {
        "simple": "'***' AS email",
        "range": "CONCAT(LEFT(email, 3), '***@***.com') AS email",
        "complex": "CASE WHEN LENGTH(email) > 5 THEN CONCAT(LEFT(email, 2), REPEAT('*', LENGTH(email) - 4), RIGHT(email, 2)) ELSE '***' END AS email",
    }
    return mapping[complexity]


def query_patterns(table: str, complexity: str) -> list[dict[str, str]]:
    where_clause = rap_where_equivalent(complexity)
    return [
        {
            "name": "point",
            "rap": f"SELECT id, email, amount FROM {table} LIMIT 50000",
            "where": f"SELECT id, email, amount FROM {table} WHERE {where_clause} LIMIT 50000",
        },
        {
            "name": "range_scan",
            "rap": f"SELECT id, email FROM {table}",
            "where": f"SELECT id, email FROM {table} WHERE {where_clause}",
        },
        {
            "name": "order_limit",
            "rap": f"SELECT id, email, amount FROM {table} ORDER BY created_at DESC LIMIT 100",
            "where": f"SELECT id, email, amount FROM {table} WHERE {where_clause} ORDER BY created_at DESC LIMIT 100",
        },
        {
            "name": "agg",
            "rap": f"SELECT COUNT(*), SUM(amount) FROM {table}",
            "where": f"SELECT COUNT(*), SUM(amount) FROM {table} WHERE {where_clause}",
        },
    ]


def mask_query_patterns(table: str, complexity: str) -> list[dict[str, str]]:
    func_expr = mask_func_equivalent(complexity)
    return [
        {
            "name": "full_scan",
            "mask": f"SELECT email, id FROM {table} LIMIT 50000",
            "func": f"SELECT {func_expr}, id FROM {table} LIMIT 50000",
        },
        {
            "name": "full_scan_large",
            "mask": f"SELECT email FROM {table}",
            "func": f"SELECT {func_expr} FROM {table}",
        },
    ]


# ---------------------------------------------------------------------------
# Metrics collection
# ---------------------------------------------------------------------------

def collect_query_metrics(bs: BendSQL, query_id: str) -> dict[str, Any]:
    sql = f"""
SELECT query_duration_ms, scan_rows, scan_bytes, scan_partitions,
       total_partitions, result_rows, result_bytes
FROM system.query_log
WHERE query_id = '{query_id}' AND exception_code = 0
ORDER BY event_time DESC LIMIT 1
"""
    res = bs.try_run(sql, f"metrics-{query_id[:8]}", echo_stdout=False)
    rows = csv_records(res.stdout)
    if not rows:
        return {}
    r = rows[0]
    if len(r) < 7:
        return {}
    return {
        "duration_ms": float(r[0]),
        "scan_rows": int(r[1]),
        "scan_bytes": int(r[2]),
        "scan_partitions": int(r[3]),
        "total_partitions": int(r[4]),
        "result_rows": int(r[5]),
        "result_bytes": int(r[6]),
    }


def parse_pruning_stats(explain_output: str) -> dict[str, Any]:
    stats: dict[str, Any] = {}
    seg_match = __import__("re").search(
        r"segments:\s*<range pruning:\s*(\d+)\s*to\s*(\d+)", explain_output
    )
    if seg_match:
        before, after = int(seg_match.group(1)), int(seg_match.group(2))
        stats["segments_before"] = before
        stats["segments_after"] = after
        stats["segments_pruned_pct"] = round((before - after) / before * 100, 2) if before > 0 else 0

    block_match = __import__("re").search(
        r"blocks:\s*<\s*range pruning:\s*(\d+)\s*to\s*(\d+)", explain_output
    )
    if block_match:
        before, after = int(block_match.group(1)), int(block_match.group(2))
        stats["blocks_before"] = before
        stats["blocks_after"] = after
        stats["blocks_pruned_pct"] = round((before - after) / before * 100, 2) if before > 0 else 0

    return stats


def collect_cache_stats(bs: BendSQL) -> dict[str, dict[str, int]]:
    res = bs.try_run(
        "SELECT name, num_items, size, access, hit, miss FROM system.caches",
        "caches", echo_stdout=False,
    )
    out: dict[str, dict[str, int]] = {}
    for row in csv_records(res.stdout):
        if len(row) >= 6:
            out[row[0]] = {
                "num_items": int(row[1]),
                "size": int(row[2]),
                "access": int(row[3]),
                "hit": int(row[4]),
                "miss": int(row[5]),
            }
    return out


# ---------------------------------------------------------------------------
# Data generation
# ---------------------------------------------------------------------------

def create_table_sql(database: str, table: str, row_per_block: int, block_per_segment: int) -> str:
    return f"""
CREATE TABLE IF NOT EXISTS {database}.{table} (
    id         INT NOT NULL,
    tenant_id  INT NOT NULL,
    email      VARCHAR NOT NULL,
    amount     DECIMAL(18,2) NOT NULL,
    status     VARCHAR NOT NULL,
    created_at TIMESTAMP NOT NULL
)
CLUSTER BY (created_at)
ROW_PER_BLOCK = {row_per_block}
BLOCK_PER_SEGMENT = {block_per_segment}
"""


def insert_sql(database: str, table: str, batch: int, rows_per_batch: int) -> str:
    offset = batch * rows_per_batch
    return f"""
INSERT INTO {database}.{table}
SELECT
    (number + {offset})::INT AS id,
    (number % 100)::INT AS tenant_id,
    CONCAT('user', to_string(number + {offset}), '@example.com') AS email,
    ((number % 10000) * 1.5)::DECIMAL(18,2) AS amount,
    CASE WHEN number % 3 = 0 THEN 'active' WHEN number % 3 = 1 THEN 'inactive' ELSE 'pending' END AS status,
    to_timestamp({offset} + number) AS created_at
FROM numbers({rows_per_batch})
"""


# ---------------------------------------------------------------------------
# Benchmark execution
# ---------------------------------------------------------------------------

def run_query_timed(bs: BendSQL, sql: str, label: str, rounds: int) -> tuple[list[float], list[float]]:
    """Returns (client_times_sec, server_times_sec)."""
    client_times = []
    server_times = []
    for i in range(rounds):
        res = bs.run(sql, f"{label}-r{i}", output="null", echo_stdout=False)
        client_times.append(res.elapsed_sec)
        # Get server-side time via --time=server
        if not bs.dry_run:
            srv = bs.run_time_server(sql, f"{label}-srv-r{i}")
            if srv is not None:
                server_times.append(srv)
    return client_times, server_times


def bench_rap(
    bs: BendSQL,
    database: str,
    table: str,
    complexity: str,
    warmup: int,
    rounds: int,
) -> list[dict[str, Any]]:
    full_table = f"{database}.{table}"
    policy_name = f"bench_rap_{complexity}"
    policy_def = RAP_POLICIES[complexity]
    results = []

    patterns = query_patterns(full_table, complexity)
    for pat in patterns:
        # Mode: where (baseline)
        run_query_timed(bs, pat["where"], f"warmup-where-{complexity}-{pat['name']}", warmup)
        client_times, server_times = run_query_timed(bs, pat["where"], f"where-{complexity}-{pat['name']}", rounds)
        metrics: dict[str, Any] = {
            "duration_ms": round(min(client_times) * 1000, 2),
            "duration_avg_ms": round(sum(client_times) / len(client_times) * 1000, 2),
            "all_ms": [round(t * 1000, 2) for t in client_times],
        }
        if server_times:
            metrics["server_min_ms"] = round(min(server_times) * 1000, 2)
            metrics["server_avg_ms"] = round(sum(server_times) / len(server_times) * 1000, 2)
        results.append({
            "complexity": complexity,
            "pattern": pat["name"],
            "mode": "where",
            "cache": "warm",
            "metrics": metrics,
        })

        # Mode: rap
        bs.try_run(
            f"DROP ROW ACCESS POLICY IF EXISTS {policy_name}",
            f"drop-rap-pre-{complexity}", output="null",
        )
        bs.run(
            f"CREATE ROW ACCESS POLICY {policy_name} AS {policy_def['body']}",
            f"create-rap-{complexity}", output="null",
        )
        bs.run(
            f"ALTER TABLE {full_table} ADD ROW ACCESS POLICY {policy_name} ON {policy_def['columns']}",
            f"attach-rap-{complexity}", output="null",
        )

        run_query_timed(bs, pat["rap"], f"warmup-rap-{complexity}-{pat['name']}", warmup)
        client_times, server_times = run_query_timed(bs, pat["rap"], f"rap-{complexity}-{pat['name']}", rounds)
        metrics = {
            "duration_ms": round(min(client_times) * 1000, 2),
            "duration_avg_ms": round(sum(client_times) / len(client_times) * 1000, 2),
            "all_ms": [round(t * 1000, 2) for t in client_times],
        }
        if server_times:
            metrics["server_min_ms"] = round(min(server_times) * 1000, 2)
            metrics["server_avg_ms"] = round(sum(server_times) / len(server_times) * 1000, 2)
        results.append({
            "complexity": complexity,
            "pattern": pat["name"],
            "mode": "rap",
            "cache": "warm",
            "metrics": metrics,
        })

        # Collect EXPLAIN pruning stats for the rap query
        explain_res = bs.try_run(
            f"EXPLAIN ANALYZE {pat['rap']}", f"explain-rap-{complexity}-{pat['name']}",
            echo_stdout=False,
        )
        pruning = parse_pruning_stats(explain_res.stdout)
        if pruning:
            results[-1]["pruning"] = pruning

        bs.run(
            f"ALTER TABLE {full_table} DROP ROW ACCESS POLICY {policy_name}",
            f"drop-rap-{complexity}-{pat['name']}", output="null",
        )

    return results


def bench_mask(
    bs: BendSQL,
    database: str,
    table: str,
    complexity: str,
    warmup: int,
    rounds: int,
) -> list[dict[str, Any]]:
    full_table = f"{database}.{table}"
    policy_name = f"bench_mask_{complexity}"
    mask_def = MASK_POLICIES[complexity]
    results = []

    patterns = mask_query_patterns(full_table, complexity)
    for pat in patterns:
        # Mode: func (baseline)
        run_query_timed(bs, pat["func"], f"warmup-func-{complexity}-{pat['name']}", warmup)
        client_times, server_times = run_query_timed(bs, pat["func"], f"func-{complexity}-{pat['name']}", rounds)
        metrics: dict[str, Any] = {
            "duration_ms": round(min(client_times) * 1000, 2),
            "duration_avg_ms": round(sum(client_times) / len(client_times) * 1000, 2),
            "all_ms": [round(t * 1000, 2) for t in client_times],
        }
        if server_times:
            metrics["server_min_ms"] = round(min(server_times) * 1000, 2)
            metrics["server_avg_ms"] = round(sum(server_times) / len(server_times) * 1000, 2)
        results.append({
            "complexity": complexity,
            "pattern": pat["name"],
            "mode": "func",
            "cache": "warm",
            "metrics": metrics,
        })

        # Mode: mask
        bs.try_run(
            f"DROP MASKING POLICY IF EXISTS {policy_name}",
            f"drop-mask-pre-{complexity}", output="null",
        )
        bs.run(
            f"CREATE MASKING POLICY {policy_name} AS {mask_def['body']}",
            f"create-mask-{complexity}", output="null",
        )
        bs.run(
            f"ALTER TABLE {full_table} MODIFY COLUMN {mask_def['column']} SET MASKING POLICY {policy_name}",
            f"attach-mask-{complexity}", output="null",
        )

        run_query_timed(bs, pat["mask"], f"warmup-mask-{complexity}-{pat['name']}", warmup)
        client_times, server_times = run_query_timed(bs, pat["mask"], f"mask-{complexity}-{pat['name']}", rounds)
        metrics = {
            "duration_ms": round(min(client_times) * 1000, 2),
            "duration_avg_ms": round(sum(client_times) / len(client_times) * 1000, 2),
            "all_ms": [round(t * 1000, 2) for t in client_times],
        }
        if server_times:
            metrics["server_min_ms"] = round(min(server_times) * 1000, 2)
            metrics["server_avg_ms"] = round(sum(server_times) / len(server_times) * 1000, 2)
        results.append({
            "complexity": complexity,
            "pattern": pat["name"],
            "mode": "mask",
            "cache": "warm",
            "metrics": metrics,
        })

        bs.run(
            f"ALTER TABLE {full_table} MODIFY COLUMN {mask_def['column']} UNSET MASKING POLICY",
            f"drop-mask-{complexity}-{pat['name']}", output="null",
        )

    return results


def bench_planner_cache(
    bs: BendSQL,
    database: str,
    table: str,
    complexity: str,
    rounds: int,
) -> list[dict[str, Any]]:
    full_table = f"{database}.{table}"
    policy_name = f"bench_rap_cache_{complexity}"
    policy_def = RAP_POLICIES[complexity]
    results = []

    test_sql = f"SELECT COUNT(*) FROM {full_table}"

    # Without policy: cold vs warm (baseline)
    bs.run("SET enable_planner_cache = 1", "cache-enable", output="null")

    cold = bs.run(test_sql, "cache-cold-no-policy", output="null", echo_stdout=False)
    warm_times, _ = run_query_timed(bs, test_sql, "cache-warm-no-policy", rounds)
    results.append({
        "complexity": complexity,
        "pattern": "planner_cache",
        "mode": "where",
        "cache": "cold",
        "metrics": {"duration_ms": round(cold.elapsed_sec * 1000, 2)},
    })
    results.append({
        "complexity": complexity,
        "pattern": "planner_cache",
        "mode": "where",
        "cache": "warm",
        "metrics": {
            "duration_ms": round(min(warm_times) * 1000, 2),
            "duration_avg_ms": round(sum(warm_times) / len(warm_times) * 1000, 2),
        },
    })

    # With RAP policy: 3-layer cache isolation
    # Layer 1: both_cold - first query after fresh policy attach (policy cache miss + planner cache miss)
    # Layer 2: policy_warm_planner_cold - different SQL text forces planner miss, policy cache hit
    # Layer 3: both_warm - same SQL repeated (both caches hit)

    # Fresh policy: drop and recreate to ensure policy cache miss
    bs.try_run(
        f"DROP ROW ACCESS POLICY IF EXISTS {policy_name}",
        "cache-drop-rap-pre", output="null",
    )
    bs.run(
        f"CREATE ROW ACCESS POLICY {policy_name} AS {policy_def['body']}",
        "cache-create-rap", output="null",
    )
    bs.run(
        f"ALTER TABLE {full_table} ADD ROW ACCESS POLICY {policy_name} ON {policy_def['columns']}",
        "cache-attach-rap", output="null",
    )

    # Layer 1: both_cold (first query - policy cache miss + planner cache miss)
    # Use SQL variant A
    sql_a = f"SELECT COUNT(*) FROM {full_table}"
    both_cold = bs.run(sql_a, "cache-both-cold", output="null", echo_stdout=False)
    results.append({
        "complexity": complexity,
        "pattern": "planner_cache",
        "mode": "rap",
        "cache": "both_cold",
        "metrics": {"duration_ms": round(both_cold.elapsed_sec * 1000, 2)},
    })

    # Layer 2: policy_warm_planner_cold
    # Different SQL text → planner cache miss; same policy_id → policy cache hit
    sql_b = f"SELECT COUNT(1) FROM {full_table}"
    policy_warm = bs.run(sql_b, "cache-policy-warm-planner-cold", output="null", echo_stdout=False)
    results.append({
        "complexity": complexity,
        "pattern": "planner_cache",
        "mode": "rap",
        "cache": "policy_warm",
        "metrics": {"duration_ms": round(policy_warm.elapsed_sec * 1000, 2)},
    })

    # Layer 3: both_warm (repeat sql_a - both caches hit from Layer 1)
    warm_times, _ = run_query_timed(bs, sql_a, "cache-both-warm", rounds)
    results.append({
        "complexity": complexity,
        "pattern": "planner_cache",
        "mode": "rap",
        "cache": "both_warm",
        "metrics": {
            "duration_ms": round(min(warm_times) * 1000, 2),
            "duration_avg_ms": round(sum(warm_times) / len(warm_times) * 1000, 2),
        },
    })

    bs.run(
        f"ALTER TABLE {full_table} DROP ROW ACCESS POLICY {policy_name}",
        "cache-drop-rap", output="null",
    )

    return results


# ---------------------------------------------------------------------------
# Summary computation
# ---------------------------------------------------------------------------

def compute_summary(results: list[dict[str, Any]]) -> dict[str, Any]:
    rap_vs_where: list[dict[str, Any]] = []
    mask_vs_func: list[dict[str, Any]] = []
    cache_results: list[dict[str, Any]] = []

    by_key: dict[tuple, dict] = {}
    for r in results:
        key = (r["complexity"], r["pattern"], r["mode"], r["cache"])
        by_key[key] = r

    complexities = {r["complexity"] for r in results}
    patterns = {r["pattern"] for r in results}

    for c in sorted(complexities):
        for p in sorted(patterns):
            if p == "planner_cache":
                # Planner cache: output 3-layer comparison
                for mode in ("where",):
                    cold_r = by_key.get((c, p, mode, "cold"))
                    warm_r = by_key.get((c, p, mode, "warm"))
                    if cold_r and warm_r:
                        cold_ms = cold_r["metrics"]["duration_ms"]
                        warm_ms = warm_r["metrics"]["duration_ms"]
                        speedup = ((cold_ms - warm_ms) / cold_ms * 100) if cold_ms > 0 else 0
                        cache_results.append({
                            "complexity": c, "mode": mode, "layer": "planner",
                            "cold_ms": cold_ms, "warm_ms": warm_ms,
                            "speedup_pct": round(speedup, 2),
                        })

                both_cold_r = by_key.get((c, p, "rap", "both_cold"))
                policy_warm_r = by_key.get((c, p, "rap", "policy_warm"))
                both_warm_r = by_key.get((c, p, "rap", "both_warm"))
                if both_cold_r and policy_warm_r and both_warm_r:
                    both_cold_ms = both_cold_r["metrics"]["duration_ms"]
                    policy_warm_ms = policy_warm_r["metrics"]["duration_ms"]
                    both_warm_ms = both_warm_r["metrics"]["duration_ms"]
                    meta_cost = round(both_cold_ms - policy_warm_ms, 2)
                    planning_cost = round(policy_warm_ms - both_warm_ms, 2)
                    cache_results.append({
                        "complexity": c, "mode": "rap", "layer": "3-layer",
                        "both_cold_ms": both_cold_ms,
                        "policy_warm_ms": policy_warm_ms,
                        "both_warm_ms": both_warm_ms,
                        "meta_rpc_cost_ms": meta_cost,
                        "planning_cost_ms": planning_cost,
                    })
                continue

            rap_r = by_key.get((c, p, "rap", "warm"))
            where_r = by_key.get((c, p, "where", "warm"))
            if rap_r and where_r:
                rap_ms = rap_r["metrics"].get("server_min_ms", rap_r["metrics"]["duration_ms"])
                where_ms = where_r["metrics"].get("server_min_ms", where_r["metrics"]["duration_ms"])
                overhead = ((rap_ms - where_ms) / where_ms * 100) if where_ms > 0 else 0
                rap_vs_where.append({
                    "complexity": c, "pattern": p,
                    "rap_ms": rap_ms, "where_ms": where_ms,
                    "overhead_pct": round(overhead, 2),
                    "source": "server" if "server_min_ms" in rap_r["metrics"] else "client",
                })

            mask_r = by_key.get((c, p, "mask", "warm"))
            func_r = by_key.get((c, p, "func", "warm"))
            if mask_r and func_r:
                mask_ms = mask_r["metrics"].get("server_min_ms", mask_r["metrics"]["duration_ms"])
                func_ms = func_r["metrics"].get("server_min_ms", func_r["metrics"]["duration_ms"])
                overhead = ((mask_ms - func_ms) / func_ms * 100) if func_ms > 0 else 0
                mask_vs_func.append({
                    "complexity": c, "pattern": p,
                    "mask_ms": mask_ms, "func_ms": func_ms,
                    "overhead_pct": round(overhead, 2),
                    "source": "server" if "server_min_ms" in mask_r["metrics"] else "client",
                })

    return {
        "rap_vs_where": rap_vs_where,
        "mask_vs_func": mask_vs_func,
        "planner_cache": cache_results,
    }


# ---------------------------------------------------------------------------
# CLI and main
# ---------------------------------------------------------------------------

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Benchmark RAP and Data Mask policy overhead."
    )
    parser.add_argument("--dsn", default=os.getenv("BENDSQL_DSN"), help="bendsql DSN")
    parser.add_argument("--bendsql", default=os.getenv("BENDSQL", "bendsql"))
    parser.add_argument("--database", default="tmp_security_policy_bench")
    parser.add_argument("--table", default="bench_security")
    parser.add_argument("--rows", type=int, default=5_000_000)
    parser.add_argument("--rows-per-batch", type=int, default=1_000_000)
    parser.add_argument("--rows-per-block", type=int, default=10_000)
    parser.add_argument("--blocks-per-segment", type=int, default=10)
    parser.add_argument(
        "--complexity",
        choices=["simple", "range", "complex", "all"],
        default="all",
    )
    parser.add_argument(
        "--pattern",
        choices=["point", "range_scan", "order_limit", "agg", "full_scan", "all"],
        default="all",
    )
    parser.add_argument("--warmup-rounds", type=int, default=2)
    parser.add_argument("--bench-rounds", type=int, default=5)
    parser.add_argument("--output", default=None, help="Write JSON result to this file")
    parser.add_argument("--keep", action="store_true", help="Do not drop database first")
    parser.add_argument("--drop-after", action="store_true", help="Drop database after run")
    parser.add_argument("--dry-run", action="store_true")
    return parser.parse_args()


def print_summary(summary: dict[str, Any]) -> None:
    print("\n" + "=" * 70)
    print("SUMMARY: RAP vs WHERE (overhead %)")
    print("-" * 70)
    for item in summary.get("rap_vs_where", []):
        flag = "!!!" if item["overhead_pct"] > 5 else "   "
        print(f"  {flag} {item['complexity']:8s} | {item['pattern']:12s} | "
              f"rap={item['rap_ms']:8.1f}ms  where={item['where_ms']:8.1f}ms  "
              f"overhead={item['overhead_pct']:+.1f}%")

    print("\nSUMMARY: MASK vs FUNC (overhead %)")
    print("-" * 70)
    for item in summary.get("mask_vs_func", []):
        flag = "!!!" if item["overhead_pct"] > 5 else "   "
        print(f"  {flag} {item['complexity']:8s} | {item['pattern']:12s} | "
              f"mask={item['mask_ms']:8.1f}ms  func={item['func_ms']:8.1f}ms  "
              f"overhead={item['overhead_pct']:+.1f}%")

    print("\nSUMMARY: CACHE LAYERS (cost isolation)")
    print("-" * 70)
    for item in summary.get("planner_cache", []):
        if item.get("layer") == "planner":
            print(f"       {item['complexity']:8s} | {item['mode']:5s} | "
                  f"cold={item['cold_ms']:8.1f}ms  warm={item['warm_ms']:8.1f}ms  "
                  f"speedup={item['speedup_pct']:+.1f}%")
        elif item.get("layer") == "3-layer":
            print(f"       {item['complexity']:8s} | rap   | "
                  f"both_cold={item['both_cold_ms']:.1f}ms  "
                  f"policy_warm={item['policy_warm_ms']:.1f}ms  "
                  f"both_warm={item['both_warm_ms']:.1f}ms")
            print(f"       {'':8s} |       | "
                  f"meta_rpc={item['meta_rpc_cost_ms']:+.1f}ms  "
                  f"planning={item['planning_cost_ms']:+.1f}ms")
    print("=" * 70)


def main() -> int:
    args = parse_args()
    bs = BendSQL(args.bendsql, args.dsn, args.dry_run)

    complexities = ["simple", "range", "complex"] if args.complexity == "all" else [args.complexity]

    # Setup
    if not args.keep:
        bs.try_run(f"DROP DATABASE IF EXISTS {args.database}", "drop-db", output="null")
    bs.run(f"CREATE DATABASE IF NOT EXISTS {args.database}", "create-db", output="null")
    bs.run(
        create_table_sql(args.database, args.table, args.rows_per_block, args.blocks_per_segment),
        "create-table", output="null",
    )

    # Insert data (skip if --keep and table already has data)
    if not args.keep:
        batches = (args.rows + args.rows_per_batch - 1) // args.rows_per_batch
        print(f"\nInserting {args.rows} rows in {batches} batches...", flush=True)
        for i in range(batches):
            batch_rows = min(args.rows_per_batch, args.rows - i * args.rows_per_batch)
            bs.run(
                insert_sql(args.database, args.table, i, batch_rows),
                f"insert-batch-{i}", output="null",
            )

    # Verify data setup
    full_table = f"{args.database}.{args.table}"
    count_res = bs.try_run(f"SELECT COUNT(*) FROM {full_table}", "count-rows", echo_stdout=False)
    snapshot_res = bs.try_run(
        f"SELECT segment_count, block_count, row_count FROM fuse_snapshot('{args.database}', '{args.table}') LIMIT 1",
        "snapshot-info", echo_stdout=False,
    )
    data_setup: dict[str, Any] = {"rows": args.rows}
    snap_rows = csv_records(snapshot_res.stdout)
    if snap_rows and len(snap_rows[0]) >= 3:
        data_setup["segments"] = int(snap_rows[0][0])
        data_setup["blocks"] = int(snap_rows[0][1])
        data_setup["actual_rows"] = int(snap_rows[0][2])
    print(f"Data setup: {data_setup}", flush=True)

    # Switch to the benchmark database for policy DDL (policies don't support db.name syntax)
    bs.run(f"USE {args.database}", "use-db", output="null")
    all_results: list[dict[str, Any]] = []

    for complexity in complexities:
        print(f"\n{'='*40} RAP complexity={complexity} {'='*40}", flush=True)
        all_results.extend(bench_rap(bs, args.database, args.table, complexity, args.warmup_rounds, args.bench_rounds))

        print(f"\n{'='*40} MASK complexity={complexity} {'='*40}", flush=True)
        all_results.extend(bench_mask(bs, args.database, args.table, complexity, args.warmup_rounds, args.bench_rounds))

        print(f"\n{'='*40} CACHE complexity={complexity} {'='*40}", flush=True)
        all_results.extend(bench_planner_cache(bs, args.database, args.table, complexity, args.bench_rounds))

    # Filter by pattern if specified
    if args.pattern != "all":
        all_results = [r for r in all_results if r["pattern"] == args.pattern]

    # Compute summary
    summary = compute_summary(all_results)
    print_summary(summary)

    # Output
    output = {
        "args": vars(args),
        "data_setup": data_setup,
        "results": all_results,
        "summary": summary,
    }

    if args.output:
        Path(args.output).parent.mkdir(parents=True, exist_ok=True)
        with open(args.output, "w", encoding="utf-8") as f:
            json.dump(output, f, indent=2)
        print(f"\nResults written to: {args.output}")

    if args.drop_after:
        bs.try_run(f"DROP DATABASE IF EXISTS {args.database}", "drop-db-after", output="null")

    return 0


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