use super::*;

pub(super) fn emit_type_instances(
    newtype: &ty::Newtype,
    context: &Context,
    output: &mut Vec<ty::Declaration>,
) {
    let Some(instances) = context.type_instances.get(&newtype.name) else {
        return;
    };

    for bindings in instances {
        let mangled = mangle_name(&newtype.name, bindings, &newtype.generic_parameters);
        let mut inner_type = substitute_type(&newtype.inner_type, bindings);
        rewrite_type(&mut inner_type, context);
        output.push(ty::Declaration::Type(ty::Newtype {
            name: mangled,
            generic_parameters: vec![],
            inner_type,
            public: newtype.public,
        }));
    }
}

pub(super) fn emit_impl_instances(
    impl_block: &ty::ImplBlock,
    context: &Context,
    output: &mut Vec<ty::Declaration>,
) -> Result<(), String> {
    let ty_types::Type::Named(target_name, _) = &impl_block.target_type else {
        return Err("generic impl target must be a named type".into());
    };

    let Some(instances) = context.type_instances.get(target_name) else {
        return Ok(());
    };

    for bindings in instances {
        let mangled_target = mangle_name(target_name, bindings, &impl_block.generic_parameters);

        let mut return_type = substitute_type(&impl_block.callable.return_type, bindings);
        rewrite_type(&mut return_type, context);
        let mut body = substitute_block(&impl_block.callable.body, bindings);
        rewrite_block(&mut body, context);

        let original_target = &impl_block.target_type;
        if let ty_types::Type::Named(_, target_args) = original_target
            && !target_args.is_empty()
        {
            let mut substituted_target = substitute_type(original_target, bindings);
            rewrite_type(&mut substituted_target, context);
            if let Some(original_inner) = context.generic_types.get(target_name) {
                let original_inner = &original_inner.inner_type;
                let mut substituted_inner = substitute_type(original_inner, bindings);
                rewrite_type(&mut substituted_inner, context);
                apply_field_renames_in_block(&mut body, original_inner, &substituted_inner);
            }
        }

        let parameters = impl_block.callable.parameters.as_ref().map(|params| {
            params
                .iter()
                .map(|parameter| {
                    let mut parameter_type = substitute_type(&parameter.parameter_type, bindings);
                    rewrite_type(&mut parameter_type, context);
                    ty::Parameter {
                        name: parameter.name.clone(),
                        parameter_type,
                    }
                })
                .collect()
        });

        output.push(ty::Declaration::Impl(ty::ImplBlock {
            generic_parameters: vec![],
            target_type: ty_types::Type::Named(mangled_target, vec![]),
            callable: ty::Callable {
                return_type,
                body,
                parameters,
            },
            const_fn: false,
        }));
    }

    Ok(())
}

fn apply_field_renames_in_block(
    block: &mut ty_expression::Block,
    original: &ty_types::Type,
    substituted: &ty_types::Type,
) {
    let renames = match (original, substituted) {
        (ty_types::Type::Tuple(orig), ty_types::Type::Tuple(subst))
            if orig.len() == subst.len() =>
        {
            let mut map = HashMap::new();
            for (orig_field, subst_field) in orig.iter().zip(subst) {
                if let (ty_types::Type::Named(orig_name, _), ty_types::Type::Named(subst_name, _)) =
                    (orig_field, subst_field)
                    && orig_name != subst_name
                {
                    map.insert(orig_name.clone(), subst_name.clone());
                }
            }
            map
        }
        _ => return,
    };
    if renames.is_empty() {
        return;
    }
    rename_fields_in_block(block, &renames);
}

fn rename_fields_in_block(block: &mut ty_expression::Block, renames: &HashMap<String, String>) {
    for statement in &mut block.statements {
        rename_fields_in_statement(statement, renames);
    }
    if let Some(result) = &mut block.result {
        rename_fields_in_expression(result, renames);
    }
}

fn rename_fields_in_statement(
    statement: &mut ty_expression::Statement,
    renames: &HashMap<String, String>,
) {
    use ty_expression::Statement;
    match statement {
        Statement::Expression(expression)
        | Statement::Return(Some(expression))
        | Statement::Let {
            value: expression, ..
        }
        | Statement::Assign(_, expression) => rename_fields_in_expression(expression, renames),
        Statement::While { condition, body } => {
            rename_fields_in_expression(condition, renames);
            rename_fields_in_block(body, renames);
        }
        Statement::Defer(inner) => rename_fields_in_statement(inner, renames),
        _ => {}
    }
}

fn rename_fields_in_expression(
    expression: &mut ty_expression::Expression,
    renames: &HashMap<String, String>,
) {
    use ty_expression::ExpressionKind as E;
    match &mut expression.kind {
        E::Field(object, field_name) => {
            if let Some(new_name) = renames.get(field_name.as_str()) {
                *field_name = new_name.clone();
            }
            rename_fields_in_expression(object, renames);
        }
        E::Call(callee, arguments) => {
            rename_fields_in_expression(callee, renames);
            for argument in arguments {
                rename_fields_in_expression(argument, renames);
            }
        }
        E::BinaryOperation(_, left, right) | E::Replace(left, right) => {
            rename_fields_in_expression(left, renames);
            rename_fields_in_expression(right, renames);
        }
        E::UnaryOperation(_, operand)
        | E::Dereference(operand)
        | E::Convert(operand, _)
        | E::Eval(operand)
        | E::ArrayRepeat(operand, _) => {
            rename_fields_in_expression(operand, renames);
        }
        E::Index(object, index) => {
            rename_fields_in_expression(object, renames);
            rename_fields_in_expression(index, renames);
        }
        E::Block(block) => rename_fields_in_block(block, renames),
        E::If {
            condition,
            then_branch,
            else_branch,
        } => {
            rename_fields_in_expression(condition, renames);
            rename_fields_in_block(then_branch, renames);
            if let Some(else_branch) = else_branch {
                rename_fields_in_block(else_branch, renames);
            }
        }
        E::ArrayLiteral(elements) | E::TupleLiteral(elements) | E::Print(elements) => {
            for element in elements {
                rename_fields_in_expression(element, renames);
            }
        }
        E::TypeConstruction(_, _, fields) => {
            for (field_name, value) in fields {
                if let Some(new_name) = renames.get(field_name.as_str()) {
                    *field_name = new_name.clone();
                }
                rename_fields_in_expression(value, renames);
            }
        }
        _ => {}
    }
}
