//! Deferred-write micro + DB-backed Criterion benchmarks.
//!
//! ## Quick run (no DB)
//! - `cargo bench --bench deferred_write --no-run`
//! - `cargo bench --bench deferred_write`
//!
//! Includes `deferred_write_gateway_ack_sequence` (WAL append then buffer push, production order).
//!
//! ## Optional DB-backed flush benchmark
//! Disabled by default to keep local runs safe.
//!
//! Enable with:
//! - `ATHENA_RUN_DEFERRED_WRITE_DB_BENCH=1`
//!
//! Optional knobs:
//! - `ATHENA_DEFERRED_WRITE_BENCH_TABLE` (default: `athena_deferred_write_bench`)
//! - `ATHENA_DEFERRED_WRITE_BENCH_BATCH_SIZES` (default: `10,100,1000`)
//! - `ATHENA_DEFERRED_WRITE_BENCH_SAMPLE_SIZE` (default: `20`)
//! - `ATHENA_DEFERRED_WRITE_BENCH_BUFFER_SLACK` (default: `16`)
//!
//! Reuses URI/client resolution used by other benches:
//! - `ATHENA_BENCHMARK_CLIENT`
//! - `ATHENA_GATEWAY_BENCH_POSTGRES_URI`
//! - `POSTGRES_ATHENA_LOGGING_URI` fallback when client is `athena_logging`.

use std::sync::Arc;

use athena_rs::config::Config;
use athena_rs::deferred_write::{DeferredEntry, WriteBuffer, flush_pending};
use athena_rs::drivers::postgresql::sqlx_driver::PostgresClientRegistry;
use athena_rs::features::connection_pooler::ConnectionPoolManager;
use athena_rs::parser::resolve_postgres_uri;
use criterion::{Criterion, black_box, criterion_group, criterion_main};
use dotenv::dotenv;
use moka::future::Cache;
use serde_json::json;
use sqlx::{Pool, Postgres};
use tokio::runtime::Runtime;

fn bench_deferred_write_buffer(c: &mut Criterion) {
    let runtime: Runtime = Runtime::new().expect("tokio runtime");
    let buffer: Arc<WriteBuffer> = Arc::new(WriteBuffer::new(10_000));

    c.bench_function("deferred_write_buffer_push", |b| {
        b.iter(|| {
            runtime.block_on(async {
                let entry = DeferredEntry::new_insert(
                    "athena_logging".to_string(),
                    "athena_gateway_bench".to_string(),
                    json!({
                        "workspace_id": "bench-ws",
                        "name": "buffer-push",
                        "status": "active",
                    }),
                );
                buffer.push(entry).await;
                black_box(buffer.len().await);
                let _ = buffer.drain_all().await;
            });
        });
    });
}

fn bench_deferred_write_wal_append(c: &mut Criterion) {
    let runtime: Runtime = Runtime::new().expect("tokio runtime");
    let wal_root = std::env::temp_dir().join("athena_bench_deferred_write");
    let wal = athena_rs::deferred_write::WalManager::new(wal_root.to_string_lossy().as_ref())
        .expect("create wal manager");

    c.bench_function("deferred_write_wal_append", |b| {
        b.iter(|| {
            runtime.block_on(async {
                let entry = DeferredEntry::new_insert(
                    "athena_logging".to_string(),
                    "athena_gateway_bench".to_string(),
                    json!({
                        "workspace_id": "bench-ws",
                        "name": "wal-append",
                        "status": "active",
                    }),
                );
                wal.append(&entry).await.expect("append wal");
                black_box(entry.id);
                wal.clear().await.expect("clear wal");
            });
        });
    });
}

/// Same ordering as production gateway deferred insert: WAL append, then buffer push.
fn bench_deferred_write_gateway_ack_sequence(c: &mut Criterion) {
    let runtime: Runtime = Runtime::new().expect("tokio runtime");
    let wal_root = std::env::temp_dir().join("athena_bench_deferred_ack_seq");
    let wal = athena_rs::deferred_write::WalManager::new(wal_root.to_string_lossy().as_ref())
        .expect("create wal manager for ack sequence");
    let buffer: Arc<WriteBuffer> = Arc::new(WriteBuffer::new(10_000));

    c.bench_function("deferred_write_gateway_ack_sequence", |b| {
        b.iter(|| {
            runtime.block_on(async {
                let entry = DeferredEntry::new_insert(
                    "athena_logging".to_string(),
                    "athena_gateway_bench".to_string(),
                    json!({
                        "workspace_id": "bench-ws",
                        "name": "ack-sequence",
                        "status": "active",
                    }),
                );
                wal.append(&entry).await.expect("append wal");
                buffer.push(entry).await;
                black_box(buffer.len().await);
                let _ = buffer.drain_all().await;
                wal.clear().await.expect("clear wal");
            });
        });
    });
}

const DEFAULT_DEFERRED_BENCH_TABLE: &str = "athena_deferred_write_bench";
const DEFAULT_BATCH_SIZES: &[usize] = &[10, 100, 1000];
const DEFAULT_SAMPLE_SIZE: usize = 20;
const DEFAULT_BUFFER_SLACK: usize = 16;

fn deferred_bench_table_name() -> String {
    std::env::var("ATHENA_DEFERRED_WRITE_BENCH_TABLE")
        .ok()
        .map(|v| v.trim().to_string())
        .filter(|v| !v.is_empty())
        .unwrap_or_else(|| DEFAULT_DEFERRED_BENCH_TABLE.to_string())
}

fn deferred_bench_sample_size() -> usize {
    std::env::var("ATHENA_DEFERRED_WRITE_BENCH_SAMPLE_SIZE")
        .ok()
        .and_then(|v| v.trim().parse::<usize>().ok())
        .filter(|v| *v > 0)
        .unwrap_or(DEFAULT_SAMPLE_SIZE)
}

fn deferred_bench_buffer_slack() -> usize {
    std::env::var("ATHENA_DEFERRED_WRITE_BENCH_BUFFER_SLACK")
        .ok()
        .and_then(|v| v.trim().parse::<usize>().ok())
        .unwrap_or(DEFAULT_BUFFER_SLACK)
}

fn deferred_bench_batch_sizes() -> Vec<usize> {
    std::env::var("ATHENA_DEFERRED_WRITE_BENCH_BATCH_SIZES")
        .ok()
        .map(|raw| {
            raw.split(',')
                .filter_map(|s| s.trim().parse::<usize>().ok())
                .filter(|v| *v > 0)
                .collect::<Vec<usize>>()
        })
        .filter(|v| !v.is_empty())
        .unwrap_or_else(|| DEFAULT_BATCH_SIZES.to_vec())
}

fn benchmark_client_name() -> 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(|cfg| cfg.get_gateway_benchmark_client().cloned())
        .unwrap_or_else(|| "athena_logging".to_string())
}

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
    }
}

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
}

async fn provision_deferred_bench_table(
    pool: &Pool<Postgres>,
    table_name: &str,
) -> Result<(), sqlx::Error> {
    let create_sql = format!(
        "CREATE TABLE IF NOT EXISTS {table_name} (
            id SERIAL PRIMARY KEY,
            workspace_id TEXT NOT NULL DEFAULT 'bench-ws',
            name TEXT,
            status TEXT NOT NULL DEFAULT 'active'
        )"
    );
    sqlx::query(&create_sql).execute(pool).await?;

    let truncate_sql = format!("TRUNCATE TABLE {table_name} RESTART IDENTITY CASCADE");
    sqlx::query(&truncate_sql).execute(pool).await?;

    Ok(())
}

fn bench_deferred_write_flush_db(c: &mut Criterion) {
    let _ = dotenv();
    let enabled = std::env::var("ATHENA_RUN_DEFERRED_WRITE_DB_BENCH")
        .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
        .unwrap_or(false);
    if !enabled {
        eprintln!(
            "Skipping DB deferred-write bench (set ATHENA_RUN_DEFERRED_WRITE_DB_BENCH=1 to enable)"
        );
        return;
    }

    let config: Config = Config::load().unwrap_or_else(|e| {
        panic!("Deferred-write DB bench: failed to load config.yaml: {e}");
    });
    let client_name: String = benchmark_client_name();
    let uri: String = resolve_bench_postgres_uri(&config, &client_name).unwrap_or_else(|| {
        panic!(
            "Deferred-write DB bench: no usable Postgres URI for client '{client_name}'. \
             Set postgres_clients.{client_name}, ATHENA_GATEWAY_BENCH_POSTGRES_URI, or \
             POSTGRES_ATHENA_LOGGING_URI (for athena_logging)."
        )
    });

    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!("Deferred-write DB bench: failed to create registry: {e}"));

    if let Some((failed_client, err)) = errors.first() {
        panic!("Deferred-write DB bench: failed to connect client '{failed_client}': {err}");
    }

    let pool: Pool<Postgres> = registry.get_pool(&client_name).unwrap_or_else(|| {
        panic!(
            "Deferred-write DB bench: registry missing pool for client '{}'.",
            client_name
        );
    });

    let table_name = deferred_bench_table_name();
    let batch_sizes = deferred_bench_batch_sizes();
    let sample_size = deferred_bench_sample_size();
    let buffer_slack = deferred_bench_buffer_slack();

    runtime
        .block_on(provision_deferred_bench_table(&pool, &table_name))
        .unwrap_or_else(|e| panic!("Deferred-write DB bench: provision failed: {e}"));

    let pg_registry = Arc::new(registry);
    let cache: Arc<Cache<String, serde_json::Value>> = Arc::new(Cache::builder().build());

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

    for batch_size in batch_sizes {
        let bench_name = format!("flush_batch_{batch_size}");
        let pg_registry = Arc::clone(&pg_registry);
        let cache = Arc::clone(&cache);
        let client_name = client_name.clone();
        let table_name = table_name.clone();
        group.bench_function(bench_name, |b| {
            b.iter(|| {
                runtime.block_on(async {
                    let buffer = WriteBuffer::new(batch_size + buffer_slack);
                    for idx in 0..batch_size {
                        let entry = DeferredEntry::new_insert(
                            client_name.clone(),
                            table_name.to_string(),
                            json!({
                                "workspace_id": "bench-ws",
                                "name": format!("row-{batch_size}-{idx}"),
                                "status": "active",
                            }),
                        );
                        buffer.push(entry).await;
                    }

                    let summary = flush_pending(&buffer, None, &pg_registry, &cache).await;
                    assert_eq!(summary.total, batch_size, "unexpected flushed total");
                    assert_eq!(summary.failed, 0, "unexpected failed flush entries");
                    black_box(summary.succeeded);
                });
            });
        });
    }

    group.finish();
}

criterion_group!(
    benches,
    bench_deferred_write_buffer,
    bench_deferred_write_wal_append,
    bench_deferred_write_gateway_ack_sequence,
    bench_deferred_write_flush_db
);
criterion_main!(benches);
