//! End-to-end Criterion benchmarks for gateway HTTP handlers against a **real Postgres** database.
//!
//! ## Prereqs
//!
//! The benchmark client (see below) must have a resolvable URI and a live server reachable from
//! the machine running `cargo bench`:
//!
//! - `postgres_clients.<name>` in `config.yaml` (values may use `${ENV}` placeholders), and/or
//! - env **`ATHENA_GATEWAY_BENCH_POSTGRES_URI`** (full URI, overrides config for this bench only), and/or
//! - env **`POSTGRES_ATHENA_LOGGING_URI`** applied when the benchmark client is `athena_logging` and config still has `${POSTGRES_ATHENA_LOGGING_URI}`.
//!
//! On first connect the bench creates **`athena_gateway_bench`** and **`athena_gateway_bench_sink`**
//! if missing, truncates them, seeds rows, and installs **`bench_gateway_rpc_stub()`**. It assumes the
//! database may be empty or missing these tables (idempotent `CREATE TABLE IF NOT EXISTS`).
//!
//! ## Benchmark client header
//!
//! Same resolution order as production tooling:
//! 1. **`ATHENA_BENCHMARK_CLIENT`** (non-empty)
//! 2. **`gateway.benchmark_client`** in config (else **`gateway.logging_client`**)
//! 3. default **`athena_logging`**
//!
//! Run: `cargo bench --bench gateway_e2e`
//!
//! Quick smoke: `cargo bench --bench gateway_e2e -- --quick`

use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};

use actix_web::dev::ServiceResponse;
use actix_web::test;
use actix_web::web::Data;
use actix_web::{App, http::StatusCode};
use athena_rs::AppState;
use athena_rs::api::gateway::insert::{InsertWindowCoordinator, InsertWindowSettings};
use athena_rs::api::gateway::rpc::{gateway_rpc_route, rpc_get_route, rpc_post_route};
use athena_rs::api::gateway::{
    delete::delete_data,
    fetch::{fetch_data_route, gateway_update_route, get_data_route, proxy_fetch_data_route},
    insert::insert_data,
    postgrest::{
        postgrest_delete_route, postgrest_get_route, postgrest_patch_route, postgrest_post_route,
    },
    query::gateway_query_route,
};
use athena_rs::api::pipelines::run_pipeline;
use athena_rs::config::Config;
use athena_rs::drivers::postgresql::pool_manager::ConnectionPoolManager;
use athena_rs::drivers::postgresql::sqlx_driver::PostgresClientRegistry;
use athena_rs::parser::resolve_postgres_uri;
use criterion::{Criterion, black_box, criterion_group, criterion_main};
use dotenv::dotenv;
use moka::future::Cache;
use reqwest::Client;
use serde_json::{Value, json};
use sqlx::{Pool, Postgres};
use tokio::runtime::Runtime;
use tokio::task::LocalSet;

const BENCH_TABLE: &str = "athena_gateway_bench";
const BENCH_SINK_TABLE: &str = "athena_gateway_bench_sink";

static INSERT_SEQ: AtomicU64 = AtomicU64::new(0);

fn benchmark_athena_client_for_requests() -> String {
    if let Ok(value) = std::env::var("ATHENA_BENCHMARK_CLIENT") {
        let trimmed: String = value.trim().to_string();
        if !trimmed.is_empty() {
            return trimmed;
        }
    }
    Config::load()
        .ok()
        .and_then(|config| config.get_gateway_benchmark_client().cloned())
        .unwrap_or_else(|| "athena_logging".to_string())
}

fn resolve_bench_postgres_uri(config: &Config, client_name: &str) -> Option<String> {
    if let Ok(uri) = std::env::var("ATHENA_GATEWAY_BENCH_POSTGRES_URI") {
        let trimmed = uri.trim();
        if !trimmed.is_empty() {
            if let Some(u) = normalize_bench_postgres_candidate(trimmed) {
                return Some(u);
            }
        }
    }
    if let Some(template) = config.get_postgres_uri(client_name) {
        if let Some(u) = normalize_bench_postgres_candidate(template) {
            return Some(u);
        }
    }
    if client_name.eq_ignore_ascii_case("athena_logging") {
        if let Ok(uri) = std::env::var("POSTGRES_ATHENA_LOGGING_URI") {
            let trimmed = uri.trim();
            if !trimmed.is_empty() {
                if let Some(u) = normalize_bench_postgres_candidate(trimmed) {
                    return Some(u);
                }
            }
        }
    }
    None
}

fn bench_connection_panic_message(client_name: &str, err: &str) -> String {
    format!(
        "Gateway E2E bench could not connect Postgres client '{client_name}'.\n\
         Fix: set postgres_clients.{client_name} in config.yaml (with env vars resolved), or set \
         ATHENA_GATEWAY_BENCH_POSTGRES_URI to a full connection URI, or ensure POSTGRES_ATHENA_LOGGING_URI \
         is set when using athena_logging.\n\
         The bench loads `.env` from the current directory (same as `cargo run`).\n\
         Config load: use ATHENA_CONFIG_PATH or a writable config search path.\n\
         Error: {err}",
    )
}

/// True when `sqlx` can use this string as a Postgres URL (not an unresolved `${VAR}` placeholder).
fn postgres_uri_usable(uri: &str) -> bool {
    let t = uri.trim();
    if t.is_empty() || t.starts_with("${") {
        return false;
    }
    t.starts_with("postgres://") || t.starts_with("postgresql://")
}

fn normalize_bench_postgres_candidate(raw: &str) -> Option<String> {
    let resolved = resolve_postgres_uri(raw.trim());
    if postgres_uri_usable(&resolved) {
        Some(resolved)
    } else {
        None
    }
}

async fn provision_and_seed_benchmark_tables(pool: &Pool<Postgres>) -> Result<(), sqlx::Error> {
    sqlx::query(
        "CREATE TABLE IF NOT EXISTS athena_gateway_bench (
            id SERIAL PRIMARY KEY,
            workspace_id TEXT NOT NULL DEFAULT 'bench-ws',
            name TEXT,
            status TEXT NOT NULL DEFAULT 'active'
        )",
    )
    .execute(pool)
    .await?;

    sqlx::query(
        "CREATE TABLE IF NOT EXISTS athena_gateway_bench_sink (
            id SERIAL PRIMARY KEY,
            workspace_id TEXT,
            name TEXT,
            status TEXT
        )",
    )
    .execute(pool)
    .await?;

    sqlx::query("TRUNCATE TABLE athena_gateway_bench RESTART IDENTITY CASCADE")
        .execute(pool)
        .await?;
    sqlx::query("TRUNCATE TABLE athena_gateway_bench_sink RESTART IDENTITY CASCADE")
        .execute(pool)
        .await?;

    sqlx::query(
        "INSERT INTO athena_gateway_bench (workspace_id, name, status) VALUES
            ('bench-ws', 'seed-a', 'active'),
            ('bench-ws', 'seed-b', 'active'),
            ('bench-ws', 'seed-c', 'active')",
    )
    .execute(pool)
    .await?;

    sqlx::query(
        r#"CREATE OR REPLACE FUNCTION public.bench_gateway_rpc_stub()
        RETURNS TABLE(val integer)
        LANGUAGE sql
        STABLE
        AS $$ SELECT 1 AS val $$"#,
    )
    .execute(pool)
    .await?;

    Ok(())
}

fn build_gateway_bench_state(
    pg_registry: Arc<PostgresClientRegistry>,
    benchmark_client_label: &str,
) -> Data<AppState> {
    let cache: Arc<Cache<String, Value>> = Arc::new(Cache::builder().build());
    let immortal_cache: Arc<Cache<String, Value>> = Arc::new(Cache::builder().build());
    let client: Client = Client::builder()
        .pool_idle_timeout(Duration::from_secs(30))
        .build()
        .expect("Failed to build HTTP client");
    let jdbc_pool_cache: Arc<Cache<String, Pool<Postgres>>> =
        Arc::new(Cache::builder().max_capacity(64).build());
    let insert_window_coordinator: Arc<InsertWindowCoordinator> =
        InsertWindowCoordinator::new(InsertWindowSettings {
            max_batch: 100,
            max_queued: 10_000,
            deny_tables: Default::default(),
        });
    let data: Data<AppState> = Data::new(AppState {
        cache,
        immortal_cache,
        client,
        process_start_time_seconds: 0_i64,
        process_started_at: Instant::now(),
        pg_registry,
        jdbc_pool_cache,
        #[cfg(feature = "deadpool_experimental")]
        deadpool_registry: Arc::new(
            athena_rs::drivers::postgresql::deadpool_registry::DeadpoolPostgresRegistry::empty(),
        ),
        #[cfg(feature = "deadpool_experimental")]
        jdbc_deadpool_cache: Arc::new(Cache::builder().max_capacity(4).build()),
        gateway_force_camel_case_to_snake_case: false,
        gateway_auto_cast_uuid_filter_values_to_text: true,
        gateway_allow_schema_names_prefixed_as_table_name: true,
        pipeline_registry: None,
        logging_client_name: None,
        gateway_auth_client_name: None,
        gateway_benchmark_client_name: Some(benchmark_client_label.to_string()),
        gateway_api_key_fail_mode: "fail_closed".to_string(),
        gateway_jdbc_allow_private_hosts: true,
        gateway_jdbc_allowed_hosts: vec![],
        gateway_resilience_timeout_secs: 30,
        gateway_resilience_read_max_retries: 1,
        gateway_resilience_initial_backoff_ms: 100,
        gateway_admission_store_backend: "redis".to_string(),
        gateway_admission_store_fail_mode: "fail_closed".to_string(),
        prometheus_metrics_enabled: false,
        metrics_state: Arc::new(athena_rs::api::metrics::MetricsState::new()),
        gateway_insert_execution_window_ms: 0,
        gateway_insert_window_max_batch: 100,
        gateway_insert_window_max_queued: 10_000,
        gateway_insert_merge_deny_tables: Default::default(),
        insert_window_coordinator: insert_window_coordinator.clone(),
        client_stats_batcher: None,
        gateway_log_batcher: None,
        logging_task_limiter: None,
    });
    data
}

macro_rules! gateway_bench_app {
    ($state:expr) => {
        App::new()
            .app_data($state.clone())
            .service(fetch_data_route)
            .service(proxy_fetch_data_route)
            .service(gateway_update_route)
            .service(get_data_route)
            .service(insert_data)
            .service(delete_data)
            .service(postgrest_get_route)
            .service(postgrest_post_route)
            .service(postgrest_patch_route)
            .service(postgrest_delete_route)
            .service(gateway_query_route)
            .service(gateway_rpc_route)
            .service(rpc_post_route)
            .service(rpc_get_route)
            .service(run_pipeline)
    };
}

fn assert_http_success(resp: &ServiceResponse, context: &str) {
    let code = resp.status();
    assert!(code.is_success(), "{context}: expected 2xx, got {code}");
}

fn bench_gateway_e2e(c: &mut Criterion) {
    let _ = dotenv();
    let client_name: String = benchmark_athena_client_for_requests();
    let config: Config = Config::load().unwrap_or_else(|e| {
        panic!("Gateway E2E bench: failed to load config.yaml (try ATHENA_CONFIG_PATH): {e}");
    });
    let uri: String =
        resolve_bench_postgres_uri(&config, client_name.as_str()).unwrap_or_else(|| {
            let unresolved_hint: String = config
            .get_postgres_uri(client_name.as_str())
            .map(|template| resolve_postgres_uri(template))
            .filter(|resolved| !postgres_uri_usable(resolved))
            .map(|resolved| {
                format!(
                    "\nConfig `postgres_clients.{client_name}` currently resolves to `{resolved}`, \
                     which is not a usable postgres:// URL. If you use `${{VAR}}` placeholders, \
                     define VAR (e.g. in repo-root `.env`) before `cargo bench`, or set \
                     ATHENA_GATEWAY_BENCH_POSTGRES_URI.",
                    client_name = client_name,
                    resolved = resolved
                )
            })
            .unwrap_or_default();
            panic!(
                "Gateway E2E bench: no Postgres URI for client '{client_name}'. \
             Set postgres_clients in config, ATHENA_GATEWAY_BENCH_POSTGRES_URI, or \
             POSTGRES_ATHENA_LOGGING_URI (for athena_logging).{unresolved_hint}",
            );
        });

    let runtime: Runtime = Runtime::new().expect("tokio runtime");
    let pool_manager: ConnectionPoolManager = ConnectionPoolManager::default();
    let (registry, errors) = runtime
        .block_on(async {
            PostgresClientRegistry::from_entries(vec![(client_name.clone(), uri)], pool_manager)
                .await
        })
        .unwrap_or_else(|e| {
            panic!(
                "{}",
                bench_connection_panic_message(&client_name, &e.to_string())
            )
        });

    if let Some((failed_client, err)) = errors.first() {
        panic!(
            "{}",
            bench_connection_panic_message(failed_client, &err.to_string())
        );
    }

    let pool: Pool<Postgres> = registry.get_pool(&client_name).unwrap_or_else(|| {
        panic!(
            "Gateway E2E bench: connected but no pool for client '{}' (registry empty).",
            client_name
        );
    });

    runtime
        .block_on(provision_and_seed_benchmark_tables(&pool))
        .unwrap_or_else(|e| panic!("Gateway E2E bench: provision/seed failed: {e}"));

    let pg_registry: Arc<PostgresClientRegistry> = Arc::new(registry);
    let app_state: Data<AppState> = build_gateway_bench_state(pg_registry, &client_name);
    // `InsertWindowCoordinator::bind_app_state` uses `tokio::spawn`; must run inside the runtime.
    runtime.block_on(async {
        app_state
            .insert_window_coordinator
            .bind_app_state(app_state.clone());
    });

    let local: LocalSet = LocalSet::new();
    let app = local.block_on(&runtime, async {
        test::init_service(gateway_bench_app!(app_state)).await
    });

    let client: &str = client_name.as_str();

    let fetch_body: Value = json!({
        "table_name": BENCH_TABLE,
        "conditions": [{"eq_column": "workspace_id", "eq_value": "bench-ws"}],
        "limit": 50
    });
    let update_body: Value = json!({
        "table_name": BENCH_TABLE,
        "columns": [{ "status": "bench-updated" }],
        "conditions": [{ "eq_column": "id", "eq_value": "1" }]
    });
    let pipeline_body: Value = json!({
        "source": {
            "table_name": BENCH_TABLE,
            "columns": ["id", "workspace_id", "name", "status"],
            "conditions": [],
            "limit": 5
        },
        "sink": { "table_name": BENCH_SINK_TABLE }
    });
    let rpc_body: Value = json!({ "function": "bench_gateway_rpc_stub" });

    let mut group = c.benchmark_group("gateway_e2e_actix");
    group.sample_size(30);

    group.bench_function("post_gateway_data_fetch_ok", |b| {
        b.iter(|| {
            local.block_on(&runtime, async {
                let req = test::TestRequest::post()
                    .uri("/gateway/data")
                    .insert_header(("X-Athena-Client", client))
                    .insert_header(("Content-Type", "application/json"))
                    .set_json(&fetch_body)
                    .to_request();
                let resp: ServiceResponse = test::call_service(&app, req).await;
                assert_http_success(&resp, "POST /gateway/data");
                let body = test::read_body(resp).await;
                black_box(body.len());
            });
        });
    });

    group.bench_function("post_gateway_fetch_ok", |b| {
        b.iter(|| {
            local.block_on(&runtime, async {
                let req = test::TestRequest::post()
                    .uri("/gateway/fetch")
                    .insert_header(("X-Athena-Client", client))
                    .insert_header(("Content-Type", "application/json"))
                    .set_json(&fetch_body)
                    .to_request();
                let resp: ServiceResponse = test::call_service(&app, req).await;
                assert_http_success(&resp, "POST /gateway/fetch");
                let body = test::read_body(resp).await;
                black_box(body.len());
            });
        });
    });

    let view = BENCH_TABLE;
    // Match the POST /gateway/fetch filter (TEXT column + string param). Query-param GET /data
    // supplies eq_value as text; filtering on SERIAL `id` can fail against some PG/sqlx pairings.
    let get_data_uri: String = format!(
        "/data?view={view}&eq_column=workspace_id&eq_value=bench-ws&limit=10",
        view = view
    );
    group.bench_function("get_data_route_ok", |b| {
        let uri: String = get_data_uri.clone();
        b.iter(|| {
            local.block_on(&runtime, async {
                let req = test::TestRequest::get()
                    .uri(&uri)
                    .insert_header(("X-Athena-Client", client))
                    .to_request();
                let resp: ServiceResponse = test::call_service(&app, req).await;
                assert_http_success(&resp, "GET /data");
                let body = test::read_body(resp).await;
                black_box(body.len());
            });
        });
    });

    group.bench_function("get_data_route_missing_view_400", |b| {
        b.iter(|| {
            local.block_on(&runtime, async {
                let req = test::TestRequest::get()
                    .uri("/data")
                    .insert_header(("X-Athena-Client", client))
                    .to_request();
                let resp: ServiceResponse = test::call_service(&app, req).await;
                assert_eq!(resp.status(), StatusCode::BAD_REQUEST, "missing view");
                let body = test::read_body(resp).await;
                black_box(body.len());
            });
        });
    });

    group.bench_function("post_gateway_query_select_1_ok", |b| {
        let query_body: Value = json!({ "query": "SELECT 1 AS one" });
        b.iter(|| {
            local.block_on(&runtime, async {
                let req = test::TestRequest::post()
                    .uri("/gateway/query")
                    .insert_header(("X-Athena-Client", client))
                    .insert_header(("Content-Type", "application/json"))
                    .set_json(&query_body)
                    .to_request();
                let resp: ServiceResponse = test::call_service(&app, req).await;
                assert_http_success(&resp, "POST /gateway/query");
                let body = test::read_body(resp).await;
                black_box(body.len());
            });
        });
    });

    let rest_list_uri: String = format!("/rest/v1/{BENCH_TABLE}?limit=10");
    group.bench_function("get_rest_v1_list_ok", |b| {
        let uri: String = rest_list_uri.clone();
        b.iter(|| {
            local.block_on(&runtime, async {
                let req = test::TestRequest::get()
                    .uri(&uri)
                    .insert_header(("X-Athena-Client", client))
                    .to_request();
                let resp: ServiceResponse = test::call_service(&app, req).await;
                assert_http_success(&resp, "GET /rest/v1");
                let body = test::read_body(resp).await;
                black_box(body.len());
            });
        });
    });

    let rest_patch_uri: String = format!("/rest/v1/{BENCH_TABLE}?id=eq.2");
    let rest_patch_json: Value = json!({ "status": "patched" });
    group.bench_function("patch_rest_v1_ok", |b| {
        let uri: String = rest_patch_uri.clone();
        let payload: Value = rest_patch_json.clone();
        b.iter(|| {
            local.block_on(&runtime, async {
                let req = test::TestRequest::patch()
                    .uri(&uri)
                    .insert_header(("X-Athena-Client", client))
                    .insert_header(("Content-Type", "application/json"))
                    .set_json(&payload)
                    .to_request();
                let resp: ServiceResponse = test::call_service(&app, req).await;
                assert_http_success(&resp, "PATCH /rest/v1");
                let body = test::read_body(resp).await;
                black_box(body.len());
            });
        });
    });

    group.bench_function("put_gateway_insert_ok", |b| {
        b.iter(|| {
            local.block_on(&runtime, async {
                let n: u64 = INSERT_SEQ.fetch_add(1, Ordering::Relaxed);
                let insert_body: Value = json!({
                    "table_name": BENCH_TABLE,
                    "insert_body": {
                        "name": format!("bench-row-{n}"),
                        "workspace_id": "bench-ws",
                        "status": "inserted"
                    }
                });
                let req = test::TestRequest::put()
                    .uri("/gateway/insert")
                    .insert_header(("X-Athena-Client", client))
                    .insert_header(("Content-Type", "application/json"))
                    .set_json(&insert_body)
                    .to_request();
                let resp: ServiceResponse = test::call_service(&app, req).await;
                assert_http_success(&resp, "PUT /gateway/insert");
                let body = test::read_body(resp).await;
                black_box(body.len());
            });
        });
    });

    group.bench_function("post_gateway_update_ok", |b| {
        let payload: Value = update_body.clone();
        b.iter(|| {
            local.block_on(&runtime, async {
                let req = test::TestRequest::post()
                    .uri("/gateway/update")
                    .insert_header(("X-Athena-Client", client))
                    .insert_header(("Content-Type", "application/json"))
                    .set_json(&payload)
                    .to_request();
                let resp: ServiceResponse = test::call_service(&app, req).await;
                assert_http_success(&resp, "POST /gateway/update");
                let body = test::read_body(resp).await;
                black_box(body.len());
            });
        });
    });

    group.bench_function("post_gateway_rpc_ok", |b| {
        let payload: Value = rpc_body.clone();
        b.iter(|| {
            local.block_on(&runtime, async {
                let req = test::TestRequest::post()
                    .uri("/gateway/rpc")
                    .insert_header(("X-Athena-Client", client))
                    .insert_header(("Content-Type", "application/json"))
                    .set_json(&payload)
                    .to_request();
                let resp: ServiceResponse = test::call_service(&app, req).await;
                assert_http_success(&resp, "POST /gateway/rpc");
                let body = test::read_body(resp).await;
                black_box(body.len());
            });
        });
    });

    group.bench_function("post_rpc_path_compat_ok", |b| {
        b.iter(|| {
            local.block_on(&runtime, async {
                let req = test::TestRequest::post()
                    .uri("/rpc/bench_gateway_rpc_stub")
                    .insert_header(("X-Athena-Client", client))
                    .insert_header(("Content-Type", "application/json"))
                    .set_json(json!({}))
                    .to_request();
                let resp: ServiceResponse = test::call_service(&app, req).await;
                assert_http_success(&resp, "POST /rpc/{fn}");
                let body = test::read_body(resp).await;
                black_box(body.len());
            });
        });
    });

    group.bench_function("get_rpc_path_compat_ok", |b| {
        b.iter(|| {
            local.block_on(&runtime, async {
                let req = test::TestRequest::get()
                    .uri("/rpc/bench_gateway_rpc_stub")
                    .insert_header(("X-Athena-Client", client))
                    .to_request();
                let resp: ServiceResponse = test::call_service(&app, req).await;
                assert_http_success(&resp, "GET /rpc/{fn}");
                let body = test::read_body(resp).await;
                black_box(body.len());
            });
        });
    });

    let pipe: Value = pipeline_body.clone();
    group.bench_function("post_pipelines_inline_ok", |b| {
        b.iter(|| {
            local.block_on(&runtime, async {
                let req = test::TestRequest::post()
                    .uri("/pipelines")
                    .insert_header(("X-Athena-Client", client))
                    .insert_header(("Content-Type", "application/json"))
                    .set_json(&pipe)
                    .to_request();
                let resp: ServiceResponse = test::call_service(&app, req).await;
                assert_http_success(&resp, "POST /pipelines");
                let body = test::read_body(resp).await;
                black_box(body.len());
            });
        });
    });

    group.finish();
}

criterion_group!(benches, bench_gateway_e2e);
criterion_main!(benches);
