#!/usr/bin/env python3
"""
EXPLAIN ANALYZE helper for RLS verification.

NOTE: As of Alembic 0030 (`rls_via_pg_session_jwt`), tenant RLS policies
key on `auth.session() ->> 'activeRestaurantId'`, not
`current_setting('app.tenant_id', …)`. This script still uses the legacy
GUC for benchmarking convenience because reproducing the BA JWT path
locally is heavy — `set_config('app.tenant_id', …)` is a no-op against
production policies but lets us compare plans for "with tenant key
filter" vs "without" without rebuilding pg_session_jwt state per query.
Treat the output as a plan-shape comparison only, not a security check.

Usage:
    # Environment:
    #   NEON_DATABASE_URL_DIRECT: Direct DSN to the Neon branch (required)
    #   TEST_TENANT_ID: Restaurant/tenant id to set via set_config (required)
    #   RUN_RLS_TESTS=1 to mirror CI gating (optional)
    #
    # Single query:
    #   python backend/scripts/explain_bench.py \
    #     "SELECT id FROM reservation WHERE restaurant_id =" \
    #     "  current_setting('app.tenant_id', true)" \
    #     "LIMIT 10"
    #
    # Multiple queries:
    #   python backend/scripts/explain_bench.py \
    #     "SELECT id FROM reservation ORDER BY created_at DESC LIMIT 50" \
    #     "SELECT id FROM common_question WHERE restaurant_id =" \
    #     "  current_setting('app.tenant_id', true)" \
    #     "  ORDER BY created_at DESC LIMIT 50"

Notes:
- Runs each query twice:
  1) Without tenant context (no set_config)
  2) With tenant context injected via set_config('app.tenant_id', ...)
- Prints plans in a diff-friendly format. Use a real tenant id that has data.
"""

from __future__ import annotations

import asyncio
import os
import sys
from collections.abc import Sequence

import asyncpg


def _fmt_title(title: str) -> str:
    bar = "=" * len(title)
    return f"{title}\n{bar}"


async def _explain(conn: asyncpg.Connection, sql: str) -> list[str]:
    # EXPLAIN options tuned for stable output across runs
    plan_rows = await conn.fetch("EXPLAIN (ANALYZE, COSTS, VERBOSE, BUFFERS, TIMING OFF) " + sql)
    return [r[0] for r in plan_rows]


async def main(argv: Sequence[str]) -> int:
    dsn = os.getenv("NEON_DATABASE_URL_DIRECT")
    rid = os.getenv("TEST_TENANT_ID")
    if not dsn or not rid:
        print("ERROR: NEON_DATABASE_URL_DIRECT and TEST_TENANT_ID must be set", file=sys.stderr)
        return 2

    if len(argv) < 2:
        print("Usage: explain_bench.py <SQL> [<SQL> ...]", file=sys.stderr)
        return 2

    queries = list(argv[1:])

    conn = await asyncpg.connect(dsn)
    try:
        for idx, sql in enumerate(queries, 1):
            title = f"Query {idx}: {sql.strip()}"
            print(_fmt_title(title))

            print("-- Without tenant context --")
            plan_no_tenant = await _explain(conn, sql)
            for line in plan_no_tenant:
                print(line)

            print("\n-- With tenant context (set_config app.tenant_id) --")
            await conn.execute("SELECT set_config('app.tenant_id', $1, true)", rid)
            plan_with_tenant = await _explain(conn, sql)
            for line in plan_with_tenant:
                print(line)

            # Reset for next iteration
            await conn.execute("RESET ALL")
            print()
        return 0
    finally:
        await conn.close()


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