use pg_query::{NodeEnum, protobuf::ConstrType};

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

use super::range_var_to_key;

pub(super) fn apply_table_fk_constraint(
    c: &pg_query::protobuf::Constraint,
    cols: &mut indexmap::IndexMap<String, ColumnDef>,
) {
    if ConstrType::try_from(c.contype) != Ok(ConstrType::ConstrForeign) {
        return;
    }
    let Some(rv) = c.pktable.as_ref() else { return };
    let ref_table = range_var_to_key(rv);
    let ref_col = extract_first_name(&c.pk_attrs);
    let src_col_name = extract_first_name(&c.fk_attrs);
    if let Some(col) = src_col_name.and_then(|name| cols.get_mut(&name)) {
        col.fk_target = Some((ref_table, ref_col.unwrap_or_default()));
    }
}

pub(super) fn apply_add_fk_constraint(
    key: &QualifiedName,
    cmd: &pg_query::protobuf::AlterTableCmd,
    model: &mut SchemaModel,
) {
    let Some(def_node) = cmd.def.as_ref() else {
        return;
    };
    let Some(NodeEnum::Constraint(c)) = def_node.node.as_ref() else {
        return;
    };
    if ConstrType::try_from(c.contype) != Ok(ConstrType::ConstrForeign) {
        return;
    }
    let Some(rv) = c.pktable.as_ref() else { return };
    let ref_table = range_var_to_key(rv);
    let ref_col = extract_first_name(&c.pk_attrs).unwrap_or_default();
    let Some(col_name) = extract_first_name(&c.fk_attrs) else {
        return;
    };
    let Some(table) = model.tables.get_mut(key) else {
        return;
    };
    if let Some(col) = table.columns.get_mut(&col_name) {
        col.fk_target = Some((ref_table, ref_col));
    }
}

pub(super) fn extract_inline_fk(
    col: &pg_query::protobuf::ColumnDef,
) -> Option<(QualifiedName, String)> {
    for cnode in &col.constraints {
        let Some(NodeEnum::Constraint(c)) = cnode.node.as_ref() else {
            continue;
        };
        if ConstrType::try_from(c.contype) != Ok(ConstrType::ConstrForeign) {
            continue;
        }
        let rv = c.pktable.as_ref()?;
        let ref_table = range_var_to_key(rv);
        let ref_col = extract_first_name(&c.pk_attrs).unwrap_or_default();
        return Some((ref_table, ref_col));
    }
    None
}

fn extract_first_name(nodes: &[pg_query::protobuf::Node]) -> Option<String> {
    nodes.first().and_then(|n| n.node.as_ref()).and_then(|n| {
        if let NodeEnum::String(s) = n {
            Some(s.sval.clone())
        } else {
            None
        }
    })
}

#[cfg(test)]
mod tests {
    use crate::{core::schema::SchemaModel, parse::split_and_parse};

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

    #[test]
    fn inline_fk_sets_fk_target() {
        use crate::core::schema::QualifiedName;
        let model = walk(concat!(
            "CREATE TABLE public.items (id bigint NOT NULL);",
            "CREATE TABLE public.orders (id bigint NOT NULL, item_id bigint NOT NULL REFERENCES public.items(id));",
        ));
        let col = &model.tables[&QualifiedName::new(Some("public"), "orders")].columns["item_id"];
        let (rt, rc) = col.fk_target.as_ref().expect("must have fk_target");
        assert_eq!(rt.name, "items");
        assert_eq!(rc, "id");
    }

    #[test]
    fn alter_table_add_fk_sets_fk_target() {
        use crate::core::schema::QualifiedName;
        let model = walk(concat!(
            "CREATE TABLE public.items (id bigint NOT NULL);",
            "CREATE TABLE public.orders (id bigint NOT NULL, item_id bigint NOT NULL);",
            "ALTER TABLE public.orders ADD CONSTRAINT orders_item_id_fkey FOREIGN KEY (item_id) REFERENCES public.items(id);",
        ));
        let col = &model.tables[&QualifiedName::new(Some("public"), "orders")].columns["item_id"];
        let (rt, rc) = col.fk_target.as_ref().expect("must have fk_target");
        assert_eq!(rt.name, "items");
        assert_eq!(rc, "id");
    }
}
