mod columns;
mod fk;
mod type_names;

use pg_query::{
    NodeEnum,
    protobuf::{AlterTableType, Node},
};
use text_size::TextRange;

use crate::core::schema::{ColumnDef, QualifiedName, SchemaModel, TableDef};

pub(crate) use columns::build_column_def;
pub(crate) use type_names::{type_name_to_qn, type_name_to_string};

pub(super) fn handle_create_table(
    stmt: &pg_query::protobuf::CreateStmt,
    range: TextRange,
    model: &mut SchemaModel,
) {
    let Some(rel) = &stmt.relation else {
        tracing::warn!("CreateStmt has no relation, skipping");
        return;
    };
    let key = range_var_to_key(rel);
    let columns = extract_columns(&stmt.table_elts, range.start());
    model.tables.insert(
        key,
        TableDef {
            columns,
            range,
            doc: None,
        },
    );
}

pub(super) fn handle_alter_table(
    stmt: &pg_query::protobuf::AlterTableStmt,
    stmt_start: text_size::TextSize,
    model: &mut SchemaModel,
) {
    let Some(rel) = &stmt.relation else { return };
    let key = range_var_to_key(rel);

    for cmd_node in &stmt.cmds {
        let Some(NodeEnum::AlterTableCmd(cmd)) = cmd_node.node.as_ref() else {
            continue;
        };
        let Ok(subtype) = AlterTableType::try_from(cmd.subtype) else {
            continue;
        };

        match subtype {
            AlterTableType::AtAddColumn | AlterTableType::AtAddColumnToView => {
                apply_add_column(&key, cmd, model, stmt_start);
            }
            AlterTableType::AtDropColumn => {
                apply_drop_column(&key, &cmd.name, model);
            }
            AlterTableType::AtAddConstraint => {
                fk::apply_add_fk_constraint(&key, cmd, model);
            }
            other => {
                tracing::debug!(?other, "AlterTableCmd subtype not handled, skipping");
            }
        }
    }
}

fn extract_columns(
    table_elts: &[Node],
    stmt_start: text_size::TextSize,
) -> indexmap::IndexMap<String, ColumnDef> {
    let mut cols = indexmap::IndexMap::new();
    for node in table_elts {
        let Some(inner) = node.node.as_ref() else {
            continue;
        };
        match inner {
            NodeEnum::ColumnDef(c) => {
                let (name, col) = build_column_def(c, stmt_start);
                cols.insert(name, col);
            }
            NodeEnum::Constraint(c) => {
                fk::apply_table_fk_constraint(c, &mut cols);
            }
            _ => {}
        }
    }
    cols
}

fn apply_add_column(
    key: &QualifiedName,
    cmd: &pg_query::protobuf::AlterTableCmd,
    model: &mut SchemaModel,
    stmt_start: text_size::TextSize,
) {
    let Some(NodeEnum::ColumnDef(col_def)) = cmd.def.as_ref().and_then(|b| b.node.as_ref()) else {
        return;
    };
    let (col_name, col) = build_column_def(col_def, stmt_start);
    if let Some(table) = model.tables.get_mut(key) {
        table.columns.insert(col_name, col);
    }
}

fn apply_drop_column(key: &QualifiedName, col_name: &str, model: &mut SchemaModel) {
    if let Some(table) = model.tables.get_mut(key) {
        table.columns.shift_remove(col_name);
    }
}

pub(crate) fn range_var_to_key(rel: &pg_query::protobuf::RangeVar) -> QualifiedName {
    let schema = if rel.schemaname.is_empty() {
        None
    } else {
        Some(rel.schemaname.as_str())
    };
    QualifiedName::new(schema, &rel.relname)
}

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

    fn walk(sql: &str) -> SchemaModel {
        let (stmts, _) = split_and_parse(sql).expect("must parse");
        let mut model = SchemaModel::default();
        super::super::walk_ddl(&stmts, &mut model);
        model
    }

    #[test]
    fn create_table_inserts_table_def() {
        let model = walk("CREATE TABLE public.t1 (id bigint NOT NULL);");
        assert!(
            model
                .tables
                .contains_key(&QualifiedName::new(Some("public"), "t1"))
        );
    }
}
