use super::*;

pub(super) fn compute_nongeneric_field_renames(
    type_name: &str,
    original: &ty_types::Type,
    rewritten: &ty_types::Type,
    context: &mut Context,
) {
    let (original_fields, rewritten_fields) = match (original, rewritten) {
        (ty_types::Type::Tuple(orig), ty_types::Type::Tuple(rewr)) if orig.len() == rewr.len() => {
            (orig, rewr)
        }
        _ => return,
    };
    let mut renames = HashMap::new();
    for (orig, rewr) in original_fields.iter().zip(rewritten_fields) {
        if let (ty_types::Type::Named(orig_name, _), ty_types::Type::Named(rewr_name, _)) =
            (orig, rewr)
            && orig_name != rewr_name
        {
            renames.insert(orig_name.clone(), rewr_name.clone());
        }
    }
    if !renames.is_empty() {
        context.field_renames.insert(type_name.into(), renames);
    }
}

pub(super) fn rewrite_block(block: &mut ty_expression::Block, context: &Context) {
    for statement in &mut block.statements {
        rewrite_statement(statement, context);
    }
    if let Some(result) = &mut block.result {
        rewrite_expression(result, context);
    }
}

pub(super) fn rewrite_statement(statement: &mut ty_expression::Statement, context: &Context) {
    use ty_expression::Statement as S;
    match statement {
        S::Expression(expression) => rewrite_expression(expression, context),
        S::Let {
            value,
            declared_type,
            ..
        } => {
            rewrite_expression(value, context);
            if let Some(declared_type) = declared_type {
                rewrite_type(declared_type, context);
            }
        }
        S::Assign(target, value) => {
            rewrite_expression(target, context);
            rewrite_expression(value, context);
        }
        S::Return(expression) => {
            if let Some(expression) = expression {
                rewrite_expression(expression, context);
            }
        }
        S::Label {
            initial_arguments, ..
        } => {
            for argument in initial_arguments {
                rewrite_expression(argument, context);
            }
        }
        S::Jump { arguments, .. } => {
            for argument in arguments {
                rewrite_expression(argument, context);
            }
        }
        S::While { condition, body } => {
            rewrite_expression(condition, context);
            rewrite_block(body, context);
        }
        S::Defer(inner) => rewrite_statement(inner, context),
    }
}

pub(super) fn rewrite_expression(expression: &mut ty_expression::Expression, context: &Context) {
    rewrite_expression_kind(&mut expression.kind, context);
    rewrite_type(&mut expression.resolved_type, context);
}

pub(super) fn rewrite_expression_kind(kind: &mut ty_expression::ExpressionKind, context: &Context) {
    use ty_expression::ExpressionKind as E;
    match kind {
        E::Call(callee, arguments) => {
            let mangled = try_mangle_call(callee, arguments, context);
            if let Some(mangled_name) = mangled {
                callee.kind = E::Variable(mangled_name);
            }
            rewrite_expression(callee, context);
            for argument in arguments {
                rewrite_expression(argument, context);
            }
        }
        E::BinaryOperation(_, left, right) => {
            rewrite_expression(left, context);
            rewrite_expression(right, context);
        }
        E::UnaryOperation(_, operand) => rewrite_expression(operand, context),
        E::Field(object, field_name) => {
            rewrite_expression(object, context);
            if let ty_types::Type::Named(type_name, _) = &object.resolved_type
                && let Some(renames) = context.field_renames.get(type_name.as_str())
                && let Some(mangled_field) = renames.get(field_name.as_str())
            {
                *field_name = mangled_field.clone();
            }
        }
        E::Index(object, index) => {
            rewrite_expression(object, context);
            rewrite_expression(index, context);
        }
        E::Specialize(inner, _)
        | E::Dereference(inner)
        | E::Eval(inner)
        | E::ArrayRepeat(inner, _) => {
            rewrite_expression(inner, context);
        }
        E::Convert(inner, target) => {
            rewrite_expression(inner, context);
            rewrite_type(target, context);
        }
        E::TypeConstruction(name, generic_arguments, fields) => {
            if !generic_arguments.is_empty()
                && let Some(newtype) = context.generic_types.get(name.as_str())
                && let Some(bindings) =
                    bindings_from_explicit_arguments(&newtype.generic_parameters, generic_arguments)
            {
                *name = mangle_name(name, &bindings, &newtype.generic_parameters);
                *generic_arguments = vec![];
                if let Some(renames) = context.field_renames.get(name.as_str()) {
                    for (field_name, _) in fields.iter_mut() {
                        if let Some(mangled_field) = renames.get(field_name.as_str()) {
                            *field_name = mangled_field.clone();
                        }
                    }
                }
            }
            for (_, expression) in fields {
                rewrite_expression(expression, context);
            }
        }
        E::ArrayLiteral(elements) | E::TupleLiteral(elements) | E::Print(elements) => {
            for element in elements {
                rewrite_expression(element, context);
            }
        }
        E::Block(block) => rewrite_block(block, context),
        E::If {
            condition,
            then_branch,
            else_branch,
        } => {
            rewrite_expression(condition, context);
            rewrite_block(then_branch, context);
            if let Some(else_branch) = else_branch {
                rewrite_block(else_branch, context);
            }
        }
        E::Match { value, arms } => {
            rewrite_expression(value, context);
            for arm in arms {
                rewrite_block(&mut arm.body, context);
            }
        }
        E::Replace(target, value) => {
            rewrite_expression(target, context);
            rewrite_expression(value, context);
        }
        E::SizeOf(target) => rewrite_type(target, context),
        E::Literal(_) | E::Variable(_) => {}
    }
}

pub(super) fn rewrite_type(rewritten_type: &mut ty_types::Type, context: &Context) {
    match rewritten_type {
        ty_types::Type::Named(name, generic_arguments) if !generic_arguments.is_empty() => {
            for argument in generic_arguments.iter_mut() {
                if let ty_types::GenericArgument::Type(inner) = argument {
                    rewrite_type(inner, context);
                }
            }
            if let Some(newtype) = context.generic_types.get(name.as_str())
                && let Some(bindings) =
                    bindings_from_explicit_arguments(&newtype.generic_parameters, generic_arguments)
            {
                *name = mangle_name(name, &bindings, &newtype.generic_parameters);
                *generic_arguments = vec![];
            }
        }
        ty_types::Type::Pointer(_, _, inner)
        | ty_types::Type::Reference(_, _, inner)
        | ty_types::Type::Array(inner, _)
        | ty_types::Type::Vector(inner, _)
        | ty_types::Type::Slice(_, inner) => rewrite_type(inner, context),
        ty_types::Type::Tuple(fields) | ty_types::Type::Enum(fields) => {
            for field in fields {
                rewrite_type(field, context);
            }
        }
        ty_types::Type::Function(signature) => {
            for parameter in &mut signature.parameters {
                rewrite_type(parameter, context);
            }
            rewrite_type(&mut signature.return_type, context);
        }
        _ => {}
    }
}

pub(super) fn try_mangle_call(
    callee: &ty_expression::Expression,
    arguments: &[ty_expression::Expression],
    context: &Context,
) -> Option<String> {
    let (raw_name, explicit_arguments) = match &callee.kind {
        ty_expression::ExpressionKind::Variable(name) => (name.as_str(), None),
        ty_expression::ExpressionKind::Specialize(inner, generic_arguments) => {
            if let ty_expression::ExpressionKind::Variable(name) = &inner.kind {
                (name.as_str(), Some(generic_arguments.as_slice()))
            } else {
                return None;
            }
        }
        _ => return None,
    };

    let name = raw_name.strip_prefix("impl_").unwrap_or(raw_name);

    if let Some(generic_type) = context.generic_types.get(name) {
        let generic_parameters = if let Some(impl_block) = context.generic_impls.get(name) {
            &impl_block.generic_parameters
        } else {
            &generic_type.generic_parameters
        };

        let generic_names: Vec<&str> = generic_parameters
            .iter()
            .map(|parameter| parameter.name.as_str())
            .collect();

        let mut bindings = if let Some(explicit) = explicit_arguments {
            bindings_from_explicit_arguments(generic_parameters, explicit)?
        } else {
            Bindings::new()
        };

        unify_type_fields(
            &generic_type.inner_type,
            arguments,
            &mut bindings,
            &generic_names,
        );

        if generic_parameters
            .iter()
            .all(|parameter| bindings.contains(&parameter.name))
        {
            let mangled = mangle_name(name, &bindings, generic_parameters);
            if raw_name.starts_with("impl_") {
                return Some(format!("impl_{mangled}"));
            }
            return Some(mangled);
        }
    }

    None
}

pub(super) fn unify_type_fields(
    inner_type: &ty_types::Type,
    arguments: &[ty_expression::Expression],
    bindings: &mut Bindings,
    generic_names: &[&str],
) {
    match inner_type {
        ty_types::Type::Tuple(fields) => {
            for (field_type, argument) in fields.iter().zip(arguments) {
                let _ = unify(field_type, &argument.resolved_type, bindings, generic_names);
            }
        }
        single => {
            if let Some(argument) = arguments.first() {
                let _ = unify(single, &argument.resolved_type, bindings, generic_names);
            }
        }
    }
}
