#!/usr/bin/env python3
"""Policy-overhead benchmark — measures the p50/p95/p99 latency Gateco's
policy layer adds on top of the raw vector-DB leg.

Methodology (matches the published description on /docs/performance):
  - pgvector connector against a local table of --vectors embeddings (1536-dim)
  - --resources registered GatedResources with classification/sensitivity/
    domain metadata, 10 chunks each (sidecar metadata resolution, the default)
  - one ACTIVE 5-rule RBAC policy (allow-by-group, classification ceiling,
    restricted deny, department/domain allow, critical-sensitivity deny)
  - one active principal with groups, resolved per request
  - requests supply a precomputed query_vector so server-side embedding is
    excluded — the measured overhead is principal load + policy evaluation +
    sidecar registry resolution + audit/usage writes + response assembly

Two overhead series are reported per request:
  service_overhead = SecuredRetrieval.latency_ms - connector_latency_ms
      (inside the retrieval service, excludes HTTP middleware)
  full_overhead    = client wall-clock ms - connector_latency_ms
      (everything Gateco adds, including auth middleware + serialization)

Usage (backend running on --base-url with DEBUG=true and ADMIN_SETUP_ENABLED=true):
    cd gateco/apps/backend && python scripts/policy_overhead_bench.py \
        --requests 300 --json-out /tmp/bench.json
"""

from __future__ import annotations

import argparse
import asyncio
import json
import math
import os
import random
import secrets
import statistics
import sys
import time
import uuid
from pathlib import Path

import asyncpg
import httpx

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))

BENCH_TABLE = "bench_vectors"
EMBED_DIM = 1536
CHUNKS_PER_RESOURCE = 10

# Required: the admin token of the target environment (used once, to flip the
# bench org's plan tier). No default on purpose — set ADMIN_TOKEN explicitly.
ADMIN_TOKEN = os.environ.get("ADMIN_TOKEN", "")


def _percentile(values: list[float], pct: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(values)
    k = (len(ordered) - 1) * pct / 100.0
    lo, hi = math.floor(k), math.ceil(k)
    if lo == hi:
        return ordered[lo]
    return ordered[lo] + (ordered[hi] - ordered[lo]) * (k - lo)


def _stats(name: str, values: list[float]) -> dict:
    return {
        "series": name,
        "n": len(values),
        "p50_ms": round(_percentile(values, 50), 2),
        "p95_ms": round(_percentile(values, 95), 2),
        "p99_ms": round(_percentile(values, 99), 2),
        "mean_ms": round(statistics.fmean(values), 2) if values else 0.0,
        "max_ms": round(max(values), 2) if values else 0.0,
    }


def _unit_vector(rng: random.Random) -> list[float]:
    v = [rng.gauss(0, 1) for _ in range(EMBED_DIM)]
    norm = math.sqrt(sum(x * x for x in v)) or 1.0
    return [x / norm for x in v]


def _vec_literal(v: list[float]) -> str:
    return "[" + ",".join(f"{x:.6f}" for x in v) + "]"


async def _setup_bench_table(dsn: str, n_vectors: int, rng: random.Random) -> list[str]:
    conn = await asyncpg.connect(dsn)
    try:
        await conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
        await conn.execute(f"DROP TABLE IF EXISTS {BENCH_TABLE}")
        await conn.execute(
            f"CREATE TABLE {BENCH_TABLE} ("
            f"id text PRIMARY KEY, embedding vector({EMBED_DIM}), content text)"
        )
        vector_ids = [f"bench-v-{i:05d}" for i in range(n_vectors)]
        rows = [
            (vid, _vec_literal(_unit_vector(rng)), f"bench content {vid}")
            for vid in vector_ids
        ]
        await conn.executemany(
            f"INSERT INTO {BENCH_TABLE} (id, embedding, content) "
            f"VALUES ($1, $2::vector, $3)",
            rows,
        )
        await conn.execute(
            f"CREATE INDEX ON {BENCH_TABLE} USING hnsw (embedding vector_cosine_ops)"
        )
        return vector_ids
    finally:
        await conn.close()


async def _register_resources(
    org_id: str, connector_id: str, vector_ids: list[str], n_resources: int
) -> None:
    """Insert GatedResource + ResourceChunk registry rows directly."""
    from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

    from gateco.database.enums import Classification, ResourceType, Sensitivity
    from gateco.database.models.resource import GatedResource
    from gateco.database.models.resource_chunk import ResourceChunk
    from gateco.database.settings import DatabaseSettings

    engine = create_async_engine(DatabaseSettings().database_url)
    factory = async_sessionmaker(engine)
    classifications = [
        Classification.public,
        Classification.internal,
        Classification.confidential,
        Classification.restricted,
    ]
    sensitivities = [Sensitivity.low, Sensitivity.medium, Sensitivity.high, Sensitivity.critical]
    domains = ["engineering", "finance", "marketing", "management"]

    async with factory() as session:
        vid_iter = iter(vector_ids)
        for i in range(n_resources):
            resource = GatedResource(
                organization_id=uuid.UUID(org_id),
                type=ResourceType.file,
                title=f"Bench Resource {i:03d}",
                content_url=f"bench://resource/{i:03d}",
                classification=classifications[i % 4],
                sensitivity=sensitivities[i % 4],
                domain=domains[i % 4],
                source_connector_id=uuid.UUID(connector_id),
            )
            session.add(resource)
            await session.flush()
            for c in range(CHUNKS_PER_RESOURCE):
                vid = next(vid_iter, None)
                if vid is None:
                    break
                session.add(
                    ResourceChunk(
                        resource_id=resource.id,
                        index=c,
                        preview=f"bench chunk {vid}",
                        vector_id=vid,
                        source_connector_id=uuid.UUID(connector_id),
                    )
                )
        await session.commit()
    await engine.dispose()


async def run_bench(args: argparse.Namespace) -> dict:
    rng = random.Random(20260816)
    base = args.base_url.rstrip("/")
    suffix = secrets.token_hex(4)

    async with httpx.AsyncClient(base_url=base, timeout=60.0) as http:
        # 1. Bench org + admin user
        email = f"bench-admin-{suffix}@gateco-bench.example"
        password = secrets.token_urlsafe(16) + "!Aa1"
        r = await http.post(
            "/api/auth/signup",
            json={
                "name": "Bench Admin",
                "email": email,
                "password": password,
                "organization_name": f"Policy Bench {suffix}",
            },
        )
        r.raise_for_status()
        data = r.json()
        access_token = data["tokens"]["access_token"]
        org_id = data["user"]["organization"]["id"]

        r = await http.patch(
            f"/api/admin/db/organizations/{org_id}/plan",
            json={"plan": "team"},
            headers={"X-Admin-Token": ADMIN_TOKEN},
        )
        r.raise_for_status()

        # Re-login so the JWT carries the upgraded plan claim
        r = await http.post("/api/auth/login", json={"email": email, "password": password})
        r.raise_for_status()
        access_token = r.json()["tokens"]["access_token"]
        auth = {"Authorization": f"Bearer {access_token}"}

        # 2. Stub IDP + principal
        r = await http.post(
            "/api/identity-providers",
            headers=auth,
            json={
                "name": f"bench-idp-{suffix}",
                # The placeholder/vh-fake markers make _is_stub_config() true:
                # the StubAdapter serves the embedded principals and the
                # vendor_iam (Growth+) gate is bypassed — no real vendor calls.
                "type": "okta",
                "config": {
                    "domain": "bench.placeholder",
                    "api_token": "vh-fake-bench-token",
                    "principals": [
                        {
                            "external_id": f"bench-p-{suffix}",
                            "display_name": "Bench Principal",
                            "email": f"bench-p-{suffix}@gateco-bench.example",
                            "groups": ["bench-eng"],
                            "roles": ["viewer"],
                            "department": "engineering",
                        }
                    ]
                },
            },
        )
        r.raise_for_status()
        idp_id = r.json().get("data", r.json())["id"]
        r = await http.post(f"/api/identity-providers/{idp_id}/sync", headers=auth)
        r.raise_for_status()
        await asyncio.sleep(1)
        r = await http.post(
            "/api/principals/resolve",
            headers=auth,
            json={"email": f"bench-p-{suffix}@gateco-bench.example"},
        )
        r.raise_for_status()
        principal_id = r.json().get("data", r.json())["id"]

        # 3. pgvector connector
        pg = {
            "host": os.environ.get("BENCH_PG_HOST", "localhost"),
            "port": int(os.environ.get("BENCH_PG_PORT", "5432")),
            "database": os.environ.get("BENCH_PG_DB", "gateco_db"),
            "user": os.environ.get("BENCH_PG_USER", "postgres"),
            "password": os.environ.get("BENCH_PG_PASSWORD", "postgres"),
        }
        r = await http.post(
            "/api/connectors",
            headers=auth,
            json={
                "name": f"bench-pgvector-{suffix}",
                "type": "pgvector",
                "config": pg,
                "metadata_resolution_mode": "sidecar",
            },
        )
        r.raise_for_status()
        connector_id = r.json().get("data", r.json())["id"]

        dsn = f"postgresql://{pg['user']}:{pg['password']}@{pg['host']}:{pg['port']}/{pg['database']}"
        vector_ids = await _setup_bench_table(dsn, args.vectors, rng)

        r = await http.patch(
            f"/api/connectors/{connector_id}/search-config",
            headers=auth,
            json={
                "search_config": {
                    "table_name": BENCH_TABLE,
                    "vector_column": "embedding",
                    "id_column": "id",
                    "content_column": "content",
                    "top_k": 50,
                    "expected_dimension": EMBED_DIM,
                }
            },
        )
        r.raise_for_status()

        # 4. Registry rows (100 resources x 10 chunks -> covers 1000 vectors)
        await _register_resources(org_id, connector_id, vector_ids, args.resources)

        # 5. One active 5-rule RBAC policy
        r = await http.post(
            "/api/policies",
            headers=auth,
            json={
                "name": f"bench-rbac-{suffix}",
                "description": "Benchmark policy: 5 RBAC rules",
                "type": "rbac",
                "effect": "allow",
                "resource_selectors": [
                    {"field": "connector_id", "op": "eq", "value": connector_id}
                ],
                "rules": [
                    {
                        "description": "group read public",
                        "effect": "allow",
                        "priority": 10,
                        "conditions": [
                            {"field": "principal.groups", "operator": "contains", "value": "bench-eng"},
                            {"field": "resource.classification", "operator": "eq", "value": "public"},
                        ],
                    },
                    {
                        "description": "group read internal",
                        "effect": "allow",
                        "priority": 20,
                        "conditions": [
                            {"field": "principal.groups", "operator": "contains", "value": "bench-eng"},
                            {"field": "resource.classification", "operator": "eq", "value": "internal"},
                        ],
                    },
                    {
                        "description": "deny restricted",
                        "effect": "deny",
                        "priority": 30,
                        "conditions": [
                            {"field": "resource.classification", "operator": "eq", "value": "restricted"},
                        ],
                    },
                    {
                        "description": "department domain allow",
                        "effect": "allow",
                        "priority": 40,
                        "conditions": [
                            {"field": "principal.attributes.department", "operator": "eq", "value": "engineering"},
                            {"field": "resource.domain", "operator": "eq", "value": "engineering"},
                        ],
                    },
                    {
                        "description": "deny critical sensitivity",
                        "effect": "deny",
                        "priority": 50,
                        "conditions": [
                            {"field": "resource.sensitivity", "operator": "eq", "value": "critical"},
                        ],
                    },
                ],
            },
        )
        r.raise_for_status()
        policy_id = r.json().get("data", r.json())["id"]
        r = await http.post(f"/api/policies/{policy_id}/activate", headers=auth)
        r.raise_for_status()

        # 6. Measure
        async def one_request() -> dict | None:
            qv = _unit_vector(rng)
            t0 = time.perf_counter()
            resp = await http.post(
                "/api/retrievals/execute",
                headers=auth,
                json={
                    "query_vector": qv,
                    "query": "bench probe",
                    "principal_id": principal_id,
                    "connector_id": connector_id,
                    "top_k": args.top_k,
                },
            )
            wall_ms = (time.perf_counter() - t0) * 1000.0
            if resp.status_code != 200:
                return None
            body = resp.json().get("data", resp.json())
            return {
                "wall_ms": wall_ms,
                "latency_ms": body.get("latency_ms"),
                "connector_latency_ms": body.get("connector_latency_ms"),
                "allowed": body.get("allowed_chunks"),
                "denied": body.get("denied_chunks"),
            }

        for _ in range(args.warmup):
            await one_request()
            await asyncio.sleep(1.0 / args.pace)

        samples: list[dict] = []
        errors = 0
        for _ in range(args.requests):
            s = await one_request()
            if s is None or s["latency_ms"] is None or s["connector_latency_ms"] is None:
                errors += 1
            else:
                samples.append(s)
            await asyncio.sleep(1.0 / args.pace)

    service_total = [s["latency_ms"] for s in samples]
    connector_leg = [s["connector_latency_ms"] for s in samples]
    service_overhead = [s["latency_ms"] - s["connector_latency_ms"] for s in samples]
    full_overhead = [s["wall_ms"] - s["connector_latency_ms"] for s in samples]

    report = {
        "methodology": {
            "connector": "pgvector (local)",
            "vectors": args.vectors,
            "embedding_dim": EMBED_DIM,
            "resources_registered": args.resources,
            "chunks_per_resource": CHUNKS_PER_RESOURCE,
            "policy": "1 active RBAC policy, 5 rules (2 group-allow, 1 classification deny, 1 department allow, 1 sensitivity deny)",
            "metadata_resolution": "sidecar",
            "top_k": args.top_k,
            "requests": args.requests,
            "warmup": args.warmup,
            "pace_rps": args.pace,
            "query_embedding": "client-supplied (server-side embedding excluded)",
            "notes": (
                "Run against a DEBUG=true dev backend (verbose logging and "
                "relaxed rate limits) — overhead numbers are conservative "
                "relative to production settings."
            ),
        },
        "errors": errors,
        "sample_allowed_denied": (
            {"allowed": samples[0]["allowed"], "denied": samples[0]["denied"]} if samples else {}
        ),
        "series": [
            _stats("connector_leg_ms (raw vector query)", connector_leg),
            _stats("service_total_ms (retrieval service end-to-end)", service_total),
            _stats("service_overhead_ms (policy layer, service-side)", service_overhead),
            _stats("full_overhead_ms (policy layer incl. HTTP middleware)", full_overhead),
        ],
    }
    return report


def main() -> int:
    if not ADMIN_TOKEN:
        print(
            "ERROR: ADMIN_TOKEN is not set. Export the admin token of the "
            "target environment before running (the benchmark uses it once, "
            "to set the bench org's plan tier).",
            file=sys.stderr,
        )
        return 2

    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base-url", default="http://localhost:8000")
    parser.add_argument("--requests", type=int, default=300)
    parser.add_argument("--warmup", type=int, default=20)
    parser.add_argument("--top-k", type=int, default=10)
    parser.add_argument("--vectors", type=int, default=1000)
    parser.add_argument("--resources", type=int, default=100)
    parser.add_argument("--pace", type=float, default=8.0, help="requests per second cap")
    parser.add_argument("--json-out", default="")
    args = parser.parse_args()

    report = asyncio.run(run_bench(args))

    print()
    print(f"{'series':<52} {'n':>4} {'p50':>8} {'p95':>8} {'p99':>8} {'mean':>8} {'max':>8}")
    for s in report["series"]:
        print(
            f"{s['series']:<52} {s['n']:>4} {s['p50_ms']:>8} {s['p95_ms']:>8} "
            f"{s['p99_ms']:>8} {s['mean_ms']:>8} {s['max_ms']:>8}"
        )
    print(f"\nerrors: {report['errors']}")

    if args.json_out:
        Path(args.json_out).write_text(json.dumps(report, indent=2))
        print(f"report written to {args.json_out}")

    overhead_p95 = next(
        s["p95_ms"] for s in report["series"] if s["series"].startswith("service_overhead")
    )
    print(f"\nservice policy-overhead p95: {overhead_p95} ms "
          f"({'WITHIN' if overhead_p95 < 25 else 'EXCEEDS'} the <25ms claim)")
    return 0


if __name__ == "__main__":
    sys.exit(main())