use pg_query::{NodeEnum, protobuf::FunctionParameterMode};
use text_size::TextRange;

use super::util::parse_name;

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

use super::tables::{type_name_to_qn, type_name_to_string};

pub(super) fn handle_create_function(
    stmt: &pg_query::protobuf::CreateFunctionStmt,
    range: TextRange,
    model: &mut SchemaModel,
) {
    let Some((schema, name)) = parse_name(&stmt.funcname) else {
        tracing::warn!("CreateFunctionStmt has no parseable name, skipping");
        return;
    };
    let key = QualifiedName::new(schema.as_deref(), &name);
    let return_type = stmt.return_type.as_ref().and_then(type_name_to_qn);
    let param_types = extract_param_types(&stmt.parameters);
    let params = extract_params(&stmt.parameters);
    model.functions.insert(
        key,
        FunctionDef {
            range,
            doc: None,
            return_type,
            param_types,
            params,
        },
    );
}

// RETURNS TABLE(...) columns are encoded as FuncParamTable — exclude from param_types and signature
fn is_table_column(p: &pg_query::protobuf::FunctionParameter) -> bool {
    FunctionParameterMode::try_from(p.mode) == Ok(FunctionParameterMode::FuncParamTable)
}

fn extract_param_types(parameters: &[pg_query::protobuf::Node]) -> Vec<QualifiedName> {
    parameters
        .iter()
        .filter_map(|n| n.node.as_ref())
        .filter_map(|n| {
            if let NodeEnum::FunctionParameter(p) = n {
                if is_table_column(p) {
                    return None;
                }
                p.arg_type.as_ref().and_then(type_name_to_qn)
            } else {
                None
            }
        })
        .collect()
}

fn extract_params(parameters: &[pg_query::protobuf::Node]) -> Vec<(String, String)> {
    parameters
        .iter()
        .filter_map(|n| n.node.as_ref())
        .filter_map(|n| {
            if let NodeEnum::FunctionParameter(p) = n {
                if is_table_column(p) {
                    return None;
                }
                let ty = p
                    .arg_type
                    .as_ref()
                    .map(type_name_to_string)
                    .unwrap_or_default();
                Some((p.name.clone(), ty))
            } else {
                None
            }
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use crate::{
        core::schema::{QualifiedName, SchemaModel},
        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_function_inserts_function_def() {
        let model = walk(
            "CREATE FUNCTION public.f1(x bigint) RETURNS bigint LANGUAGE sql AS $$ SELECT x $$;",
        );
        assert!(
            model
                .functions
                .contains_key(&QualifiedName::new(Some("public"), "f1"))
        );
    }

    #[test]
    fn function_with_custom_return_type_sets_return_type() {
        let model = walk(concat!(
            "CREATE TYPE public.item_color AS ENUM ('red');",
            "CREATE FUNCTION public.get_color(p bigint) RETURNS public.item_color LANGUAGE sql AS $$ SELECT 'red'::public.item_color $$;",
        ));
        let fn_key = QualifiedName::new(Some("public"), "get_color");
        let func = &model.functions[&fn_key];
        let rt = func.return_type.as_ref().expect("must have return_type");
        assert_eq!(rt.schema, "public");
        assert_eq!(rt.name, "item_color");
    }

    #[test]
    fn function_with_builtin_return_type_has_no_return_type() {
        let model = walk(
            "CREATE FUNCTION public.f(p bigint) RETURNS bigint LANGUAGE sql AS $$ SELECT p $$;",
        );
        let fn_key = QualifiedName::new(Some("public"), "f");
        assert!(model.functions[&fn_key].return_type.is_none());
    }

    #[test]
    fn returns_void_has_no_return_type() {
        let model = walk("CREATE FUNCTION public.f() RETURNS void LANGUAGE sql AS $$ SELECT 1 $$;");
        let fn_key = QualifiedName::new(Some("public"), "f");
        assert!(
            model.functions[&fn_key].return_type.is_none(),
            "void must not produce a return_type"
        );
    }

    #[test]
    fn returns_trigger_has_no_return_type() {
        let model = walk(
            "CREATE FUNCTION public.f() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN RETURN NEW; END $$;",
        );
        let fn_key = QualifiedName::new(Some("public"), "f");
        assert!(
            model.functions[&fn_key].return_type.is_none(),
            "trigger must not produce a return_type"
        );
    }

    #[test]
    fn returns_table_columns_excluded_from_param_types() {
        let model = walk(concat!(
            "CREATE TYPE public.my_type AS ENUM ('a');",
            "CREATE FUNCTION public.f() RETURNS TABLE(x bigint, y public.my_type) LANGUAGE sql AS $$ SELECT 1, 'a' $$;",
        ));
        let fn_key = QualifiedName::new(Some("public"), "f");
        let func = &model.functions[&fn_key];
        assert!(
            func.param_types.is_empty(),
            "TABLE output columns must not appear in param_types; got {:?}",
            func.param_types,
        );
    }

    #[test]
    fn in_param_with_user_type_still_included() {
        let model = walk(concat!(
            "CREATE TYPE public.my_type AS ENUM ('a');",
            "CREATE FUNCTION public.f(p public.my_type) RETURNS bigint LANGUAGE sql AS $$ SELECT 1 $$;",
        ));
        let fn_key = QualifiedName::new(Some("public"), "f");
        let func = &model.functions[&fn_key];
        assert_eq!(
            func.param_types.len(),
            1,
            "IN param with user type must be in param_types"
        );
        assert_eq!(func.param_types[0].name, "my_type");
    }

    #[test]
    fn function_params_populated_with_name_and_type() {
        let model = walk(
            "CREATE FUNCTION public.f(p_id bigint, p_name text) RETURNS void LANGUAGE sql AS $$ SELECT 1 $$;",
        );
        let fn_key = QualifiedName::new(Some("public"), "f");
        let func = &model.functions[&fn_key];
        assert_eq!(func.params.len(), 2);
        assert_eq!(func.params[0].0, "p_id");
        assert_eq!(func.params[0].1, "int8");
        assert_eq!(func.params[1].0, "p_name");
        assert_eq!(func.params[1].1, "text");
    }

    #[test]
    fn returns_table_columns_excluded_from_params() {
        let model = walk(
            "CREATE FUNCTION public.f() RETURNS TABLE(x bigint) LANGUAGE sql AS $$ SELECT 1 $$;",
        );
        let fn_key = QualifiedName::new(Some("public"), "f");
        assert!(
            model.functions[&fn_key].params.is_empty(),
            "TABLE output columns must not appear in params"
        );
    }

    #[test]
    fn zero_param_function_has_empty_params() {
        let model = walk("CREATE FUNCTION public.f() RETURNS void LANGUAGE sql AS $$ SELECT 1 $$;");
        let fn_key = QualifiedName::new(Some("public"), "f");
        assert!(model.functions[&fn_key].params.is_empty());
    }
}
