pg-api 0.3.5

use deadpool_postgres::{Config, Pool, Runtime};
use serde_json::Value;
use tokio_postgres::{NoTls, Row};
use tokio_postgres::types::{FromSql, Type};
use uuid::Uuid;

use crate::models::{
    Account, AccountRole, AppState, DatabaseAccess, FieldMetadata, Permission, QueryOptions,
    QueryResult,
};

/// Wrapper to extract raw binary bytes from any PostgreSQL column type.
/// Used as a fallback when no native Rust type mapping exists (e.g. NUMERIC).
struct RawBytes(Vec<u8>);

impl<'a> FromSql<'a> for RawBytes {
    fn from_sql(_ty: &Type, raw: &'a [u8]) -> Result<Self, Box<dyn std::error::Error + Sync + Send>> {
        Ok(RawBytes(raw.to_vec()))
    }

    fn accepts(_ty: &Type) -> bool {
        true
    }
}

/// Parse PostgreSQL NUMERIC binary format to a decimal string.
///
/// Binary format: [ndigits:i16][weight:i16][sign:u16][dscale:u16][digits:i16*ndigits]
/// Each digit group represents a base-10000 value (0..9999).
/// The value is: sign * sum(digit[i] * 10000^(weight - i)) for i in 0..ndigits
fn numeric_binary_to_string(raw: &[u8]) -> Option<String> {
    if raw.len() < 8 {
        return None;
    }

    let ndigits = i16::from_be_bytes([raw[0], raw[1]]) as usize;
    let weight = i16::from_be_bytes([raw[2], raw[3]]);
    let sign = u16::from_be_bytes([raw[4], raw[5]]);
    let dscale = u16::from_be_bytes([raw[6], raw[7]]) as usize;

    if raw.len() != 8 + ndigits * 2 {
        return None;
    }

    // NaN
    if sign == 0xC000 {
        return Some("NaN".to_string());
    }

    let is_negative = sign == 0x4000;

    let mut digits = Vec::with_capacity(ndigits);
    for i in 0..ndigits {
        let offset = 8 + i * 2;
        digits.push(i16::from_be_bytes([raw[offset], raw[offset + 1]]));
    }

    let mut result = String::new();
    if is_negative {
        result.push('-');
    }

    let int_groups = (weight + 1).max(0) as usize;

    if ndigits == 0 {
        result.push('0');
    } else {
        // Integer part
        for i in 0..int_groups {
            let d = if i < ndigits { digits[i] } else { 0 };
            if i == 0 {
                result.push_str(&d.to_string());
            } else {
                result.push_str(&format!("{:04}", d));
            }
        }
        if int_groups == 0 {
            result.push('0');
        }
    }

    // Decimal part
    if dscale > 0 {
        result.push('.');
        let mut dec = String::new();
        for i in int_groups..ndigits {
            dec.push_str(&format!("{:04}", digits[i]));
        }
        while dec.len() < dscale {
            dec.push('0');
        }
        if dec.len() > dscale {
            dec.truncate(dscale);
        }
        result.push_str(&dec);
    }

    Some(result)
}

pub async fn get_or_create_pool(
    state: &AppState,
    pool_key: &str,
    account: &Account,
    db_access: &DatabaseAccess,
) -> Result<Pool, String> {
    if let Some(pool) = state.connections.get(pool_key) {
        return Ok(pool.clone());
    }
    
    let instances = state.instances.read().await;
    let instance = instances.get(&account.instance_id)
        .ok_or("Instance not found")?;
    
    let mut config = Config::new();
    config.host = Some(instance.host.clone());
    config.port = Some(instance.port);
    config.user = Some(db_access.username.clone());
    config.password = Some(db_access.password.clone());
    config.dbname = Some(db_access.database.clone());
    
    let pool = config.create_pool(Some(Runtime::Tokio1), NoTls)
        .map_err(|e| e.to_string())?;
    
    state.connections.insert(pool_key.to_string(), pool.clone());
    
    Ok(pool)
}

/// v0.3 query gate: literal-aware validation replacing the old
/// first-keyword prefix check.
///
/// Guarantees (best-effort, fail-closed on ambiguity):
/// - exactly one statement (`;` outside literals is rejected);
/// - comments are stripped before classification, so `/* x */ DROP`
///   or `-- hi\nDELETE` cannot smuggle a different statement class;
/// - `WITH` bodies are scanned: a CTE performing writes requires the
///   matching write permission (agreed v0.3 behavior);
/// - dangerous statement types (`COPY`, transaction control, `VACUUM`,
///   `CALL/DO`, `EXPLAIN ANALYZE`, …) and dangerous functions
///   (`pg_read_file`, `dblink`, large-object import/export, …) are denied
///   for every role except `Superuser`.
///
/// Error strings carry a machine-readable prefix (`MULTI_STATEMENT`,
/// `STATEMENT_NOT_ALLOWED`, `MISSING_PERMISSION`) mapped to API codes by
/// the handlers.
pub fn validate_query_permissions(
    query: &str,
    permissions: &[Permission],
    role: &AccountRole,
) -> Result<(), String> {
    if *role == AccountRole::Superuser {
        return Ok(());
    }

    if query.len() > 1_000_000 {
        return Err("STATEMENT_NOT_ALLOWED: query exceeds 1MB limit".to_string());
    }

    // 1. Strip comments and split top-level statements (literal-aware).
    let code = strip_sql_comments(query);
    let statements: Vec<&str> = split_statements(&code)
        .into_iter()
        .map(str::trim)
        .filter(|s| !s.is_empty())
        .collect();
    if statements.len() != 1 {
        return Err(
            "MULTI_STATEMENT: only a single statement per request is allowed".to_string(),
        );
    }
    let stmt = statements[0];

    // 2. Deny-list of dangerous functions, checked on code with string
    // literals blanked so that `'dblink'` as a value cannot cause a false
    // denial, while real calls cannot hide inside literals.
    let blanked = blank_string_literals(stmt);
    let bupper = blanked.to_uppercase();
    for func in DENIED_FUNCTIONS {
        if contains_word(&bupper, func) {
            return Err(format!(
                "STATEMENT_NOT_ALLOWED: function '{func}' is not permitted through this API"
            ));
        }
    }

    // 3. Classify the statement and collect required permissions.
    // Statement types that are never proxied (transaction control, DDL
    // outside the permission model, session commands, …) are rejected by
    // first keyword — not by global scan, so that e.g. every UPDATE
    // (which contains the word SET) and innocent column names keep working.
    let keyword = first_keyword(stmt);
    if DENIED_FIRST_KEYWORDS.contains(&keyword.as_str()) {
        return Err(format!(
            "STATEMENT_NOT_ALLOWED: '{keyword}' is not permitted through this API"
        ));
    }
    let mut required: Vec<Permission> = Vec::new();
    match keyword.as_str() {
        "SELECT" | "VALUES" | "TABLE" | "SHOW" => {
            required.push(Permission::Select);
        }
        "INSERT" => required.push(Permission::Insert),
        "UPDATE" => required.push(Permission::Update),
        "DELETE" => required.push(Permission::Delete),
        "MERGE" => {
            required.push(Permission::Insert);
            required.push(Permission::Update);
            required.push(Permission::Delete);
        }
        "CREATE" => required.push(Permission::Create),
        "DROP" => required.push(Permission::Drop),
        "TRUNCATE" => required.push(Permission::Truncate),
        "EXPLAIN" => {
            if contains_word(&bupper, "ANALYZE") {
                // EXPLAIN ANALYZE really executes the statement.
                return Err(
                    "STATEMENT_NOT_ALLOWED: use EXPLAIN without ANALYZE".to_string(),
                );
            }
            required.push(Permission::Select);
        }
        "WITH" => {
            // Data-modifying CTEs require the matching write permission.
            if contains_word(&bupper, "INSERT") {
                required.push(Permission::Insert);
            }
            if contains_word(&bupper, "UPDATE") {
                required.push(Permission::Update);
            }
            if contains_word(&bupper, "DELETE") {
                required.push(Permission::Delete);
            }
            if contains_word(&bupper, "MERGE") {
                required.push(Permission::Insert);
                required.push(Permission::Update);
                required.push(Permission::Delete);
            }
            required.push(Permission::Select);
        }
        _ => {
            return Err(format!(
                "STATEMENT_NOT_ALLOWED: unsupported statement '{keyword}'"
            ));
        }
    }

    // 4. Authorize.
    if permissions.contains(&Permission::All) {
        return Ok(());
    }
    let missing: Vec<String> = required
        .iter()
        .filter(|p| !permissions.contains(p))
        .map(|p| format!("{p:?}"))
        .collect();
    if missing.is_empty() {
        Ok(())
    } else {
        Err(format!("MISSING_PERMISSION: missing permissions: {}", missing.join(", ")))
    }
}

/// Statement types rejected by first keyword for all non-superuser roles.
/// Checked positionally (not by global scan) so that e.g. `UPDATE … SET …`
/// and innocent identifiers keep working.
const DENIED_FIRST_KEYWORDS: &[&str] = &[
    "COPY", "GRANT", "REVOKE", "ALTER", "SET", "RESET", "LISTEN", "NOTIFY", "UNLISTEN", "VACUUM",
    "ANALYZE", "CLUSTER", "REINDEX", "DISCARD", "CHECKPOINT", "CALL", "DO", "EXECUTE", "PREPARE",
    "DEALLOCATE", "BEGIN", "START", "COMMIT", "ROLLBACK", "SAVEPOINT", "LOCK", "COMMENT",
    "SECURITY", "LOAD", "IMPORT",
];

/// Functions rejected for all non-superuser roles (word-boundary match).
const DENIED_FUNCTIONS: &[&str] = &[
    "PG_READ_FILE",
    "PG_READ_BINARY_FILE",
    "PG_LS_DIR",
    "PG_STAT_FILE",
    "DBLINK",
    "DBLINK_EXEC",
    "DBLINK_CONNECT",
    "LO_IMPORT",
    "LO_EXPORT",
    "PG_EXECUTE_FROM_FILE",
];

/// Remove `--` line comments and `/* */` block comments, keeping string
/// literals (including dollar-quoted) intact.
fn strip_sql_comments(q: &str) -> String {
    let bytes = q.as_bytes();
    let mut out = String::with_capacity(q.len());
    let mut i = 0;
    let mut state = ScanState::Code;
    let mut dollar_tag = String::new();

    while i < bytes.len() {
        let c = bytes[i] as char;
        match state {
            ScanState::Code => {
                if c == '-' && peek(bytes, i + 1) == Some('-') {
                    state = ScanState::LineComment;
                    i += 2;
                } else if c == '/' && peek(bytes, i + 1) == Some('*') {
                    state = ScanState::BlockComment(1);
                    i += 2;
                } else if c == '\'' {
                    state = ScanState::SingleQuote;
                    out.push(c);
                    i += 1;
                } else if c == '"' {
                    state = ScanState::DoubleQuote;
                    out.push(c);
                    i += 1;
                } else if c == '$' {
                    if let Some(tag) = match_dollar_tag(bytes, i) {
                        dollar_tag = tag.clone();
                        state = ScanState::DollarQuote;
                        out.push_str(&tag);
                        i += tag.len();
                    } else {
                        // Positional parameter ($1) — keep as code.
                        out.push(c);
                        i += 1;
                    }
                } else {
                    out.push(c);
                    i += 1;
                }
            }
            ScanState::LineComment => {
                if c == '\n' {
                    state = ScanState::Code;
                    out.push(c);
                }
                i += 1;
            }
            ScanState::BlockComment(depth) => {
                if c == '/' && peek(bytes, i + 1) == Some('*') {
                    state = ScanState::BlockComment(depth + 1);
                    i += 2;
                } else if c == '*' && peek(bytes, i + 1) == Some('/') {
                    if depth == 1 {
                        state = ScanState::Code;
                    } else {
                        state = ScanState::BlockComment(depth - 1);
                    }
                    i += 2;
                } else {
                    i += 1;
                }
            }
            ScanState::SingleQuote => {
                out.push(c);
                if c == '\'' {
                    if peek(bytes, i + 1) == Some('\'') {
                        out.push('\'');
                        i += 2;
                    } else {
                        state = ScanState::Code;
                        i += 1;
                    }
                } else {
                    i += 1;
                }
            }
            ScanState::DoubleQuote => {
                out.push(c);
                if c == '"' {
                    if peek(bytes, i + 1) == Some('"') {
                        out.push('"');
                        i += 2;
                    } else {
                        state = ScanState::Code;
                        i += 1;
                    }
                } else {
                    i += 1;
                }
            }
            ScanState::DollarQuote => {
                if c == '$' {
                    if let Some(tag) = match_dollar_tag(bytes, i) {
                        if tag == dollar_tag {
                            out.push_str(&tag);
                            state = ScanState::Code;
                            i += tag.len();
                            continue;
                        }
                    }
                }
                out.push(c);
                i += 1;
            }
        }
    }
    out
}

#[derive(Clone, Copy)]
enum ScanState {
    Code,
    LineComment,
    BlockComment(u32),
    SingleQuote,
    DoubleQuote,
    DollarQuote,
}

fn peek(bytes: &[u8], i: usize) -> Option<char> {
    bytes.get(i).map(|b| *b as char)
}

/// If `bytes[i] == '$'`, try to match a dollar-quote tag `$tag$`.
/// Returns the full tag including both `$`. `$1` (parameters) yield None.
fn match_dollar_tag(bytes: &[u8], i: usize) -> Option<String> {
    let mut j = i + 1;
    while j < bytes.len() {
        let c = bytes[j] as char;
        if c == '$' {
            let tag: String = bytes[i..=j].iter().map(|b| *b as char).collect();
            let inner = &tag[1..tag.len() - 1];
            // A lone `$` or `$1` is not a quote tag.
            if inner.is_empty() {
                return Some(tag);
            }
            if inner.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'_')
                && !inner.bytes().next().is_some_and(|b| b.is_ascii_digit())
            {
                return Some(tag);
            }
            return None;
        }
        if !(c.is_ascii_alphanumeric() || c == '_') {
            return None;
        }
        j += 1;
        if j - i > 64 {
            return None;
        }
    }
    None
}

/// Split on `;` that appear outside literals/comments (comments already stripped).
fn split_statements(code: &str) -> Vec<&str> {
    let bytes = code.as_bytes();
    let mut parts = Vec::new();
    let mut start = 0;
    let mut i = 0;
    let mut state = ScanState::Code;
    let mut dollar_tag = String::new();

    while i < bytes.len() {
        let c = bytes[i] as char;
        match state {
            ScanState::Code => {
                if c == ';' {
                    parts.push(&code[start..i]);
                    start = i + 1;
                    i += 1;
                } else if c == '\'' {
                    state = ScanState::SingleQuote;
                    i += 1;
                } else if c == '"' {
                    state = ScanState::DoubleQuote;
                    i += 1;
                } else if c == '$' {
                    if let Some(tag) = match_dollar_tag(bytes, i) {
                        dollar_tag = tag.clone();
                        state = ScanState::DollarQuote;
                        i += tag.len();
                    } else {
                        i += 1;
                    }
                } else {
                    i += 1;
                }
            }
            ScanState::SingleQuote => {
                if c == '\'' {
                    if peek(bytes, i + 1) == Some('\'') {
                        i += 2;
                    } else {
                        state = ScanState::Code;
                        i += 1;
                    }
                } else {
                    i += 1;
                }
            }
            ScanState::DoubleQuote => {
                if c == '"' {
                    if peek(bytes, i + 1) == Some('"') {
                        i += 2;
                    } else {
                        state = ScanState::Code;
                        i += 1;
                    }
                } else {
                    i += 1;
                }
            }
            ScanState::DollarQuote => {
                if c == '$' {
                    if let Some(tag) = match_dollar_tag(bytes, i) {
                        if tag == dollar_tag {
                            state = ScanState::Code;
                            i += tag.len();
                            continue;
                        }
                    }
                }
                i += 1;
            }
            _ => {
                i += 1;
            }
        }
    }
    parts.push(&code[start..]);
    parts
}

/// Replace string-literal contents with spaces (keeps offsets, removes
/// false-positive denylist hits like `'DROP TABLE'` as a value).
fn blank_string_literals(stmt: &str) -> String {
    let bytes = stmt.as_bytes();
    let mut out = String::with_capacity(stmt.len());
    let mut i = 0;
    let mut state = ScanState::Code;
    let mut dollar_tag = String::new();

    while i < bytes.len() {
        let c = bytes[i] as char;
        match state {
            ScanState::Code => {
                if c == '\'' {
                    state = ScanState::SingleQuote;
                    out.push(' ');
                    i += 1;
                } else if c == '"' {
                    // Quoted identifiers can hide keywords ("drop"); blank them too.
                    state = ScanState::DoubleQuote;
                    out.push(' ');
                    i += 1;
                } else if c == '$' {
                    if let Some(tag) = match_dollar_tag(bytes, i) {
                        dollar_tag = tag.clone();
                        state = ScanState::DollarQuote;
                        for _ in 0..tag.len() {
                            out.push(' ');
                        }
                        i += tag.len();
                    } else {
                        out.push(c);
                        i += 1;
                    }
                } else {
                    out.push(c);
                    i += 1;
                }
            }
            ScanState::SingleQuote | ScanState::DoubleQuote => {
                let quote = if matches!(state, ScanState::SingleQuote) {
                    '\''
                } else {
                    '"'
                };
                if c == quote {
                    if peek(bytes, i + 1) == Some(quote) {
                        out.push_str("  ");
                        i += 2;
                    } else {
                        out.push(' ');
                        state = ScanState::Code;
                        i += 1;
                    }
                } else {
                    out.push(if c == '\n' { '\n' } else { ' ' });
                    i += 1;
                }
            }
            ScanState::DollarQuote => {
                if c == '$' {
                    if let Some(tag) = match_dollar_tag(bytes, i) {
                        if tag == dollar_tag {
                            for _ in 0..tag.len() {
                                out.push(' ');
                            }
                            state = ScanState::Code;
                            i += tag.len();
                            continue;
                        }
                    }
                }
                out.push(if c == '\n' { '\n' } else { ' ' });
                i += 1;
            }
            _ => {
                out.push(c);
                i += 1;
            }
        }
    }
    out
}

/// First SQL keyword, skipping leading whitespace and parentheses
/// (so `(SELECT …)` and `((WITH …))` classify correctly).
fn first_keyword(stmt: &str) -> String {
    let s = stmt
        .trim_start()
        .trim_start_matches(|c| c == '(')
        .trim_start();
    let end = s
        .find(|c: char| !c.is_ascii_alphanumeric() && c != '_')
        .unwrap_or(s.len());
    s[..end].to_uppercase()
}

/// Word-boundary substring search on already-uppercased text.
fn contains_word(haystack: &str, word: &str) -> bool {
    let h = haystack.as_bytes();
    let w = word.as_bytes();
    if w.len() > h.len() {
        return false;
    }
    let is_word = |c: u8| c.is_ascii_alphanumeric() || c == b'_' || c == b'$';
    let mut i = 0;
    while i + w.len() <= h.len() {
        if &h[i..i + w.len()] == w {
            let left_ok = i == 0 || !is_word(h[i - 1]);
            let right_ok = i + w.len() == h.len() || !is_word(h[i + w.len()]);
            if left_ok && right_ok {
                return true;
            }
        }
        i += 1;
    }
    false
}

pub fn rows_to_json(rows: &[Row]) -> Vec<Value> {
    rows.iter().map(|row| {
        let mut obj = serde_json::Map::new();
        
        for (i, column) in row.columns().iter().enumerate() {
            let value = match column.type_().name() {
                "bool" => row.try_get::<_, bool>(i).ok().map(Value::Bool),
                "int2" | "int4" => row.try_get::<_, i32>(i).ok().map(|v| Value::Number(v.into())),
                "int8" => row.try_get::<_, i64>(i).ok().map(|v| Value::Number(v.into())),
                "float4" => row.try_get::<_, f32>(i).ok().map(|v| Value::Number(serde_json::Number::from_f64(v as f64).unwrap_or_else(|| 0.into()))),
                "float8" => row.try_get::<_, f64>(i).ok().map(|v| Value::Number(serde_json::Number::from_f64(v).unwrap_or_else(|| 0.into()))),
                "text" | "varchar" | "char" | "bpchar" => row.try_get::<_, String>(i).ok().map(Value::String),
                "json" | "jsonb" => row.try_get::<_, Value>(i).ok(),
                // PostgreSQL timestamp has no timezone; timestamptz is decoded as DateTime<Utc>.
                "timestamp" => row.try_get::<_, chrono::NaiveDateTime>(i).ok().map(|dt| Value::String(dt.format("%Y-%m-%d %H:%M:%S").to_string())),
                "timestamptz" => row.try_get::<_, chrono::DateTime<chrono::Utc>>(i).ok().map(|dt| Value::String(dt.to_rfc3339())),
                "date" => row.try_get::<_, chrono::NaiveDate>(i).ok().map(|d| Value::String(d.format("%Y-%m-%d").to_string())),
                "uuid" => row.try_get::<_, Uuid>(i).ok().map(|u| Value::String(u.to_string())),
                "numeric" | "decimal" => {
                    row.try_get::<_, RawBytes>(i)
                        .ok()
                        .and_then(|raw| numeric_binary_to_string(&raw.0))
                        .map(Value::String)
                },
                "time" | "timetz" => row.try_get::<_, chrono::NaiveTime>(i).ok().map(|t| Value::String(t.format("%H:%M:%S").to_string())),
                _ => row.try_get::<_, String>(i).ok().map(Value::String),
            };
            
            obj.insert(column.name().to_string(), value.unwrap_or(Value::Null));
        }
        
        Value::Object(obj)
    }).collect()
}

/// Encode a JSON value as a SQL literal for safe interpolation (v0.3.2).
///
/// Why interpolation instead of bind parameters: the 0.1.x server
/// interpolated params into the query text, so clients were written
/// against tolerant semantics (extra params ignored, JSON integers work
/// against any numeric column, bare `$1` in SELECT lists). Strict binary
/// binding broke all of that (i64 vs int2/int4/float mismatch, unknown
/// type inference). Correct escaping below restores 0.1.x semantics
/// without its injection flaw:
/// - quote-doubling makes literal breakout impossible;
/// - `E''` is used when the value contains a backslash, so escaping is
///   well-defined regardless of the server's `standard_conforming_strings`;
/// - numbers/bools come from parsed JSON and cannot carry SQL text.
/// The v0.3 statement gate (single-statement, deny-lists) still applies
/// to the original query, on top.
fn json_to_sql_literal(value: &Value) -> Result<String, String> {
    match value {
        Value::Null => Ok("NULL".to_string()),
        Value::Bool(b) => Ok(if *b { "TRUE".to_string() } else { "FALSE".to_string() }),
        Value::Number(n) => {
            if let Some(i) = n.as_i64() {
                Ok(i.to_string())
            } else if let Some(u) = n.as_u64() {
                Ok(u.to_string())
            } else if let Some(f) = n.as_f64() {
                if !f.is_finite() {
                    return Err(
                        "STATEMENT_NOT_ALLOWED: non-finite number parameter".to_string(),
                    );
                }
                Ok(format!("{f:?}"))
            } else {
                Err("STATEMENT_NOT_ALLOWED: unsupported number parameter".to_string())
            }
        }
        Value::String(s) => Ok(quote_string_literal(s)),
        // Arrays/objects travel as their JSON spelling inside a literal;
        // Postgres coerces to json/jsonb/text targets as needed.
        Value::Array(_) | Value::Object(_) => Ok(quote_string_literal(&value.to_string())),
    }
}

fn quote_string_literal(s: &str) -> String {
    if s.contains('\\') {
        format!("E'{}'", s.replace('\\', "\\\\").replace('\'', "\\'"))
    } else {
        format!("'{}'", s.replace('\'', "''"))
    }
}

/// Substitute `$n` placeholders with safely-encoded literals.
///
/// - Only `$n` outside string literals/dollar-quotes/comments is replaced
///   (a literal `'$1'` or `R$` pricing text is left untouched), using the
///   same literal-aware scanner as the validator.
/// - Extra params are ignored (0.1.x compatibility); a referenced index
///   with no matching param is an error.
pub fn interpolate_query(query: &str, params: &[Value]) -> Result<String, String> {
    let bytes = query.as_bytes();
    let mut out = String::with_capacity(query.len() + params.len() * 8);
    let mut i = 0;
    let mut state = ScanState::Code;
    let mut dollar_tag = String::new();

    while i < bytes.len() {
        let c = bytes[i] as char;
        match state {
            ScanState::Code => {
                if c == '-' && peek(bytes, i + 1) == Some('-') {
                    state = ScanState::LineComment;
                    out.push_str("--");
                    i += 2;
                } else if c == '/' && peek(bytes, i + 1) == Some('*') {
                    state = ScanState::BlockComment(1);
                    out.push_str("/*");
                    i += 2;
                } else if c == '\'' {
                    state = ScanState::SingleQuote;
                    out.push(c);
                    i += 1;
                } else if c == '"' {
                    state = ScanState::DoubleQuote;
                    out.push(c);
                    i += 1;
                } else if c == '$' {
                    if let Some(tag) = match_dollar_tag(bytes, i) {
                        dollar_tag = tag.clone();
                        state = ScanState::DollarQuote;
                        out.push_str(&tag);
                        i += tag.len();
                    } else if let Some((n, len)) = match_positional_param(bytes, i) {
                        if n == 0 || n > params.len() {
                            return Err(format!(
                                "STATEMENT_NOT_ALLOWED: missing value for parameter ${n}"
                            ));
                        }
                        out.push_str(&json_to_sql_literal(&params[n - 1])?);
                        i += len;
                    } else {
                        out.push(c);
                        i += 1;
                    }
                } else {
                    out.push(c);
                    i += 1;
                }
            }
            ScanState::LineComment => {
                out.push(c);
                if c == '\n' {
                    state = ScanState::Code;
                }
                i += 1;
            }
            ScanState::BlockComment(depth) => {
                if c == '/' && peek(bytes, i + 1) == Some('*') {
                    state = ScanState::BlockComment(depth + 1);
                    out.push_str("/*");
                    i += 2;
                } else if c == '*' && peek(bytes, i + 1) == Some('/') {
                    if depth == 1 {
                        state = ScanState::Code;
                    } else {
                        state = ScanState::BlockComment(depth - 1);
                    }
                    out.push_str("*/");
                    i += 2;
                } else {
                    out.push(c);
                    i += 1;
                }
            }
            ScanState::SingleQuote => {
                out.push(c);
                if c == '\'' {
                    if peek(bytes, i + 1) == Some('\'') {
                        out.push('\'');
                        i += 2;
                    } else {
                        state = ScanState::Code;
                        i += 1;
                    }
                } else {
                    i += 1;
                }
            }
            ScanState::DoubleQuote => {
                out.push(c);
                if c == '"' {
                    if peek(bytes, i + 1) == Some('"') {
                        out.push('"');
                        i += 2;
                    } else {
                        state = ScanState::Code;
                        i += 1;
                    }
                } else {
                    i += 1;
                }
            }
            ScanState::DollarQuote => {
                if c == '$' {
                    if let Some(tag) = match_dollar_tag(bytes, i) {
                        if tag == dollar_tag {
                            out.push_str(&tag);
                            state = ScanState::Code;
                            i += tag.len();
                            continue;
                        }
                    }
                }
                out.push(c);
                i += 1;
            }
        }
    }
    Ok(out)
}

/// Match `$n` (n >= 1, decimal) at `bytes[i] == '$'`.
/// Returns (n, bytes consumed). `$` not followed by a digit is None
/// (`R$`, `$$`, `$tag$` are handled elsewhere).
fn match_positional_param(bytes: &[u8], i: usize) -> Option<(usize, usize)> {
    let mut j = i + 1;
    while j < bytes.len() && (bytes[j] as char).is_ascii_digit() {
        j += 1;
    }
    if j == i + 1 {
        return None;
    }
    let n: usize = std::str::from_utf8(&bytes[i + 1..j]).ok()?.parse().ok()?;
    Some((n, j - i))
}

/// Default per-query timeout applied when the client omits `timeout_ms`.
const DEFAULT_STATEMENT_TIMEOUT_MS: u64 = 60_000;
/// Hard ceiling: no single query may run longer than this.
const MAX_STATEMENT_TIMEOUT_MS: u64 = 300_000;

/// v0.3.4: escrita sem RETURNING deve reportar linhas afetadas reais
/// (antes: rows.len() == 0, pois nao ha linhas retornadas).
/// Deteccao literal-aware (reusa o scanner do gate): ignora a palavra
/// RETURNING dentro de literais e so casa INSERT/UPDATE/DELETE/MERGE
/// como primeira keyword.
fn is_plain_write(query: &str) -> bool {
    let blanked = blank_string_literals(query);
    let upper = blanked.to_uppercase();
    if !matches!(
        first_keyword(&upper).as_str(),
        "INSERT" | "UPDATE" | "DELETE" | "MERGE"
    ) {
        return false;
    }
    !contains_word(&upper, "RETURNING")
}

pub async fn execute_query_with_pool(
    pool: Pool,
    query: String,
    params: Vec<Value>,
    options: &QueryOptions,
) -> Result<QueryResult, String> {
    let client = pool.get().await.map_err(|e| e.to_string())?;

    // v0.3: every user query runs inside its own transaction so that
    // `SET LOCAL statement_timeout` and `READ ONLY` are scoped to the
    // query and cannot leak into pooled sessions. The validator already
    // rejected transaction-control statements, so BEGIN/COMMIT here are ours.
    let timeout_ms = options
        .timeout_ms
        .unwrap_or(DEFAULT_STATEMENT_TIMEOUT_MS)
        .clamp(1, MAX_STATEMENT_TIMEOUT_MS);
    client
        .batch_execute(if options.read_only {
            "BEGIN TRANSACTION READ ONLY"
        } else {
            "BEGIN"
        })
        .await
        .map_err(|e| e.to_string())?;
    client
        .batch_execute(&format!("SET LOCAL statement_timeout = '{timeout_ms}ms'"))
        .await
        .map_err(|e| e.to_string())?;

    // v0.3.2: interpolate params into the query text with safe escaping
    // (0.1.x-compatible semantics) instead of binary binding, which broke
    // on any JSON-integer vs int2/int4/float column and on unknown-type
    // inference. The statement gate already validated this query.
    let final_query = interpolate_query(&query, &params)?;

    // v0.3.4: escritas sem RETURNING usam execute() para obter o
    // numero real de linhas afetadas (query() retornaria 0 linhas).
    let (rows, affected) = if is_plain_write(&query) {
        match client.execute(&final_query, &[]).await {
            Ok(n) => (vec![], Some(n)),
            Err(e) => {
                let e = e.to_string();
                let _ = client.batch_execute("ROLLBACK").await;
                if e.contains("57014")
                    || e.contains("canceling statement due to statement timeout")
                {
                    return Err(
                        "STATEMENT_TIMEOUT: query exceeded its time budget".to_string(),
                    );
                }
                return Err(e);
            }
        }
    } else {
        match client.query(&final_query, &[]).await {
            Ok(rows) => (rows, None),
            Err(e) => {
                let e = e.to_string();
                let _ = client.batch_execute("ROLLBACK").await;
                // Postgres reports a cancelled statement as 57014.
                if e.contains("57014") || e.contains("canceling statement due to statement timeout") {
                    return Err("STATEMENT_TIMEOUT: query exceeded its time budget".to_string());
                }
                return Err(e);
            }
        }
    };
    client
        .batch_execute("COMMIT")
        .await
        .map_err(|e| e.to_string())?;

    let mut result = rows_to_result(rows);
    result.rows_affected = affected;
    Ok(result)
}

/// Execute several queries atomically on a single pooled connection.
/// Used by `/v1/transaction` (v0.3: real BEGIN/COMMIT instead of fan-out).
/// Permission checks must have been performed by the caller beforehand.
pub async fn execute_transaction_with_pool(
    pool: Pool,
    queries: Vec<(String, Vec<Value>, QueryOptions)>,
) -> Result<Vec<QueryResult>, String> {
    if queries.is_empty() {
        return Err("STATEMENT_NOT_ALLOWED: empty transaction".to_string());
    }
    if queries.len() > MAX_BATCH_ITEMS {
        return Err(format!(
            "STATEMENT_NOT_ALLOWED: transaction exceeds {MAX_BATCH_ITEMS} statements"
        ));
    }

    let mut client = pool.get().await.map_err(|e| e.to_string())?;
    let txn = client
        .transaction()
        .await
        .map_err(|e| e.to_string())?;
    let mut results = Vec::with_capacity(queries.len());
    for (query, params, options) in &queries {
        let timeout_ms = options
            .timeout_ms
            .unwrap_or(DEFAULT_STATEMENT_TIMEOUT_MS)
            .clamp(1, MAX_STATEMENT_TIMEOUT_MS);
        txn.batch_execute(&format!("SET LOCAL statement_timeout = '{timeout_ms}ms'"))
            .await
            .map_err(|e| e.to_string())?;

        // v0.3.2: safe interpolation (see execute_query_with_pool).
        let final_query = interpolate_query(query, params)?;
        // v0.3.4: plain writes usam execute() (linhas afetadas reais).
        if is_plain_write(query) {
            match txn.execute(&final_query, &[]).await {
                Ok(n) => {
                    let mut r = rows_to_result(vec![]);
                    r.rows_affected = Some(n);
                    results.push(r);
                }
                Err(e) => {
                    let msg = e.to_string();
                    txn.rollback().await.map_err(|rb| rb.to_string())?;
                    if msg.contains("57014")
                        || msg.contains("canceling statement due to statement timeout")
                    {
                        return Err(
                            "STATEMENT_TIMEOUT: query exceeded its time budget".to_string(),
                        );
                    }
                    return Err(msg);
                }
            }
            continue;
        }
        let rows = txn.query(&final_query, &[]).await;
        match rows {
            Ok(rows) => results.push(rows_to_result(rows)),
            Err(e) => {
                let msg = e.to_string();
                txn.rollback().await.map_err(|rb| rb.to_string())?;
                if msg.contains("57014")
                    || msg.contains("canceling statement due to statement timeout")
                {
                    return Err("STATEMENT_TIMEOUT: query exceeded its time budget".to_string());
                }
                return Err(msg);
            }
        }
    }
    txn.commit().await.map_err(|e| e.to_string())?;
    Ok(results)
}

/// Maximum number of statements accepted by `/v1/batch` and `/v1/transaction`.
pub const MAX_BATCH_ITEMS: usize = 50;

fn rows_to_result(rows: Vec<Row>) -> QueryResult {
    let fields = if !rows.is_empty() {
        rows[0].columns().iter().map(|col| {
            FieldMetadata {
                name: col.name().to_string(),
                data_type: col.type_().name().to_string(),
                nullable: true,
                max_length: None,
            }
        }).collect()
    } else {
        vec![]
    };

    let json_rows = rows_to_json(&rows);

    QueryResult {
        rows: json_rows,
        fields,
        query_plan: None,
        rows_affected: None,
    }
}

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

    #[test]
    fn test_is_plain_write() {
        // Escritas sem RETURNING => execute() (afetadas reais).
        assert!(is_plain_write("UPDATE t SET a = $1 WHERE id = $2"));
        assert!(is_plain_write("INSERT INTO t (a) VALUES ($1)"));
        assert!(is_plain_write("DELETE FROM t WHERE id = $1"));
        assert!(is_plain_write("  update t set a=null where id=$1  "));
        // Com RETURNING => query() (linhas retornadas).
        assert!(!is_plain_write(
            "UPDATE t SET a = $1 WHERE id = $2 RETURNING id"
        ));
        assert!(!is_plain_write("DELETE FROM t WHERE id = $1 RETURNING *"));
        // Leituras nunca sao plain writes.
        assert!(!is_plain_write("SELECT * FROM t WHERE a = $1"));
        assert!(!is_plain_write(
            "WITH x AS (SELECT 1) SELECT * FROM x"
        ));
        // RETURNING dentro de literal nao conta.
        assert!(is_plain_write(
            "UPDATE t SET a = 'RETURNING' WHERE id = $1"
        ));
        assert!(!is_plain_write("SELECT 'RETURNING'"));
    }

    #[test]
    fn test_validate_query_permissions_select() {
        let permissions = vec![Permission::Select];
        let role = AccountRole::Developer;
        
        let result = validate_query_permissions("SELECT * FROM users", &permissions, &role);
        assert!(result.is_ok());
    }

    #[test]
    fn test_validate_query_permissions_insert_denied() {
        let permissions = vec![Permission::Select];
        let role = AccountRole::Developer;

        let result = validate_query_permissions("INSERT INTO users VALUES (1)", &permissions, &role);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("MISSING_PERMISSION"));
    }

    #[test]
    fn test_validate_query_permissions_owner_bypass() {
        let permissions = vec![];
        let role = AccountRole::Superuser;
        
        // Owner pode fazer qualquer coisa
        let result = validate_query_permissions("DROP TABLE users", &permissions, &role);
        assert!(result.is_ok());
    }

    #[test]
    fn test_validate_query_permissions_all_permission() {
        let permissions = vec![Permission::All];
        let role = AccountRole::Developer;
        
        // Permissão All permite tudo
        let result = validate_query_permissions("DELETE FROM users", &permissions, &role);
        assert!(result.is_ok());
    }

    #[test]
    fn test_validate_query_permissions_update() {
        let permissions = vec![Permission::Update, Permission::Select];
        let role = AccountRole::Developer;
        
        let result = validate_query_permissions("UPDATE users SET name = 'test'", &permissions, &role);
        assert!(result.is_ok());
    }

    #[test]
    fn test_validate_query_permissions_delete() {
        let permissions = vec![Permission::Delete];
        let role = AccountRole::Developer;
        
        let result = validate_query_permissions("DELETE FROM users WHERE id = 1", &permissions, &role);
        assert!(result.is_ok());
    }

    #[test]
    fn test_validate_query_permissions_create() {
        let permissions = vec![Permission::Create];
        let role = AccountRole::Developer;
        
        let result = validate_query_permissions("CREATE TABLE test (id INT)", &permissions, &role);
        assert!(result.is_ok());
    }

    #[test]
    fn test_validate_query_permissions_drop() {
        let permissions = vec![Permission::Drop];
        let role = AccountRole::Developer;
        
        let result = validate_query_permissions("DROP TABLE test", &permissions, &role);
        assert!(result.is_ok());
    }

    #[test]
    fn test_validate_query_permissions_truncate() {
        let permissions = vec![Permission::Truncate];
        let role = AccountRole::Developer;
        
        let result = validate_query_permissions("TRUNCATE TABLE test", &permissions, &role);
        assert!(result.is_ok());
    }

    #[test]
    fn test_validate_query_permissions_unsupported() {
        let permissions = vec![Permission::Select];
        let role = AccountRole::Developer;
        
        let result = validate_query_permissions("ALTER TABLE users ADD COLUMN x INT", &permissions, &role);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("STATEMENT_NOT_ALLOWED"));
    }

    #[test]
    fn test_validate_query_permissions_case_insensitive() {
        let permissions = vec![Permission::Select];
        let role = AccountRole::Developer;
        
        let result = validate_query_permissions("select * from users", &permissions, &role);
        assert!(result.is_ok());
    }

    #[test]
    fn test_validate_query_permissions_with_cte() {
        let permissions = vec![Permission::Select];
        let role = AccountRole::Developer;
        
        let result = validate_query_permissions("WITH cte AS (SELECT * FROM users) SELECT * FROM cte", &permissions, &role);
        assert!(result.is_ok());
    }

    #[test]
    fn test_interpolate_basic_types() {
        assert_eq!(
            interpolate_query("SELECT $1, $2, $3, $4", &[json!(42), json!("o'clock"), json!(true), Value::Null]).unwrap(),
            "SELECT 42, 'o''clock', TRUE, NULL"
        );
    }

    #[test]
    fn test_interpolate_backslash_uses_escape_string() {
        assert_eq!(
            interpolate_query("SELECT $1", &[json!("a\\b'c")]).unwrap(),
            "SELECT E'a\\\\b\\'c'"
        );
    }

    #[test]
    fn test_interpolate_ignores_dollar_in_literals() {
        // '$1' como texto e 'R$ 5' nao sao placeholders.
        assert_eq!(
            interpolate_query("SELECT '$1', 'R$ 5', $1", &[json!(7)]).unwrap(),
            "SELECT '$1', 'R$ 5', 7"
        );
    }

    #[test]
    fn test_interpolate_injection_attempt_stays_literal() {
        let evil = "x'; DROP TABLE users; --";
        let out = interpolate_query("SELECT * FROM t WHERE name = $1", &[json!(evil)]).unwrap();
        assert_eq!(out, "SELECT * FROM t WHERE name = 'x''; DROP TABLE users; --'");
    }

    #[test]
    fn test_interpolate_extra_params_ignored() {
        // Compat 0.1.x: params sem placeholder sao ignorados.
        assert_eq!(
            interpolate_query("SELECT 1", &[json!(1), json!(2)]).unwrap(),
            "SELECT 1"
        );
    }

    #[test]
    fn test_interpolate_missing_param_errors() {
        let r = interpolate_query("SELECT $1, $3", &[json!(1), json!(2)]);
        assert!(r.is_err());
        assert!(r.unwrap_err().contains("missing value for parameter $3"));
    }

    #[test]
    fn test_interpolate_repeated_and_multi_digit() {
        assert_eq!(
            interpolate_query("SELECT $1, $1, $10", &[
                json!(1), json!(2), json!(3), json!(4), json!(5),
                json!(6), json!(7), json!(8), json!(9), json!("dez"),
            ]).unwrap(),
            "SELECT 1, 1, 'dez'"
        );
    }

    #[test]
    fn test_interpolate_json_values() {
        assert_eq!(
            interpolate_query("SELECT $1", &[json!({"a": 1})]).unwrap(),
            "SELECT '{\"a\":1}'"
        );
    }

    #[test]
    fn test_interpolate_big_u64() {
        // u64 acima de i64::MAX tem ramo proprio (nao vira float).
        let big = serde_json::Number::from(u64::MAX);
        assert_eq!(
            interpolate_query("SELECT $1", &[Value::Number(big)]).unwrap(),
            format!("SELECT {}", u64::MAX)
        );
    }

    #[test]
    fn test_numeric_binary_to_string_zero() {
        // Teste básico com zero
        let raw = vec![0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00];
        let result = numeric_binary_to_string(&raw);
        assert_eq!(result, Some("0".to_string()));
    }

    #[test]
    fn test_numeric_binary_to_string_nan() {
        // Teste com NaN
        let raw = vec![0x00, 0x00, 0x00, 0x00, 0xC0, 0x00, 0x00, 0x00];
        let result = numeric_binary_to_string(&raw);
        assert_eq!(result, Some("NaN".to_string()));
    }

    #[test]
    fn test_numeric_binary_to_string_invalid_short() {
        // Dados muito curtos
        let raw = vec![0x00, 0x00, 0x00];
        let result = numeric_binary_to_string(&raw);
        assert_eq!(result, None);
    }

    // --- v0.3 adversarial matrix: smuggling / bypass attempts ---

    fn perms(p: &[Permission]) -> Vec<Permission> {
        p.to_vec()
    }

    #[test]
    fn test_rejects_stacked_statements() {
        let role = AccountRole::Developer;
        let all = perms(&[Permission::All]);
        for q in [
            "SELECT 1; DROP TABLE users",
            "SELECT 1 ; SELECT 2",
            "WITH a AS (SELECT 1) SELECT * FROM a; DELETE FROM users",
        ] {
            let r = validate_query_permissions(q, &all, &role);
            assert!(r.is_err(), "{q} should be rejected");
            assert!(r.unwrap_err().contains("MULTI_STATEMENT"), "{q}");
        }
    }

    #[test]
    fn test_semicolon_inside_literal_is_single_statement() {
        let role = AccountRole::Developer;
        let sel = perms(&[Permission::Select]);
        assert!(validate_query_permissions("SELECT ';' AS s", &sel, &role).is_ok());
        assert!(validate_query_permissions("SELECT $1", &sel, &role).is_ok());
    }

    #[test]
    fn test_comment_smuggling_blocked() {
        let role = AccountRole::Developer;
        let sel = perms(&[Permission::Select]);
        // Prefixo SELECT falso escondido em comentário.
        assert!(validate_query_permissions("/* SELECT */ DELETE FROM users", &sel, &role).is_err());
        assert!(validate_query_permissions("-- SELECT\nDELETE FROM users", &sel, &role).is_err());
        // Comentário após statement válido não cria segundo statement.
        assert!(validate_query_permissions("SELECT 1 /* trailing */", &sel, &role).is_ok());
    }

    #[test]
    fn test_with_writes_require_write_permission() {
        let role = AccountRole::Developer;
        let sel = perms(&[Permission::Select]);
        let del = perms(&[Permission::Select, Permission::Delete]);
        let q = "WITH moved AS (DELETE FROM a RETURNING *) SELECT * FROM moved";
        assert!(validate_query_permissions(q, &sel, &role).is_err());
        assert!(validate_query_permissions(q, &del, &role).is_ok());
        // WITH puramente de leitura continua liberado com SELECT.
        assert!(validate_query_permissions(
            "WITH cte AS (SELECT * FROM users) SELECT * FROM cte",
            &sel, &role
        ).is_ok());
    }

    #[test]
    fn test_row_lock_query_is_allowed_for_queue_claims() {
        let role = AccountRole::Developer;
        let permissions = perms(&[Permission::Select]);
        let query = "SELECT id FROM queue WHERE status = 'pending' ORDER BY id FOR UPDATE SKIP LOCKED LIMIT 1";
        assert!(validate_query_permissions(query, &permissions, &role).is_ok());
    }

    #[test]
    fn test_denied_statements_and_functions() {
        let role = AccountRole::Developer;
        let all = perms(&[Permission::All]);
        for q in [
            "COPY users TO STDOUT",
            "GRANT SELECT ON t TO u",
            "VACUUM users",
            "CALL my_proc()",
            "DO $$ BEGIN RAISE NOTICE 'hi'; END $$",
            "EXPLAIN ANALYZE SELECT * FROM users",
            "BEGIN",
            "SET ROLE postgres",
            "SELECT pg_read_file('pg_hba.conf')",
            "SELECT dblink('c', 'SELECT 1')",
            "SELECT lo_import('/etc/passwd')",
        ] {
            let r = validate_query_permissions(q, &all, &role);
            assert!(r.is_err(), "{q} should be denied even with ALL");
            assert!(r.unwrap_err().contains("STATEMENT_NOT_ALLOWED"), "{q}");
        }
    }

    #[test]
    fn test_dangerous_words_as_values_are_allowed() {
        // Falso-positivo não pode quebrar CRUD legítimo.
        let role = AccountRole::Developer;
        let w = perms(&[Permission::Select, Permission::Insert, Permission::Update]);
        assert!(validate_query_permissions(
            "SELECT * FROM posts WHERE comment = 'DROP TABLE x'", &w, &role).is_ok());
        assert!(validate_query_permissions("UPDATE users SET name = 'x'", &w, &role).is_ok());
        assert!(validate_query_permissions("SELECT set_config('a','b',false)", &w, &role).is_ok());
    }

    #[test]
    fn test_explain_without_analyze_allowed() {
        let role = AccountRole::Developer;
        let sel = perms(&[Permission::Select]);
        assert!(validate_query_permissions("EXPLAIN SELECT * FROM users", &sel, &role).is_ok());
    }

    #[test]
    fn test_superuser_bypass_kept() {
        let role = AccountRole::Superuser;
        assert!(validate_query_permissions("DROP TABLE users", &[], &role).is_ok());
    }
}