qail-pg 0.27.10

//! DDL (Data Definition Language) encoders.
//!
//! CREATE TABLE, CREATE INDEX, DROP, ALTER statements.

use bytes::BytesMut;
use qail_core::ast::{
    Action, ColumnGeneration, Constraint, Expr, Qail, TableConstraint, TriggerEvent, TriggerTiming,
};
use qail_core::migrate::policy::{PolicyPermissiveness, PolicyTarget};

/// Map QAIL types to PostgreSQL types.
#[inline]
pub fn map_type(t: &str) -> &str {
    match t {
        "str" | "text" | "string" | "TEXT" => "TEXT",
        "smallint" | "SMALLINT" | "int2" | "INT2" => "SMALLINT",
        "int" | "i32" | "INT" | "INTEGER" => "INT",
        "bigint" | "i64" | "BIGINT" => "BIGINT",
        "uuid" | "UUID" => "UUID",
        "bool" | "boolean" | "BOOLEAN" => "BOOLEAN",
        "dec" | "decimal" | "DECIMAL" => "DECIMAL",
        "float" | "f64" | "DOUBLE PRECISION" => "DOUBLE PRECISION",
        "serial" | "SERIAL" => "SERIAL",
        "bigserial" | "BIGSERIAL" => "BIGSERIAL",
        "timestamp" | "time" | "TIMESTAMP" => "TIMESTAMP",
        "timestamptz" | "TIMESTAMPTZ" => "TIMESTAMPTZ",
        "date" | "DATE" => "DATE",
        "json" | "jsonb" | "JSON" | "JSONB" => "JSONB",
        "varchar" | "VARCHAR" => "VARCHAR(255)",
        _ => t,
    }
}

/// Encode CREATE TABLE statement.
pub fn encode_make(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"CREATE TABLE ");
    buf.extend_from_slice(cmd.table.as_bytes());
    buf.extend_from_slice(b" (");

    let composite_pk_columns: Vec<&str> = cmd
        .columns
        .iter()
        .filter_map(|col| match col {
            Expr::Def {
                name, constraints, ..
            } if constraints.contains(&Constraint::PrimaryKey) => Some(name.as_str()),
            _ => None,
        })
        .collect();
    let use_composite_pk = composite_pk_columns.len() > 1;

    let mut first = true;
    for col in &cmd.columns {
        if let Expr::Def {
            name,
            data_type,
            constraints,
        } = col
        {
            if !first {
                buf.extend_from_slice(b", ");
            }
            first = false;

            buf.extend_from_slice(name.as_bytes());
            buf.extend_from_slice(b" ");
            buf.extend_from_slice(map_type(data_type).as_bytes());

            // Default to NOT NULL unless Nullable
            if !constraints.contains(&Constraint::Nullable) {
                buf.extend_from_slice(b" NOT NULL");
            }

            // DEFAULT
            for constraint in constraints {
                if let Constraint::Default(val) = constraint {
                    buf.extend_from_slice(b" DEFAULT ");
                    let sql_default = match val.as_str() {
                        "uuid()" => "gen_random_uuid()",
                        "now()" => "NOW()",
                        other => other,
                    };
                    buf.extend_from_slice(sql_default.as_bytes());
                }
                if let Constraint::Generated(generation) = constraint {
                    match generation {
                        ColumnGeneration::Stored(expr) if expr == "identity" => {
                            buf.extend_from_slice(b" GENERATED ALWAYS AS IDENTITY");
                        }
                        ColumnGeneration::Stored(expr) if expr == "identity_by_default" => {
                            buf.extend_from_slice(b" GENERATED BY DEFAULT AS IDENTITY");
                        }
                        ColumnGeneration::Stored(expr) => {
                            buf.extend_from_slice(b" GENERATED ALWAYS AS (");
                            buf.extend_from_slice(expr.as_bytes());
                            buf.extend_from_slice(b") STORED");
                        }
                        ColumnGeneration::Virtual(expr) => {
                            buf.extend_from_slice(b" GENERATED ALWAYS AS (");
                            buf.extend_from_slice(expr.as_bytes());
                            buf.extend_from_slice(b")");
                        }
                    }
                }
            }

            // PRIMARY KEY
            if constraints.contains(&Constraint::PrimaryKey) && !use_composite_pk {
                buf.extend_from_slice(b" PRIMARY KEY");
            }

            // UNIQUE
            if constraints.contains(&Constraint::Unique) {
                buf.extend_from_slice(b" UNIQUE");
            }

            // REFERENCES (foreign key)
            for constraint in constraints {
                if let Constraint::References(target) = constraint {
                    buf.extend_from_slice(b" REFERENCES ");
                    buf.extend_from_slice(target.as_bytes());
                }
            }

            // CHECK constraint
            for constraint in constraints {
                if let Constraint::Check(vals) = constraint {
                    if vals.len() == 1
                        && vals[0]
                            .trim_start()
                            .to_ascii_uppercase()
                            .starts_with("CONSTRAINT ")
                    {
                        buf.extend_from_slice(b" ");
                        buf.extend_from_slice(vals[0].as_bytes());
                        continue;
                    }

                    let looks_like_expr = vals.len() == 1
                        || vals.iter().any(|v| {
                            v.chars().any(|c| {
                                c.is_whitespace() || matches!(c, '<' | '>' | '=' | '!' | '(' | ')')
                            })
                        });

                    if looks_like_expr {
                        buf.extend_from_slice(b" CHECK (");
                        if vals.len() == 1 {
                            buf.extend_from_slice(vals[0].as_bytes());
                        } else {
                            for (i, v) in vals.iter().enumerate() {
                                if i > 0 {
                                    buf.extend_from_slice(b" ");
                                }
                                buf.extend_from_slice(v.as_bytes());
                            }
                        }
                        buf.extend_from_slice(b")");
                    } else {
                        buf.extend_from_slice(b" CHECK (");
                        buf.extend_from_slice(name.as_bytes());
                        buf.extend_from_slice(b" IN (");
                        for (i, v) in vals.iter().enumerate() {
                            if i > 0 {
                                buf.extend_from_slice(b", ");
                            }
                            buf.extend_from_slice(b"'");
                            buf.extend_from_slice(v.as_bytes());
                            buf.extend_from_slice(b"'");
                        }
                        buf.extend_from_slice(b"))");
                    }
                }
            }
        }
    }

    if use_composite_pk {
        if !first {
            buf.extend_from_slice(b", ");
        }
        first = false;
        buf.extend_from_slice(b"PRIMARY KEY (");
        for (i, col) in composite_pk_columns.iter().enumerate() {
            if i > 0 {
                buf.extend_from_slice(b", ");
            }
            buf.extend_from_slice(col.as_bytes());
        }
        buf.extend_from_slice(b")");
    }

    // Table constraints
    for tc in &cmd.table_constraints {
        if !first {
            buf.extend_from_slice(b", ");
        }
        first = false;
        match tc {
            TableConstraint::Unique(cols) => {
                buf.extend_from_slice(b"UNIQUE (");
                for (i, col) in cols.iter().enumerate() {
                    if i > 0 {
                        buf.extend_from_slice(b", ");
                    }
                    buf.extend_from_slice(col.as_bytes());
                }
                buf.extend_from_slice(b")");
            }
            TableConstraint::PrimaryKey(cols) => {
                buf.extend_from_slice(b"PRIMARY KEY (");
                for (i, col) in cols.iter().enumerate() {
                    if i > 0 {
                        buf.extend_from_slice(b", ");
                    }
                    buf.extend_from_slice(col.as_bytes());
                }
                buf.extend_from_slice(b")");
            }
        }
    }

    buf.extend_from_slice(b")");
}

/// Encode CREATE INDEX statement.
pub fn encode_index(cmd: &Qail, buf: &mut BytesMut) {
    if let Some(idx) = &cmd.index_def {
        if idx.unique {
            buf.extend_from_slice(b"CREATE UNIQUE INDEX ");
        } else {
            buf.extend_from_slice(b"CREATE INDEX ");
        }
        buf.extend_from_slice(idx.name.as_bytes());
        buf.extend_from_slice(b" ON ");
        buf.extend_from_slice(idx.table.as_bytes());
        if let Some(method) = &idx.index_type
            && !method.trim().is_empty()
        {
            buf.extend_from_slice(b" USING ");
            buf.extend_from_slice(method.trim().as_bytes());
        }
        buf.extend_from_slice(b" (");
        for (i, col) in idx.columns.iter().enumerate() {
            if i > 0 {
                buf.extend_from_slice(b", ");
            }
            buf.extend_from_slice(col.as_bytes());
        }
        buf.extend_from_slice(b")");
        if let Some(where_clause) = &idx.where_clause {
            buf.extend_from_slice(b" WHERE ");
            buf.extend_from_slice(where_clause.as_bytes());
        }
    }
}

/// Encode DROP TABLE statement.
pub fn encode_drop_table(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"DROP TABLE IF EXISTS ");
    buf.extend_from_slice(cmd.table.as_bytes());
}

/// Encode DROP INDEX statement.
pub fn encode_drop_index(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"DROP INDEX IF EXISTS ");
    buf.extend_from_slice(cmd.table.as_bytes());
}

/// Encode ALTER TABLE ADD COLUMN statement.
pub fn encode_alter_add_column(cmd: &Qail, buf: &mut BytesMut) {
    for col in &cmd.columns {
        if let Expr::Def {
            name,
            data_type,
            constraints,
        } = col
        {
            buf.extend_from_slice(b"ALTER TABLE ");
            buf.extend_from_slice(cmd.table.as_bytes());
            buf.extend_from_slice(b" ADD COLUMN ");
            buf.extend_from_slice(name.as_bytes());
            buf.extend_from_slice(b" ");
            buf.extend_from_slice(map_type(data_type).as_bytes());

            if !constraints.contains(&Constraint::Nullable) {
                buf.extend_from_slice(b" NOT NULL");
            }

            for constraint in constraints {
                if let Constraint::Default(val) = constraint {
                    buf.extend_from_slice(b" DEFAULT ");
                    let sql_default = match val.as_str() {
                        "uuid()" => "gen_random_uuid()",
                        "now()" => "NOW()",
                        other => other,
                    };
                    buf.extend_from_slice(sql_default.as_bytes());
                }
            }
        }
    }
}

/// Encode ALTER TABLE DROP COLUMN statement.
pub fn encode_alter_drop_column(cmd: &Qail, buf: &mut BytesMut) {
    for col in &cmd.columns {
        let col_name = match col {
            Expr::Named(n) => n.clone(),
            Expr::Def { name, .. } => name.clone(),
            _ => continue,
        };
        buf.extend_from_slice(b"ALTER TABLE ");
        buf.extend_from_slice(cmd.table.as_bytes());
        buf.extend_from_slice(b" DROP COLUMN ");
        buf.extend_from_slice(col_name.as_bytes());
    }
}

/// Encode ALTER TABLE ALTER COLUMN TYPE statement.
pub fn encode_alter_column_type(cmd: &Qail, buf: &mut BytesMut) {
    for col in &cmd.columns {
        if let Expr::Def {
            name, data_type, ..
        } = col
        {
            buf.extend_from_slice(b"ALTER TABLE ");
            buf.extend_from_slice(cmd.table.as_bytes());
            buf.extend_from_slice(b" ALTER COLUMN ");
            buf.extend_from_slice(name.as_bytes());
            buf.extend_from_slice(b" TYPE ");
            buf.extend_from_slice(map_type(data_type).as_bytes());
        }
    }
}

/// Encode ALTER TABLE RENAME COLUMN statement.
/// The `Mod` action stores renames as `Expr::Named("old_name -> new_name")`.
pub fn encode_rename_column(cmd: &Qail, buf: &mut BytesMut) {
    for col in &cmd.columns {
        if let Expr::Named(rename_str) = col {
            // Parse "old_name -> new_name" format
            if let Some((old, new)) = rename_str.split_once(" -> ") {
                buf.extend_from_slice(b"ALTER TABLE ");
                buf.extend_from_slice(cmd.table.as_bytes());
                buf.extend_from_slice(b" RENAME COLUMN ");
                buf.extend_from_slice(old.trim().as_bytes());
                buf.extend_from_slice(b" TO ");
                buf.extend_from_slice(new.trim().as_bytes());
            }
        }
    }
}

/// Encode CREATE VIEW statement.
/// CREATE VIEW name AS SELECT ...
pub fn encode_create_view(
    cmd: &Qail,
    buf: &mut BytesMut,
    params: &mut Vec<Option<Vec<u8>>>,
) -> Result<(), super::super::EncodeError> {
    buf.extend_from_slice(b"CREATE VIEW ");
    buf.extend_from_slice(cmd.table.as_bytes());
    buf.extend_from_slice(b" AS ");

    // The source_query contains the SELECT statement for the view
    if let Some(ref source) = cmd.source_query {
        super::dml::encode_select(source, buf, params)?;
    } else if let Some(query) = &cmd.payload {
        buf.extend_from_slice(query.as_bytes());
    } else {
        return Err(super::super::EncodeError::UnsupportedAction(
            Action::CreateView,
        ));
    }
    Ok(())
}

/// Encode DROP VIEW statement.
pub fn encode_drop_view(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"DROP VIEW IF EXISTS ");
    buf.extend_from_slice(cmd.table.as_bytes());
}

/// Encode CREATE MATERIALIZED VIEW statement.
pub fn encode_create_materialized_view(
    cmd: &Qail,
    buf: &mut BytesMut,
    params: &mut Vec<Option<Vec<u8>>>,
) -> Result<(), super::super::EncodeError> {
    buf.extend_from_slice(b"CREATE MATERIALIZED VIEW ");
    buf.extend_from_slice(cmd.table.as_bytes());
    buf.extend_from_slice(b" AS ");

    if let Some(ref source) = cmd.source_query {
        super::dml::encode_select(source, buf, params)?;
    } else if let Some(query) = &cmd.payload {
        buf.extend_from_slice(query.as_bytes());
    } else {
        return Err(super::super::EncodeError::UnsupportedAction(
            Action::CreateMaterializedView,
        ));
    }
    Ok(())
}

/// Encode REFRESH MATERIALIZED VIEW statement.
pub fn encode_refresh_materialized_view(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"REFRESH MATERIALIZED VIEW ");
    buf.extend_from_slice(cmd.table.as_bytes());
}

/// Encode DROP MATERIALIZED VIEW statement.
pub fn encode_drop_materialized_view(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"DROP MATERIALIZED VIEW IF EXISTS ");
    buf.extend_from_slice(cmd.table.as_bytes());
}

/// Encode ALTER TABLE ALTER COLUMN SET NOT NULL.
pub fn encode_alter_set_not_null(cmd: &Qail, buf: &mut BytesMut) {
    if let Some(Expr::Named(col)) = cmd.columns.first() {
        buf.extend_from_slice(b"ALTER TABLE ");
        buf.extend_from_slice(cmd.table.as_bytes());
        buf.extend_from_slice(b" ALTER COLUMN ");
        buf.extend_from_slice(col.as_bytes());
        buf.extend_from_slice(b" SET NOT NULL");
    }
}

/// Encode ALTER TABLE ALTER COLUMN DROP NOT NULL.
pub fn encode_alter_drop_not_null(cmd: &Qail, buf: &mut BytesMut) {
    if let Some(Expr::Named(col)) = cmd.columns.first() {
        buf.extend_from_slice(b"ALTER TABLE ");
        buf.extend_from_slice(cmd.table.as_bytes());
        buf.extend_from_slice(b" ALTER COLUMN ");
        buf.extend_from_slice(col.as_bytes());
        buf.extend_from_slice(b" DROP NOT NULL");
    }
}

/// Encode ALTER TABLE ALTER COLUMN SET DEFAULT.
pub fn encode_alter_set_default(cmd: &Qail, buf: &mut BytesMut) {
    if let Some(Expr::Named(col)) = cmd.columns.first() {
        buf.extend_from_slice(b"ALTER TABLE ");
        buf.extend_from_slice(cmd.table.as_bytes());
        buf.extend_from_slice(b" ALTER COLUMN ");
        buf.extend_from_slice(col.as_bytes());
        buf.extend_from_slice(b" SET DEFAULT ");
        let default_expr = cmd.payload.as_deref().unwrap_or("NULL");
        buf.extend_from_slice(default_expr.as_bytes());
    }
}

/// Encode ALTER TABLE ALTER COLUMN DROP DEFAULT.
pub fn encode_alter_drop_default(cmd: &Qail, buf: &mut BytesMut) {
    if let Some(Expr::Named(col)) = cmd.columns.first() {
        buf.extend_from_slice(b"ALTER TABLE ");
        buf.extend_from_slice(cmd.table.as_bytes());
        buf.extend_from_slice(b" ALTER COLUMN ");
        buf.extend_from_slice(col.as_bytes());
        buf.extend_from_slice(b" DROP DEFAULT");
    }
}

/// Encode ALTER TABLE ENABLE ROW LEVEL SECURITY.
pub fn encode_alter_enable_rls(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"ALTER TABLE ");
    buf.extend_from_slice(cmd.table.as_bytes());
    buf.extend_from_slice(b" ENABLE ROW LEVEL SECURITY");
}

/// Encode ALTER TABLE DISABLE ROW LEVEL SECURITY.
pub fn encode_alter_disable_rls(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"ALTER TABLE ");
    buf.extend_from_slice(cmd.table.as_bytes());
    buf.extend_from_slice(b" DISABLE ROW LEVEL SECURITY");
}

/// Encode ALTER TABLE FORCE ROW LEVEL SECURITY.
pub fn encode_alter_force_rls(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"ALTER TABLE ");
    buf.extend_from_slice(cmd.table.as_bytes());
    buf.extend_from_slice(b" FORCE ROW LEVEL SECURITY");
}

/// Encode ALTER TABLE NO FORCE ROW LEVEL SECURITY.
pub fn encode_alter_no_force_rls(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"ALTER TABLE ");
    buf.extend_from_slice(cmd.table.as_bytes());
    buf.extend_from_slice(b" NO FORCE ROW LEVEL SECURITY");
}

// ── Session & procedural commands ──────────────────────────────────

/// Encode SAVEPOINT name.
pub fn encode_savepoint(cmd: &Qail, buf: &mut BytesMut) {
    let name = cmd.savepoint_name.as_deref().unwrap_or("qail_sp");
    buf.extend_from_slice(b"SAVEPOINT ");
    buf.extend_from_slice(name.as_bytes());
}

/// Encode RELEASE SAVEPOINT name.
pub fn encode_release_savepoint(cmd: &Qail, buf: &mut BytesMut) {
    let name = cmd.savepoint_name.as_deref().unwrap_or("qail_sp");
    buf.extend_from_slice(b"RELEASE SAVEPOINT ");
    buf.extend_from_slice(name.as_bytes());
}

/// Encode ROLLBACK TO SAVEPOINT name.
pub fn encode_rollback_to_savepoint(cmd: &Qail, buf: &mut BytesMut) {
    let name = cmd.savepoint_name.as_deref().unwrap_or("qail_sp");
    buf.extend_from_slice(b"ROLLBACK TO SAVEPOINT ");
    buf.extend_from_slice(name.as_bytes());
}

/// Encode CALL procedure_name.
pub fn encode_call(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"CALL ");
    buf.extend_from_slice(cmd.table.as_bytes());
}

/// Encode DO $$ body $$ LANGUAGE lang.
pub fn encode_do(cmd: &Qail, buf: &mut BytesMut) {
    let body = cmd.payload.as_deref().unwrap_or("");
    let lang = if cmd.table.is_empty() {
        "plpgsql"
    } else {
        &cmd.table
    };
    buf.extend_from_slice(b"DO $$ ");
    buf.extend_from_slice(body.as_bytes());
    buf.extend_from_slice(b" $$ LANGUAGE ");
    buf.extend_from_slice(lang.as_bytes());
}

/// Encode SET key = 'value'.
///
/// The value is escaped: `'` → `''` to prevent SQL injection.
pub fn encode_session_set(cmd: &Qail, buf: &mut BytesMut) {
    let value = cmd.payload.as_deref().unwrap_or("");
    let escaped = value.replace('\'', "''");
    buf.extend_from_slice(b"SET ");
    buf.extend_from_slice(cmd.table.as_bytes());
    buf.extend_from_slice(b" = '");
    buf.extend_from_slice(escaped.as_bytes());
    buf.extend_from_slice(b"'");
}

/// Encode SHOW key.
pub fn encode_session_show(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"SHOW ");
    buf.extend_from_slice(cmd.table.as_bytes());
}

/// Encode RESET key.
pub fn encode_session_reset(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"RESET ");
    buf.extend_from_slice(cmd.table.as_bytes());
}

#[inline]
fn encode_identifier_maybe_quoted(buf: &mut BytesMut, ident: &str) {
    let needs_quotes = ident.is_empty()
        || ident
            .chars()
            .next()
            .map(|c| c.is_ascii_digit())
            .unwrap_or(false)
        || ident
            .chars()
            .any(|c| !c.is_ascii_alphanumeric() && c != '_');

    if needs_quotes {
        buf.extend_from_slice(b"\"");
        buf.extend_from_slice(ident.replace('"', "\"\"").as_bytes());
        buf.extend_from_slice(b"\"");
    } else {
        buf.extend_from_slice(ident.as_bytes());
    }
}

/// Encode CREATE DATABASE name.
pub fn encode_create_database(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"CREATE DATABASE ");
    encode_identifier_maybe_quoted(buf, &cmd.table);
}

/// Encode DROP DATABASE IF EXISTS name.
pub fn encode_drop_database(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"DROP DATABASE IF EXISTS ");
    encode_identifier_maybe_quoted(buf, &cmd.table);
}

/// Encode GRANT privileges ON object TO role.
pub fn encode_grant(cmd: &Qail, buf: &mut BytesMut) -> Result<(), super::super::EncodeError> {
    let Some(role) = cmd.payload.as_deref() else {
        return Err(super::super::EncodeError::UnsupportedAction(Action::Grant));
    };

    let mut first = true;
    let mut privs = String::new();
    for col in &cmd.columns {
        if let Expr::Named(p) = col {
            if !first {
                privs.push_str(", ");
            }
            first = false;
            privs.push_str(p);
        }
    }
    if privs.is_empty() || cmd.table.trim().is_empty() || role.trim().is_empty() {
        return Err(super::super::EncodeError::UnsupportedAction(Action::Grant));
    }

    buf.extend_from_slice(b"GRANT ");
    buf.extend_from_slice(privs.as_bytes());
    buf.extend_from_slice(b" ON ");
    buf.extend_from_slice(cmd.table.as_bytes());
    buf.extend_from_slice(b" TO ");
    buf.extend_from_slice(role.as_bytes());
    Ok(())
}

/// Encode REVOKE privileges ON object FROM role.
pub fn encode_revoke(cmd: &Qail, buf: &mut BytesMut) -> Result<(), super::super::EncodeError> {
    let Some(role) = cmd.payload.as_deref() else {
        return Err(super::super::EncodeError::UnsupportedAction(Action::Revoke));
    };

    let mut first = true;
    let mut privs = String::new();
    for col in &cmd.columns {
        if let Expr::Named(p) = col {
            if !first {
                privs.push_str(", ");
            }
            first = false;
            privs.push_str(p);
        }
    }
    if privs.is_empty() || cmd.table.trim().is_empty() || role.trim().is_empty() {
        return Err(super::super::EncodeError::UnsupportedAction(Action::Revoke));
    }

    buf.extend_from_slice(b"REVOKE ");
    buf.extend_from_slice(privs.as_bytes());
    buf.extend_from_slice(b" ON ");
    buf.extend_from_slice(cmd.table.as_bytes());
    buf.extend_from_slice(b" FROM ");
    buf.extend_from_slice(role.as_bytes());
    Ok(())
}

/// Encode CREATE POLICY statement.
pub fn encode_create_policy(
    cmd: &Qail,
    buf: &mut BytesMut,
) -> Result<(), super::super::EncodeError> {
    let Some(policy) = &cmd.policy_def else {
        return Err(super::super::EncodeError::UnsupportedAction(
            Action::CreatePolicy,
        ));
    };

    buf.extend_from_slice(b"CREATE POLICY ");
    buf.extend_from_slice(policy.name.as_bytes());
    buf.extend_from_slice(b" ON ");
    buf.extend_from_slice(policy.table.as_bytes());

    if policy.permissiveness == PolicyPermissiveness::Restrictive {
        buf.extend_from_slice(b" AS RESTRICTIVE");
    }

    let target = match policy.target {
        PolicyTarget::All => "ALL",
        PolicyTarget::Select => "SELECT",
        PolicyTarget::Insert => "INSERT",
        PolicyTarget::Update => "UPDATE",
        PolicyTarget::Delete => "DELETE",
    };
    buf.extend_from_slice(b" FOR ");
    buf.extend_from_slice(target.as_bytes());

    if let Some(role) = &policy.role {
        buf.extend_from_slice(b" TO ");
        buf.extend_from_slice(role.as_bytes());
    }

    if let Some(expr) = &policy.using {
        buf.extend_from_slice(b" USING (");
        buf.extend_from_slice(expr.to_string().as_bytes());
        buf.extend_from_slice(b")");
    }
    if let Some(expr) = &policy.with_check {
        buf.extend_from_slice(b" WITH CHECK (");
        buf.extend_from_slice(expr.to_string().as_bytes());
        buf.extend_from_slice(b")");
    }

    Ok(())
}

/// Encode DROP POLICY statement.
///
/// Expects table in `cmd.table` and policy name in `cmd.payload` (or `cmd.policy_def`).
pub fn encode_drop_policy(cmd: &Qail, buf: &mut BytesMut) -> Result<(), super::super::EncodeError> {
    let (policy_name, table_name) = if let Some(policy) = &cmd.policy_def {
        (policy.name.as_str(), policy.table.as_str())
    } else if let Some(name) = cmd.payload.as_deref() {
        (name, cmd.table.as_str())
    } else {
        return Err(super::super::EncodeError::UnsupportedAction(
            Action::DropPolicy,
        ));
    };

    if policy_name.trim().is_empty() || table_name.trim().is_empty() {
        return Err(super::super::EncodeError::UnsupportedAction(
            Action::DropPolicy,
        ));
    }

    buf.extend_from_slice(b"DROP POLICY IF EXISTS ");
    buf.extend_from_slice(policy_name.as_bytes());
    buf.extend_from_slice(b" ON ");
    buf.extend_from_slice(table_name.as_bytes());
    Ok(())
}

/// Encode CREATE FUNCTION statement.
pub fn encode_create_function(
    cmd: &Qail,
    buf: &mut BytesMut,
) -> Result<(), super::super::EncodeError> {
    let Some(func) = &cmd.function_def else {
        return Err(super::super::EncodeError::UnsupportedAction(
            Action::CreateFunction,
        ));
    };

    let lang = func.language.as_deref().unwrap_or("plpgsql");
    let args = func.args.join(", ");
    buf.extend_from_slice(b"CREATE OR REPLACE FUNCTION ");
    buf.extend_from_slice(func.name.as_bytes());
    buf.extend_from_slice(b"(");
    buf.extend_from_slice(args.as_bytes());
    buf.extend_from_slice(b") RETURNS ");
    buf.extend_from_slice(func.returns.as_bytes());
    buf.extend_from_slice(b" LANGUAGE ");
    buf.extend_from_slice(lang.as_bytes());
    if let Some(volatility) = &func.volatility
        && !volatility.trim().is_empty()
    {
        buf.extend_from_slice(b" ");
        buf.extend_from_slice(volatility.trim().to_ascii_uppercase().as_bytes());
    }
    buf.extend_from_slice(b" AS $$ ");
    buf.extend_from_slice(func.body.as_bytes());
    buf.extend_from_slice(b" $$");
    Ok(())
}

/// Encode DROP FUNCTION statement.
pub fn encode_drop_function(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"DROP FUNCTION IF EXISTS ");
    if let Some(signature) = &cmd.payload {
        buf.extend_from_slice(signature.as_bytes());
    } else {
        buf.extend_from_slice(cmd.table.as_bytes());
        buf.extend_from_slice(b"()");
    }
}

/// Encode CREATE TRIGGER statement.
pub fn encode_create_trigger(
    cmd: &Qail,
    buf: &mut BytesMut,
) -> Result<(), super::super::EncodeError> {
    let Some(trig) = &cmd.trigger_def else {
        return Err(super::super::EncodeError::UnsupportedAction(
            Action::CreateTrigger,
        ));
    };

    let timing = match trig.timing {
        TriggerTiming::Before => "BEFORE",
        TriggerTiming::After => "AFTER",
        TriggerTiming::InsteadOf => "INSTEAD OF",
    };

    let mut first = true;
    let mut events = String::new();
    for evt in &trig.events {
        if !first {
            events.push_str(" OR ");
        }
        first = false;
        let evt_str = match evt {
            TriggerEvent::Insert => "INSERT",
            TriggerEvent::Update => {
                if !trig.update_columns.is_empty() {
                    events.push_str("UPDATE OF ");
                    events.push_str(&trig.update_columns.join(", "));
                    continue;
                }
                "UPDATE"
            }
            TriggerEvent::Delete => "DELETE",
            TriggerEvent::Truncate => "TRUNCATE",
        };
        events.push_str(evt_str);
    }

    if events.is_empty() {
        return Err(super::super::EncodeError::UnsupportedAction(
            Action::CreateTrigger,
        ));
    }

    let for_each = if trig.for_each_row {
        "FOR EACH ROW"
    } else {
        "FOR EACH STATEMENT"
    };

    buf.extend_from_slice(b"CREATE TRIGGER ");
    buf.extend_from_slice(trig.name.as_bytes());
    buf.extend_from_slice(b" ");
    buf.extend_from_slice(timing.as_bytes());
    buf.extend_from_slice(b" ");
    buf.extend_from_slice(events.as_bytes());
    buf.extend_from_slice(b" ON ");
    buf.extend_from_slice(trig.table.as_bytes());
    buf.extend_from_slice(b" ");
    buf.extend_from_slice(for_each.as_bytes());
    buf.extend_from_slice(b" EXECUTE FUNCTION ");
    buf.extend_from_slice(trig.execute_function.as_bytes());
    buf.extend_from_slice(b"()");
    Ok(())
}

/// Encode DROP TRIGGER statement.
///
/// Expects `cmd.table` in `table.trigger` form.
pub fn encode_drop_trigger(
    cmd: &Qail,
    buf: &mut BytesMut,
) -> Result<(), super::super::EncodeError> {
    let Some((table, trigger)) = cmd.table.rsplit_once('.') else {
        return Err(super::super::EncodeError::UnsupportedAction(
            Action::DropTrigger,
        ));
    };
    buf.extend_from_slice(b"DROP TRIGGER IF EXISTS ");
    buf.extend_from_slice(trigger.as_bytes());
    buf.extend_from_slice(b" ON ");
    buf.extend_from_slice(table.as_bytes());
    Ok(())
}

/// Encode CREATE EXTENSION statement.
pub fn encode_create_extension(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"CREATE EXTENSION IF NOT EXISTS \"");
    buf.extend_from_slice(cmd.table.replace('"', "\"\"").as_bytes());
    buf.extend_from_slice(b"\"");

    for col in &cmd.columns {
        if let Expr::Named(opt) = col {
            buf.extend_from_slice(b" ");
            buf.extend_from_slice(opt.as_bytes());
        }
    }
}

/// Encode DROP EXTENSION statement.
pub fn encode_drop_extension(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"DROP EXTENSION IF EXISTS \"");
    buf.extend_from_slice(cmd.table.replace('"', "\"\"").as_bytes());
    buf.extend_from_slice(b"\"");
}

/// Encode COMMENT ON TABLE/COLUMN statement.
pub fn encode_comment_on(cmd: &Qail, buf: &mut BytesMut) {
    let comment_text = cmd
        .columns
        .first()
        .and_then(|c| match c {
            Expr::Named(s) => Some(s.as_str()),
            _ => None,
        })
        .unwrap_or("");
    let escaped = comment_text.replace('\'', "''");

    let trimmed = cmd.table.trim_start();
    let upper = trimmed.to_ascii_uppercase();
    let has_explicit_kind = upper.starts_with("TABLE ")
        || upper.starts_with("COLUMN ")
        || upper.starts_with("FUNCTION ")
        || upper.starts_with("TYPE ")
        || upper.starts_with("POLICY ")
        || upper.starts_with("CONSTRAINT ")
        || upper.starts_with("INDEX ")
        || upper.starts_with("SEQUENCE ")
        || upper.starts_with("VIEW ")
        || upper.starts_with("MATERIALIZED VIEW ")
        || upper.starts_with("SCHEMA ");

    if has_explicit_kind {
        buf.extend_from_slice(b"COMMENT ON ");
        buf.extend_from_slice(trimmed.as_bytes());
    } else if cmd.table.contains('.') {
        let mut parts = cmd.table.splitn(2, '.');
        let table = parts.next().unwrap_or_default();
        let col = parts.next().unwrap_or_default();
        buf.extend_from_slice(b"COMMENT ON COLUMN ");
        buf.extend_from_slice(table.as_bytes());
        buf.extend_from_slice(b".");
        buf.extend_from_slice(col.as_bytes());
    } else {
        buf.extend_from_slice(b"COMMENT ON TABLE ");
        buf.extend_from_slice(cmd.table.as_bytes());
    }

    buf.extend_from_slice(b" IS '");
    buf.extend_from_slice(escaped.as_bytes());
    buf.extend_from_slice(b"'");
}

/// Encode CREATE SEQUENCE statement.
pub fn encode_create_sequence(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"CREATE SEQUENCE ");
    buf.extend_from_slice(cmd.table.as_bytes());

    for col in &cmd.columns {
        if let Expr::Named(opt) = col {
            buf.extend_from_slice(b" ");
            buf.extend_from_slice(opt.as_bytes());
        }
    }
}

/// Encode DROP SEQUENCE statement.
pub fn encode_drop_sequence(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"DROP SEQUENCE IF EXISTS ");
    buf.extend_from_slice(cmd.table.as_bytes());
}

/// Encode CREATE TYPE ... AS ENUM statement.
pub fn encode_create_enum(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"CREATE TYPE ");
    buf.extend_from_slice(cmd.table.as_bytes());
    buf.extend_from_slice(b" AS ENUM (");

    let mut first = true;
    for col in &cmd.columns {
        if let Expr::Named(val) = col {
            if !first {
                buf.extend_from_slice(b", ");
            }
            first = false;
            buf.extend_from_slice(b"'");
            buf.extend_from_slice(val.replace('\'', "''").as_bytes());
            buf.extend_from_slice(b"'");
        }
    }

    buf.extend_from_slice(b")");
}

/// Encode DROP TYPE statement.
pub fn encode_drop_enum(cmd: &Qail, buf: &mut BytesMut) {
    buf.extend_from_slice(b"DROP TYPE IF EXISTS ");
    buf.extend_from_slice(cmd.table.as_bytes());
}

/// Encode ALTER TYPE ... ADD VALUE IF NOT EXISTS statement(s).
pub fn encode_alter_enum_add_value(cmd: &Qail, buf: &mut BytesMut) {
    let mut first = true;
    for col in &cmd.columns {
        if let Expr::Named(val) = col {
            if !first {
                buf.extend_from_slice(b"; ");
            }
            first = false;
            buf.extend_from_slice(b"ALTER TYPE ");
            buf.extend_from_slice(cmd.table.as_bytes());
            buf.extend_from_slice(b" ADD VALUE IF NOT EXISTS '");
            buf.extend_from_slice(val.replace('\'', "''").as_bytes());
            buf.extend_from_slice(b"'");
        }
    }
}

// ── Pub/Sub commands ───────────────────────────────────────────────

/// Encode LISTEN "channel".
pub fn encode_listen(cmd: &Qail, buf: &mut BytesMut) {
    let channel = cmd.channel.as_deref().unwrap_or("");
    buf.extend_from_slice(b"LISTEN \"");
    // Escape double-quotes in channel name
    buf.extend_from_slice(channel.replace('"', "\"\"").as_bytes());
    buf.extend_from_slice(b"\"");
}

/// Encode UNLISTEN "channel".
pub fn encode_unlisten(cmd: &Qail, buf: &mut BytesMut) {
    let channel = cmd.channel.as_deref().unwrap_or("");
    buf.extend_from_slice(b"UNLISTEN \"");
    buf.extend_from_slice(channel.replace('"', "\"\"").as_bytes());
    buf.extend_from_slice(b"\"");
}

/// Encode NOTIFY "channel", 'payload'.
pub fn encode_notify(cmd: &Qail, buf: &mut BytesMut) {
    let channel = cmd.channel.as_deref().unwrap_or("");
    buf.extend_from_slice(b"NOTIFY \"");
    buf.extend_from_slice(channel.replace('"', "\"\"").as_bytes());
    buf.extend_from_slice(b"\"");
    if let Some(ref payload) = cmd.payload {
        buf.extend_from_slice(b", '");
        buf.extend_from_slice(payload.replace('\'', "''").as_bytes());
        buf.extend_from_slice(b"'");
    }
}