//! `pgcrate sql` (alias `query`) — the guarded SQL gateway.
//!
//! Read-only by default. Writes preview themselves: with `--allow-write` a
//! statement runs inside a transaction, reports what it *would* do (affected
//! count + sample rows), then **rolls back**. `--commit` (which implies
//! `--allow-write`) is the only way to actually apply a change. Expensive
//! statements are flagged via EXPLAIN before they run, and runaway result sets
//! are capped.

use anyhow::{bail, Context, Result};
use serde::Serialize;
use std::io::Read;
use tokio_postgres::{Client, SimpleQueryMessage};

use crate::diagnostic::{DiagnosticSession, TimeoutConfig};

/// Default cap on rows printed/returned for a SELECT. `--limit 0` uncaps.
pub const DEFAULT_ROW_CAP: usize = 1000;

/// How many sample rows a dry-run write previews.
const DRY_RUN_SAMPLE_LIMIT: usize = 5;

/// Default EXPLAIN total-cost threshold above which a warning prints.
/// Postgres "cost" is in arbitrary planner units; ~50k corresponds roughly to
/// a sizable sequential scan. Tune via `[sql] cost_warn_threshold`.
pub const DEFAULT_COST_WARN_THRESHOLD: f64 = 50_000.0;

/// Options resolved from CLI flags + config for a single `sql` invocation.
pub struct SqlOptions {
    pub allow_write: bool,
    pub commit: bool,
    pub limit: Option<usize>,
    pub no_cost_check: bool,
    pub cost_warn_threshold: f64,
    /// Connect/statement/lock timeouts (from global CLI flags, same as DBA).
    pub timeouts: TimeoutConfig,
    pub quiet: bool,
    pub json: bool,
    /// Readable-output density for result tables (ignored in JSON/quiet mode).
    pub density: crate::output::Density,
}

#[derive(Serialize)]
struct SqlResponse {
    ok: bool,
    /// Present only on a write path; describes whether changes were committed.
    #[serde(skip_serializing_if = "Option::is_none")]
    write: Option<WriteOutcome>,
    results: Vec<SqlResult>,
}

#[derive(Serialize)]
struct WriteOutcome {
    /// `true` when changes were committed, `false` when rolled back (dry run).
    committed: bool,
    /// Total rows affected across all write statements.
    rows_affected: u64,
    /// EXPLAIN cost estimate for the statement, when one was computed. Present
    /// so `--json` consumers get the same cost signal humans see on stderr.
    #[serde(skip_serializing_if = "Option::is_none")]
    cost: Option<CostEstimate>,
}

/// Structured EXPLAIN cost signal surfaced in `--json` output.
#[derive(Serialize, Clone)]
struct CostEstimate {
    /// Estimated total plan cost (Postgres planner units).
    estimated: f64,
    /// The configured warning threshold this was checked against.
    threshold: f64,
    /// `true` when `estimated` exceeds `threshold` (a warning was emitted).
    exceeds_threshold: bool,
}

#[derive(Serialize)]
#[serde(tag = "type")]
enum SqlResult {
    #[serde(rename = "query")]
    Query {
        columns: Vec<String>,
        rows: Vec<Vec<Option<String>>>,
        /// Rows withheld by the row cap (0 when nothing was truncated).
        #[serde(skip_serializing_if = "is_zero")]
        truncated: usize,
    },
    #[serde(rename = "command")]
    CommandComplete { rows: u64 },
    /// Sample rows previewed for a dry-run write (the wrapped RETURNING result).
    #[serde(rename = "sample")]
    Sample {
        columns: Vec<String>,
        rows: Vec<Vec<Option<String>>>,
    },
}

fn is_zero(n: &usize) -> bool {
    *n == 0
}

/// Classification of a single parsed statement.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StmtKind {
    /// SELECT, SET, EXPLAIN, transaction control — never mutates.
    Read,
    /// INSERT/UPDATE/DELETE that the planner can EXPLAIN and that can be
    /// previewed with a RETURNING wrapper.
    Dml,
    /// Transactional DDL (CREATE TABLE/VIEW, ALTER, DROP, TRUNCATE, …). Runs
    /// inside a transaction so it can be dry-run, but EXPLAIN doesn't apply.
    Ddl,
    /// Statement that cannot run inside a transaction block (CREATE INDEX
    /// CONCURRENTLY, VACUUM). Dry-run is impossible; requires `--commit`.
    NonTransactional,
}

impl StmtKind {
    fn is_write(self) -> bool {
        !matches!(self, StmtKind::Read)
    }
}

/// Statement-level facts gathered from a single parse of the input.
struct ParsedSql {
    kinds: Vec<StmtKind>,
    /// `true` when the input is exactly one DML statement with no RETURNING
    /// clause — the only case where we wrap with RETURNING to sample rows.
    single_dml_no_returning: bool,
}

impl ParsedSql {
    fn has_write(&self) -> bool {
        self.kinds.iter().any(|k| k.is_write())
    }

    fn non_transactional(&self) -> bool {
        self.kinds
            .iter()
            .any(|k| matches!(k, StmtKind::NonTransactional))
    }

    /// Statements eligible for EXPLAIN cost gating (only DML; DDL/utility
    /// statements aren't plannable).
    fn explainable(&self) -> bool {
        self.kinds.iter().any(|k| matches!(k, StmtKind::Dml))
    }
}

pub async fn sql(database_url: &str, command: Option<&str>, opts: SqlOptions) -> Result<()> {
    let raw = match command {
        Some(c) => c.to_string(),
        None => {
            let mut buf = String::new();
            std::io::stdin()
                .read_to_string(&mut buf)
                .context("read SQL from stdin")?;
            buf
        }
    };

    let sql = raw.trim();
    if sql.is_empty() {
        bail!(
            "No SQL provided. Use: pgcrate sql -c \"SELECT 1\" or echo \"SELECT 1\" | pgcrate sql"
        );
    }

    let parsed = classify(sql)?;
    let is_write = parsed.has_write();

    // Gate: a write with neither --allow-write nor --commit is refused, and the
    // message teaches the new preview-by-default model.
    if is_write && !opts.allow_write {
        bail!(
            "This SQL writes. By default writes are blocked.\n  \
             Preview it (runs in a transaction, then rolls back):\n    \
             pgcrate sql -c \"…\" --allow-write\n  \
             Apply it for real:\n    \
             pgcrate sql -c \"…\" --commit"
        );
    }

    // Non-transactional writes can't be dry-run. Under --allow-write (preview)
    // we refuse and point at --commit; under --commit we run them directly.
    if parsed.non_transactional() && !opts.commit {
        bail!(
            "This SQL contains a statement that cannot run inside a transaction \
             (e.g. CREATE INDEX CONCURRENTLY, VACUUM), so it cannot be previewed.\n  \
             Re-run with --commit to execute it directly:\n    \
             pgcrate sql -c \"…\" --commit"
        );
    }

    let session = DiagnosticSession::connect(database_url, opts.timeouts.clone()).await?;
    crate::diagnostic::setup_ctrlc_handler(session.cancel_token());
    let client = session.client();

    // Cost gating runs before any execution. EXPLAIN only applies to DML.
    let cost = if is_write && !opts.no_cost_check && parsed.explainable() {
        cost_gate(client, sql, opts.cost_warn_threshold).await?
    } else {
        None
    };

    if is_write {
        run_write(client, sql, &parsed, &opts, cost).await
    } else {
        run_read(client, sql, &opts).await
    }
}

/// Read-only path: execute and print/return rows, capping the result set.
async fn run_read(client: &Client, sql: &str, opts: &SqlOptions) -> Result<()> {
    let messages = client.simple_query(sql).await.context("execute SQL")?;
    let cap = row_cap(opts.limit);
    let results = collect_results(messages, cap);

    if opts.json {
        emit_json(opts, None, results);
        return Ok(());
    }
    if opts.quiet {
        return Ok(());
    }
    print_results(&results, opts.density);
    Ok(())
}

/// Write path: dry-run (transaction → report → ROLLBACK) unless `--commit`.
async fn run_write(
    client: &Client,
    sql: &str,
    parsed: &ParsedSql,
    opts: &SqlOptions,
    cost: Option<CostEstimate>,
) -> Result<()> {
    // Non-transactional + --commit: execute directly, no transaction wrapper.
    if parsed.non_transactional() {
        let messages = client.simple_query(sql).await.context("execute SQL")?;
        let results = collect_results(messages, usize::MAX);
        let rows_affected = total_affected(&results);
        finish_write(opts, true, rows_affected, results, cost);
        return Ok(());
    }

    client
        .batch_execute("BEGIN")
        .await
        .context("begin transaction")?;

    // From here on, ensure we always close the transaction even on error. Any
    // error inside the transaction must roll back and propagate loudly — a
    // swallowed abort that still reports success is the worst failure for a
    // trust tool (see PGC-104).
    let outcome = execute_in_tx(client, sql, parsed, opts).await;

    let (results, rows_affected) = match outcome {
        Ok(v) => v,
        Err(e) => {
            let _ = client.batch_execute("ROLLBACK").await;
            return Err(rolled_back(e));
        }
    };

    if opts.commit {
        // Force any DEFERRED constraint checks to run now, while we can still
        // observe a violation as an error. Without this, a deferred violation
        // surfaces only at COMMIT, where Postgres silently downgrades COMMIT to
        // ROLLBACK and returns the `ROLLBACK` tag with no protocol error — the
        // exact shape that let an aborted transaction report COMMITTED.
        if let Err(e) = client.simple_query("SET CONSTRAINTS ALL IMMEDIATE").await {
            let _ = client.batch_execute("ROLLBACK").await;
            return Err(rolled_back(anyhow::Error::new(e)));
        }

        client
            .batch_execute("COMMIT")
            .await
            .context("commit transaction")?;
        finish_write(opts, true, rows_affected, results, cost);
    } else {
        client
            .batch_execute("ROLLBACK")
            .await
            .context("rollback transaction")?;
        finish_write(opts, false, rows_affected, results, cost);
    }

    Ok(())
}

/// Wrap a transaction error in the explicit, human-facing "ROLLED BACK —
/// nothing changed" framing. The chained source carries the original DB error;
/// `main` propagates this to exit code 10 for both human and JSON output.
fn rolled_back(e: anyhow::Error) -> anyhow::Error {
    e.context("ROLLED BACK — nothing changed")
}

/// Run the write inside the open transaction **exactly once**, gathering the
/// affected count and (when safe) a sample of affected rows.
///
/// The single-DML-without-RETURNING case is the dangerous one: previously it
/// ran the bare statement for the count *and* a RETURNING-wrapped copy for the
/// sample, double-applying every effect in the same transaction (PGC-104). Now
/// the wrap is the *only* execution — it yields both count and sample in one
/// pass — so each user statement executes once.
async fn execute_in_tx(
    client: &Client,
    sql: &str,
    parsed: &ParsedSql,
    opts: &SqlOptions,
) -> Result<(Vec<SqlResult>, u64)> {
    // Single DML without its own RETURNING: one execution that returns both the
    // true affected count and a bounded sample of affected rows.
    if parsed.single_dml_no_returning {
        return execute_dml_with_sample(client, sql, opts).await;
    }

    // Everything else (DML with RETURNING, DDL, multi-statement, writable CTEs):
    // one plain execution. No wrapping, no second pass.
    let messages = client.simple_query(sql).await.context("execute SQL")?;
    let results = collect_results(messages, usize::MAX);
    let rows_affected = total_affected(&results);
    Ok((results, rows_affected))
}

/// Execute a single INSERT/UPDATE/DELETE (no RETURNING) exactly once, deriving
/// the affected count and a bounded sample from one CTE-wrapped query.
///
/// `WITH x AS (<stmt> RETURNING *) SELECT (SELECT count(*) FROM x) AS affected,
///  (SELECT json_agg(t) FROM (SELECT * FROM x LIMIT N) t) AS samples`
///
/// The statement runs once inside the writable CTE; `count(*)` over `x` gives
/// the exact affected count and `json_agg` over a bounded slice gives the
/// sample. Any DB error (e.g. a unique violation) propagates from this one
/// call, so an aborted transaction can never be reported as COMMITTED.
async fn execute_dml_with_sample(
    client: &Client,
    sql: &str,
    opts: &SqlOptions,
) -> Result<(Vec<SqlResult>, u64)> {
    let stmt = sql.trim().trim_end_matches(';');
    let limit = match opts.limit {
        Some(0) => DRY_RUN_SAMPLE_LIMIT, // uncapped reads still bound the preview
        Some(n) => n.min(DRY_RUN_SAMPLE_LIMIT),
        None => DRY_RUN_SAMPLE_LIMIT,
    };
    // `row_to_json` (text `json`, not `jsonb`) preserves the statement's column
    // order; `jsonb` would re-sort keys and scramble the preview. `json_agg`
    // collects the slice; the keys array carries the column order explicitly so
    // Rust never has to infer it from a sorted map.
    let wrapped = format!(
        "WITH __pgcrate_w AS ({stmt} RETURNING *), \
              __pgcrate_s AS (SELECT * FROM __pgcrate_w LIMIT {limit}), \
              __pgcrate_first AS (SELECT * FROM __pgcrate_s LIMIT 1) \
         SELECT (SELECT count(*) FROM __pgcrate_w) AS __pgcrate_affected, \
                (SELECT json_agg(row_to_json(__s)) FROM __pgcrate_s __s) AS __pgcrate_samples, \
                (SELECT array_agg(k) FROM ( \
                    SELECT json_object_keys(row_to_json(__f)) AS k \
                    FROM __pgcrate_first __f \
                 ) __k) AS __pgcrate_columns"
    );

    let row = client
        .query_one(&wrapped, &[])
        .await
        .context("execute SQL")?;

    let affected: i64 = row.get("__pgcrate_affected");
    let rows_affected = affected.max(0) as u64;

    let columns: Option<Vec<String>> = row.get("__pgcrate_columns");

    let mut results = vec![SqlResult::CommandComplete {
        rows: rows_affected,
    }];
    if let Some(sample) = sample_from_json(row.get("__pgcrate_samples"), columns) {
        results.push(sample);
    }
    Ok((results, rows_affected))
}

/// Turn the `json_agg(row_to_json(...))` sample column into a `Sample` result.
/// `columns` carries the statement's column order (from `array_agg` of
/// `json_object_keys`); cells are stringified to match the simple-query text
/// representation used elsewhere.
fn sample_from_json(
    samples: Option<serde_json::Value>,
    columns: Option<Vec<String>>,
) -> Option<SqlResult> {
    let rows = samples?.as_array()?.clone();
    // Prefer the SQL-provided order; fall back to the first row's keys only if
    // it's somehow absent (e.g. all-NULL column array).
    let columns = columns
        .filter(|c| !c.is_empty())
        .or_else(|| Some(rows.first()?.as_object()?.keys().cloned().collect()))?;

    let table_rows: Vec<Vec<Option<String>>> = rows
        .iter()
        .filter_map(|r| r.as_object())
        .map(|obj| {
            columns
                .iter()
                .map(|col| match obj.get(col) {
                    None | Some(serde_json::Value::Null) => None,
                    Some(serde_json::Value::String(s)) => Some(s.clone()),
                    Some(other) => Some(other.to_string()),
                })
                .collect()
        })
        .collect();

    Some(SqlResult::Sample {
        columns,
        rows: table_rows,
    })
}

/// EXPLAIN the input and return the estimated total cost (relative to the
/// configured threshold). Returns `None` only when no cost could be computed
/// (multi-statement, non-plannable, or EXPLAIN failed) — never blocks.
async fn cost_gate(client: &Client, sql: &str, threshold: f64) -> Result<Option<CostEstimate>> {
    let stmt = sql.trim().trim_end_matches(';');
    // Only a single statement is reliably EXPLAIN-able here; multi-statement
    // writes skip the gate (the parser already flagged them as a write).
    if stmt.contains(';') {
        return Ok(None);
    }

    let explain_sql = format!("EXPLAIN (FORMAT JSON) {stmt}");
    let row = match client.query_one(&explain_sql, &[]).await {
        Ok(r) => r,
        // EXPLAIN can fail (non-plannable statement, permissions). Don't block.
        Err(_) => return Ok(None),
    };

    let plan_json: serde_json::Value = row.get(0);
    let total_cost = plan_json
        .get(0)
        .and_then(|p| p.get("Plan"))
        .and_then(|p| p.get("Total Cost"))
        .and_then(|c| c.as_f64());

    Ok(total_cost.map(|estimated| CostEstimate {
        estimated,
        threshold,
        exceeds_threshold: estimated > threshold,
    }))
}

/// Resolve the effective row cap. `Some(0)` means uncapped.
fn row_cap(limit: Option<usize>) -> usize {
    match limit {
        Some(0) => usize::MAX,
        Some(n) => n,
        None => DEFAULT_ROW_CAP,
    }
}

/// Sum affected-row counts across command-complete results.
fn total_affected(results: &[SqlResult]) -> u64 {
    results
        .iter()
        .filter_map(|r| match r {
            SqlResult::CommandComplete { rows } => Some(*rows),
            _ => None,
        })
        .sum()
}

/// Drive output for a completed write (human or JSON).
fn finish_write(
    opts: &SqlOptions,
    committed: bool,
    rows_affected: u64,
    results: Vec<SqlResult>,
    cost: Option<CostEstimate>,
) {
    // Humans get a stderr warning only when the estimate exceeds the threshold;
    // JSON consumers get the full structured cost inside the write object below.
    if !opts.json {
        if let Some(c) = &cost {
            if c.exceeds_threshold {
                eprintln!(
                    "pgcrate: warning: estimated query cost {:.0} exceeds threshold {:.0} \
                     — this may be expensive. Pass --no-cost-check to skip this check.",
                    c.estimated, c.threshold
                );
            }
        }
    }

    if opts.json {
        let write = Some(WriteOutcome {
            committed,
            rows_affected,
            cost,
        });
        emit_json(opts, write, results);
        return;
    }

    if opts.quiet {
        return;
    }

    // Print any sample rows first, then the summary banner.
    print_results(&results, opts.density);

    if committed {
        println!("\nCOMMITTED — {rows_affected} row(s) affected.");
    } else {
        println!(
            "\nDRY RUN — {rows_affected} row(s) would be affected. \
             Nothing was changed (rolled back)."
        );
        println!("Re-run with --commit to apply.");
    }
}

/// Convert simple-query messages into structured results, capping query rows.
fn collect_results(messages: Vec<SimpleQueryMessage>, cap: usize) -> Vec<SqlResult> {
    let mut results: Vec<SqlResult> = Vec::new();
    let mut current_columns: Option<Vec<String>> = None;
    let mut current_rows: Vec<Vec<Option<String>>> = Vec::new();
    let mut current_total: usize = 0;

    for msg in messages {
        match msg {
            SimpleQueryMessage::RowDescription(cols) => {
                current_columns = Some(cols.iter().map(|c| c.name().to_string()).collect());
            }
            SimpleQueryMessage::Row(row) => {
                if current_columns.is_none() {
                    current_columns =
                        Some(row.columns().iter().map(|c| c.name().to_string()).collect());
                }
                current_total += 1;
                if current_rows.len() < cap {
                    let values: Vec<Option<String>> = (0..row.len())
                        .map(|i| row.get(i).map(|s| s.to_string()))
                        .collect();
                    current_rows.push(values);
                }
            }
            SimpleQueryMessage::CommandComplete(rows) => {
                if let Some(cols) = current_columns.take() {
                    let truncated = current_total.saturating_sub(current_rows.len());
                    results.push(SqlResult::Query {
                        columns: cols,
                        rows: std::mem::take(&mut current_rows),
                        truncated,
                    });
                    current_total = 0;
                }
                results.push(SqlResult::CommandComplete { rows });
            }
            _ => {}
        }
    }

    if let Some(cols) = current_columns.take() {
        let truncated = current_total.saturating_sub(current_rows.len());
        results.push(SqlResult::Query {
            columns: cols,
            rows: std::mem::take(&mut current_rows),
            truncated,
        });
    }

    results
}

fn emit_json(opts: &SqlOptions, write: Option<WriteOutcome>, results: Vec<SqlResult>) {
    let payload = SqlResponse {
        ok: true,
        write,
        results,
    };
    match serde_json::to_string_pretty(&payload) {
        Ok(s) => println!("{s}"),
        Err(e) => eprintln!("pgcrate: failed to serialize JSON: {e}"),
    }
    let _ = opts;
}

fn print_results(results: &[SqlResult], density: crate::output::Density) {
    for result in results {
        match result {
            SqlResult::Query {
                columns,
                rows,
                truncated,
            } => {
                print_table(columns, rows, density);
                if *truncated > 0 {
                    println!("… +{truncated} more row(s) — use --limit N (or --limit 0 to uncap)");
                }
            }
            SqlResult::Sample { columns, rows } => {
                println!("Sample of affected rows:");
                print_table(columns, rows, density);
            }
            SqlResult::CommandComplete { rows } => {
                println!("OK ({rows} rows)");
            }
        }
    }
}

/// Statements that cannot run inside a transaction *and* that sqlparser 0.58
/// does not parse into a dedicated `Statement` variant. We detect these by
/// leading keyword so the user gets the "needs --commit" guidance rather than
/// a raw parse error.
fn leading_non_transactional(sql: &str) -> bool {
    let head = sql.trim_start().to_ascii_uppercase();
    head.starts_with("VACUUM") || head.starts_with("REINDEX")
}

/// Does a parsed `Query` (a `SELECT`-shaped statement) actually mutate data?
///
/// sqlparser routes several write constructs through `Statement::Query`, so a
/// naive `Query => Read` rule lets real writes execute unguarded. At any
/// nesting depth (CTEs can themselves contain `WITH`, set operations nest,
/// subqueries nest), this catches:
///
/// - writable CTEs (`WITH d AS (DELETE … RETURNING *) SELECT *`),
/// - leading-CTE DML (`WITH x AS (…) UPDATE …`, whose top-level body is the
///   `UPDATE`/`INSERT`/`DELETE`),
/// - `SELECT … INTO new_table` (creates a table).
fn query_writes(query: &sqlparser::ast::Query) -> bool {
    if let Some(with) = &query.with {
        if with.cte_tables.iter().any(|cte| query_writes(&cte.query)) {
            return true;
        }
    }
    set_expr_writes(&query.body)
}

fn set_expr_writes(body: &sqlparser::ast::SetExpr) -> bool {
    use sqlparser::ast::SetExpr;
    match body {
        // A DML body inside a query is a write (writable CTE / leading-CTE DML).
        SetExpr::Insert(_) | SetExpr::Update(_) | SetExpr::Delete(_) => true,
        // `SELECT … INTO t` materializes a table — that's a write.
        SetExpr::Select(select) => select.into.is_some(),
        SetExpr::Query(inner) => query_writes(inner),
        SetExpr::SetOperation { left, right, .. } => {
            set_expr_writes(left) || set_expr_writes(right)
        }
        SetExpr::Values(_) | SetExpr::Table(_) => false,
    }
}

/// Classify a single parsed statement. `Query` and `Explain` recurse because
/// either can hide a write (writable CTE; `EXPLAIN ANALYZE <write>`, which
/// *executes* the inner statement).
fn classify_statement(stmt: &sqlparser::ast::Statement) -> (StmtKind, bool) {
    use sqlparser::ast::Statement;
    match stmt {
        // A query is a write iff it embeds DML / SELECT INTO; otherwise read.
        Statement::Query(q) => {
            if query_writes(q) {
                // Embedded writes can't be safely RETURNING-wrapped, and the
                // planner can't EXPLAIN them as plain DML — treat as DDL-class
                // (transactional dry-run, no sample, no cost gate).
                (StmtKind::Ddl, false)
            } else {
                (StmtKind::Read, false)
            }
        }

        // EXPLAIN ANALYZE *runs* the inner statement; the parenthesized
        // `EXPLAIN (ANALYZE)` form parses with `analyze=false` in sqlparser yet
        // still executes, so we also classify the inner statement and treat any
        // inner write as a write regardless of the flag.
        Statement::Explain {
            statement, analyze, ..
        } => {
            let (inner_kind, _) = classify_statement(statement);
            if *analyze || inner_kind.is_write() {
                // EXPLAIN can't run inside our RETURNING wrapper or be cost-gated
                // meaningfully; route it through the transactional DDL path so a
                // dry-run still rolls back the side effects of ANALYZE.
                (StmtKind::Ddl, false)
            } else {
                (StmtKind::Read, false)
            }
        }

        Statement::Set(_)
        | Statement::StartTransaction { .. }
        | Statement::Commit { .. }
        | Statement::Rollback { .. }
        | Statement::ExplainTable { .. } => (StmtKind::Read, false),

        Statement::Insert(insert) => (StmtKind::Dml, insert.returning.is_some()),
        Statement::Update { returning, .. } => (StmtKind::Dml, returning.is_some()),
        Statement::Delete(delete) => (StmtKind::Dml, delete.returning.is_some()),

        // CREATE INDEX CONCURRENTLY cannot run in a transaction.
        Statement::CreateIndex(idx) if idx.concurrently => (StmtKind::NonTransactional, false),

        // Everything else that mutates: transactional DDL/utility.
        _ => (StmtKind::Ddl, false),
    }
}

/// Parse the input once and classify each statement.
fn classify(sql: &str) -> Result<ParsedSql> {
    // VACUUM / REINDEX aren't parseable by sqlparser; short-circuit so the
    // caller can route them to the "requires --commit" path.
    if leading_non_transactional(sql) {
        return Ok(ParsedSql {
            kinds: vec![StmtKind::NonTransactional],
            single_dml_no_returning: false,
        });
    }

    let dialect = sqlparser::dialect::PostgreSqlDialect {};
    let statements = sqlparser::parser::Parser::parse_sql(&dialect, sql).context("parse SQL")?;

    let mut kinds = Vec::with_capacity(statements.len());
    let mut dml_returnings: Vec<bool> = Vec::new();

    for stmt in &statements {
        let (kind, has_returning) = classify_statement(stmt);
        if kind == StmtKind::Dml {
            dml_returnings.push(has_returning);
        }
        kinds.push(kind);
    }

    let single_dml_no_returning = statements.len() == 1
        && kinds.first() == Some(&StmtKind::Dml)
        && dml_returnings.first() == Some(&false);

    Ok(ParsedSql {
        kinds,
        single_dml_no_returning,
    })
}

fn print_table(columns: &[String], rows: &[Vec<Option<String>>], density: crate::output::Density) {
    if columns.is_empty() {
        return;
    }

    if density.is_dense() {
        // Dense: header row plus unpadded ` | `-joined cells, no width padding
        // and no separator rule. Columns are still labelled and aligned by the
        // pipe delimiter, so the result stays unambiguous.
        println!("{}", columns.join(" | "));
        for row in rows {
            let line: Vec<&str> = columns
                .iter()
                .enumerate()
                .map(|(i, _)| row.get(i).and_then(|v| v.as_deref()).unwrap_or("NULL"))
                .collect();
            println!("{}", line.join(" | "));
        }
        return;
    }

    let mut widths: Vec<usize> = columns.iter().map(|c| c.len()).collect();
    for row in rows {
        for (i, cell) in row.iter().enumerate() {
            if i >= widths.len() {
                continue;
            }
            let s = cell.as_deref().unwrap_or("NULL");
            widths[i] = widths[i].max(s.len());
        }
    }

    let header: Vec<String> = columns
        .iter()
        .enumerate()
        .map(|(i, c)| format!("{:width$}", c, width = widths[i]))
        .collect();
    println!("{}", header.join(" | "));

    let sep: Vec<String> = widths.iter().map(|w| "-".repeat(*w)).collect();
    println!("{}", sep.join("-+-"));

    for row in rows {
        let line: Vec<String> = columns
            .iter()
            .enumerate()
            .map(|(i, _)| {
                let s = row.get(i).and_then(|v| v.as_deref()).unwrap_or("NULL");
                format!("{:width$}", s, width = widths[i])
            })
            .collect();
        println!("{}", line.join(" | "));
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn select_is_read() {
        let p = classify("SELECT 1").unwrap();
        assert!(!p.has_write());
        assert!(!p.non_transactional());
    }

    #[test]
    fn update_is_write_and_samplable() {
        let p = classify("UPDATE users SET name = 'x' WHERE id = 1").unwrap();
        assert!(p.has_write());
        assert!(p.explainable());
        assert!(p.single_dml_no_returning);
    }

    #[test]
    fn update_with_returning_not_wrapped() {
        let p = classify("UPDATE users SET name = 'x' WHERE id = 1 RETURNING id").unwrap();
        assert!(p.has_write());
        assert!(!p.single_dml_no_returning);
    }

    #[test]
    fn insert_delete_are_dml() {
        assert!(
            classify("INSERT INTO t (a) VALUES (1)")
                .unwrap()
                .single_dml_no_returning
        );
        assert!(
            classify("DELETE FROM t WHERE a = 1")
                .unwrap()
                .single_dml_no_returning
        );
    }

    #[test]
    fn multi_statement_write_not_single() {
        let p = classify("UPDATE t SET a = 1; UPDATE t SET a = 2").unwrap();
        assert!(p.has_write());
        assert!(!p.single_dml_no_returning);
    }

    #[test]
    fn ddl_is_write_not_explainable() {
        let p = classify("CREATE TABLE t (id int)").unwrap();
        assert!(p.has_write());
        assert!(!p.explainable());
        assert!(!p.non_transactional());
    }

    #[test]
    fn drop_and_truncate_are_writes() {
        assert!(classify("DROP TABLE t").unwrap().has_write());
        assert!(classify("TRUNCATE t").unwrap().has_write());
        assert!(classify("ALTER TABLE t ADD COLUMN c int")
            .unwrap()
            .has_write());
    }

    #[test]
    fn create_index_concurrently_is_non_transactional() {
        let p = classify("CREATE INDEX CONCURRENTLY idx ON t (a)").unwrap();
        assert!(p.has_write());
        assert!(p.non_transactional());
    }

    #[test]
    fn plain_create_index_is_transactional() {
        let p = classify("CREATE INDEX idx ON t (a)").unwrap();
        assert!(p.has_write());
        assert!(!p.non_transactional());
    }

    #[test]
    fn vacuum_is_non_transactional() {
        assert!(classify("VACUUM ANALYZE t").unwrap().non_transactional());
    }

    #[test]
    fn row_cap_resolution() {
        assert_eq!(row_cap(None), DEFAULT_ROW_CAP);
        assert_eq!(row_cap(Some(0)), usize::MAX);
        assert_eq!(row_cap(Some(10)), 10);
    }

    // --- Bypass regression tests: writes that hide inside a Query/Explain. ---
    // Each of these previously classified as Read and, under --allow-write,
    // executed in autocommit with no transaction wrapper and no rollback.

    #[test]
    fn writable_cte_delete_is_write() {
        let p = classify("WITH d AS (DELETE FROM t RETURNING *) SELECT * FROM d").unwrap();
        assert!(
            p.has_write(),
            "writable-CTE DELETE must classify as a write"
        );
        // Not wrappable/explainable — the embedded write goes the DDL dry-run path.
        assert!(!p.single_dml_no_returning);
        assert!(!p.explainable());
    }

    #[test]
    fn writable_cte_update_and_insert_are_writes() {
        assert!(
            classify("WITH u AS (UPDATE t SET a = 1 RETURNING *) SELECT * FROM u")
                .unwrap()
                .has_write()
        );
        assert!(
            classify("WITH i AS (INSERT INTO t VALUES (1) RETURNING *) SELECT * FROM i")
                .unwrap()
                .has_write()
        );
    }

    #[test]
    fn leading_cte_dml_is_write() {
        // Top-level body is the UPDATE; sqlparser still wraps it in Statement::Query.
        let p = classify("WITH x AS (SELECT 1) UPDATE t SET a = 1").unwrap();
        assert!(p.has_write(), "leading-CTE UPDATE must classify as a write");
    }

    #[test]
    fn nested_writable_cte_is_write() {
        // A DELETE buried inside a CTE's own WITH clause.
        let sql = "WITH outer1 AS (\
                       WITH inner1 AS (DELETE FROM t RETURNING *) SELECT * FROM inner1\
                   ) SELECT * FROM outer1";
        assert!(
            classify(sql).unwrap().has_write(),
            "DELETE nested in a CTE's WITH must classify as a write"
        );
    }

    #[test]
    fn select_into_is_write() {
        let p = classify("SELECT id INTO new_table FROM t").unwrap();
        assert!(
            p.has_write(),
            "SELECT … INTO materializes a table — a write"
        );
    }

    #[test]
    fn explain_analyze_write_is_write() {
        // EXPLAIN ANALYZE executes the inner statement.
        assert!(classify("EXPLAIN ANALYZE DELETE FROM t")
            .unwrap()
            .has_write());
        assert!(classify("EXPLAIN ANALYZE UPDATE t SET a = 1")
            .unwrap()
            .has_write());
        // Parenthesized form parses analyze=false in sqlparser but still executes;
        // the inner-statement check catches it.
        assert!(classify("EXPLAIN (ANALYZE) DELETE FROM t")
            .unwrap()
            .has_write());
        // Plain EXPLAIN of a write doesn't execute, but classifying it as a write
        // is the conservative, correct call.
        assert!(classify("EXPLAIN DELETE FROM t").unwrap().has_write());
    }

    #[test]
    fn plain_reads_stay_reads() {
        // Guard against over-eager classification breaking ordinary queries.
        assert!(!classify("SELECT 1").unwrap().has_write());
        assert!(!classify("EXPLAIN SELECT 1").unwrap().has_write());
        assert!(!classify("EXPLAIN (FORMAT JSON) SELECT 1")
            .unwrap()
            .has_write());
        assert!(!classify("WITH q AS (SELECT 1) SELECT * FROM q")
            .unwrap()
            .has_write());
        assert!(!classify("SELECT * FROM (SELECT 1) sub")
            .unwrap()
            .has_write());
        assert!(!classify("(SELECT 1) UNION (SELECT 2)").unwrap().has_write());
        assert!(!classify("VALUES (1), (2)").unwrap().has_write());
    }
}
