use super::*;

pub(super) fn substitute_const(
    constant: &ty_types::ConstExpression,
    bindings: &Bindings,
) -> ty_types::ConstExpression {
    use crate::const_expression::{EvalContext, eval};
    let mut context = EvalContext {
        bindings: Some(&bindings.constants),
        evaluator: None,
    };
    eval(constant, &mut context)
}

pub(super) fn substitute_type(source: &ty_types::Type, bindings: &Bindings) -> ty_types::Type {
    use ty_types::Type;
    match source {
        Type::Generic(name) => bindings
            .types
            .get(name)
            .cloned()
            .unwrap_or_else(|| source.clone()),
        Type::Named(name, arguments)
            if arguments.is_empty() && bindings.types.contains_key(name) =>
        {
            bindings.types[name].clone()
        }
        Type::Pointer(mutability, lifetime, inner) => Type::Pointer(
            *mutability,
            lifetime.clone(),
            Box::new(substitute_type(inner, bindings)),
        ),
        Type::Reference(mutability, lifetime, inner) => Type::Reference(
            *mutability,
            lifetime.clone(),
            Box::new(substitute_type(inner, bindings)),
        ),
        Type::Array(element, size) => Type::Array(
            Box::new(substitute_type(element, bindings)),
            substitute_const(size, bindings),
        ),
        Type::Vector(element, size) => Type::Vector(
            Box::new(substitute_type(element, bindings)),
            substitute_const(size, bindings),
        ),
        Type::Slice(mutability, element) => {
            Type::Slice(*mutability, Box::new(substitute_type(element, bindings)))
        }
        Type::Tuple(fields) => Type::Tuple(
            fields
                .iter()
                .map(|field| substitute_type(field, bindings))
                .collect(),
        ),
        Type::Enum(variants) => Type::Enum(
            variants
                .iter()
                .map(|variant| substitute_type(variant, bindings))
                .collect(),
        ),
        Type::Named(name, generic_arguments) => {
            let substituted_arguments: Vec<_> = generic_arguments
                .iter()
                .map(|argument| substitute_generic_argument(argument, bindings))
                .collect();
            Type::Named(name.clone(), substituted_arguments)
        }
        Type::Function(signature) => Type::Function(ty_types::FunctionSignature {
            parameters: signature
                .parameters
                .iter()
                .map(|parameter| substitute_type(parameter, bindings))
                .collect(),
            return_type: Box::new(substitute_type(&signature.return_type, bindings)),
        }),
        other => other.clone(),
    }
}

pub(super) fn substitute_generic_argument(
    argument: &ty_types::GenericArgument,
    bindings: &Bindings,
) -> ty_types::GenericArgument {
    match argument {
        ty_types::GenericArgument::Type(inner) => {
            ty_types::GenericArgument::Type(substitute_type(inner, bindings))
        }
        ty_types::GenericArgument::Const(constant) => {
            ty_types::GenericArgument::Const(substitute_const(constant, bindings))
        }
        lifetime @ ty_types::GenericArgument::Lifetime(_) => lifetime.clone(),
    }
}

pub(super) fn substitute_expression(
    expression: &ty_expression::Expression,
    bindings: &Bindings,
) -> ty_expression::Expression {
    ty_expression::Expression::new(
        substitute_expression_kind(&expression.kind, bindings),
        substitute_type(&expression.resolved_type, bindings),
    )
}

pub(super) fn substitute_expression_kind(
    kind: &ty_expression::ExpressionKind,
    bindings: &Bindings,
) -> ty_expression::ExpressionKind {
    use ty_expression::ExpressionKind as E;
    match kind {
        E::Literal(literal) => E::Literal(literal.clone()),
        E::Variable(name) => {
            if let Some(constant) = bindings.constants.get(name) {
                let evaluated = substitute_const(constant, bindings);
                if let ty_types::ConstExpression::Integer(value) = evaluated {
                    return E::Literal(ty_expression::Literal::Integer(value));
                }
            }
            E::Variable(name.clone())
        }
        E::BinaryOperation(operator, left, right) => E::BinaryOperation(
            *operator,
            Box::new(substitute_expression(left, bindings)),
            Box::new(substitute_expression(right, bindings)),
        ),
        E::UnaryOperation(operator, operand) => E::UnaryOperation(
            *operator,
            Box::new(substitute_expression(operand, bindings)),
        ),
        E::Call(callee, arguments) => E::Call(
            Box::new(substitute_expression(callee, bindings)),
            arguments
                .iter()
                .map(|argument| substitute_expression(argument, bindings))
                .collect(),
        ),
        E::Field(object, field) => E::Field(
            Box::new(substitute_expression(object, bindings)),
            field.clone(),
        ),
        E::Index(object, index) => E::Index(
            Box::new(substitute_expression(object, bindings)),
            Box::new(substitute_expression(index, bindings)),
        ),
        E::Specialize(inner, generic_arguments) => E::Specialize(
            Box::new(substitute_expression(inner, bindings)),
            generic_arguments
                .iter()
                .map(|argument| substitute_generic_argument(argument, bindings))
                .collect(),
        ),
        E::Dereference(inner) => E::Dereference(Box::new(substitute_expression(inner, bindings))),
        E::Eval(inner) => E::Eval(Box::new(substitute_expression(inner, bindings))),
        E::ArrayRepeat(value, count) => E::ArrayRepeat(
            Box::new(substitute_expression(value, bindings)),
            substitute_const(count, bindings),
        ),
        E::Convert(inner, target) => E::Convert(
            Box::new(substitute_expression(inner, bindings)),
            substitute_type(target, bindings),
        ),
        E::SizeOf(target) => E::SizeOf(substitute_type(target, bindings)),
        E::TypeConstruction(name, generic_arguments, fields) => E::TypeConstruction(
            name.clone(),
            generic_arguments
                .iter()
                .map(|argument| substitute_generic_argument(argument, bindings))
                .collect(),
            fields
                .iter()
                .map(|(field_name, expression)| {
                    (
                        field_name.clone(),
                        substitute_expression(expression, bindings),
                    )
                })
                .collect(),
        ),
        E::ArrayLiteral(elements) => E::ArrayLiteral(
            elements
                .iter()
                .map(|element| substitute_expression(element, bindings))
                .collect(),
        ),
        E::TupleLiteral(elements) => E::TupleLiteral(
            elements
                .iter()
                .map(|element| substitute_expression(element, bindings))
                .collect(),
        ),
        E::Block(block) => E::Block(substitute_block(block, bindings)),
        E::If {
            condition,
            then_branch,
            else_branch,
        } => E::If {
            condition: Box::new(substitute_expression(condition, bindings)),
            then_branch: substitute_block(then_branch, bindings),
            else_branch: else_branch
                .as_ref()
                .map(|branch| substitute_block(branch, bindings)),
        },
        E::Match { value, arms } => E::Match {
            value: Box::new(substitute_expression(value, bindings)),
            arms: arms
                .iter()
                .map(|arm| ty_expression::MatchArm {
                    pattern: arm.pattern.clone(),
                    body: substitute_block(&arm.body, bindings),
                })
                .collect(),
        },
        E::Replace(target, value) => E::Replace(
            Box::new(substitute_expression(target, bindings)),
            Box::new(substitute_expression(value, bindings)),
        ),
        E::Print(arguments) => E::Print(
            arguments
                .iter()
                .map(|argument| substitute_expression(argument, bindings))
                .collect(),
        ),
    }
}

pub(super) fn substitute_block(
    block: &ty_expression::Block,
    bindings: &Bindings,
) -> ty_expression::Block {
    ty_expression::Block {
        statements: block
            .statements
            .iter()
            .map(|statement| substitute_statement(statement, bindings))
            .collect(),
        result: block
            .result
            .as_ref()
            .map(|expression| Box::new(substitute_expression(expression, bindings))),
    }
}

pub(super) fn substitute_statement(
    statement: &ty_expression::Statement,
    bindings: &Bindings,
) -> ty_expression::Statement {
    use ty_expression::Statement as S;
    match statement {
        S::Expression(expression) => S::Expression(substitute_expression(expression, bindings)),
        S::Let {
            name,
            binding,
            declared_type,
            value,
            line,
        } => S::Let {
            name: name.clone(),
            binding: *binding,
            declared_type: declared_type
                .as_ref()
                .map(|declared_type| substitute_type(declared_type, bindings)),
            value: substitute_expression(value, bindings),
            line: *line,
        },
        S::Assign(target, value) => S::Assign(
            substitute_expression(target, bindings),
            substitute_expression(value, bindings),
        ),
        S::Return(expression) => S::Return(
            expression
                .as_ref()
                .map(|expression| substitute_expression(expression, bindings)),
        ),
        S::Label {
            name,
            parameters,
            initial_arguments,
        } => S::Label {
            name: name.clone(),
            parameters: parameters
                .iter()
                .map(|parameter| ty::Parameter {
                    name: parameter.name.clone(),
                    parameter_type: substitute_type(&parameter.parameter_type, bindings),
                })
                .collect(),
            initial_arguments: initial_arguments
                .iter()
                .map(|argument| substitute_expression(argument, bindings))
                .collect(),
        },
        S::Jump { label, arguments } => S::Jump {
            label: label.clone(),
            arguments: arguments
                .iter()
                .map(|argument| substitute_expression(argument, bindings))
                .collect(),
        },
        S::While { condition, body } => S::While {
            condition: substitute_expression(condition, bindings),
            body: substitute_block(body, bindings),
        },
        S::Defer(inner) => S::Defer(Box::new(substitute_statement(inner, bindings))),
    }
}

pub(super) fn substitute_named_constants_in_declaration(
    declaration: ty::Declaration,
    constants: &HashMap<String, u128>,
) -> ty::Declaration {
    match declaration {
        ty::Declaration::Type(mut newtype) => {
            substitute_named_constants_in_type(&mut newtype.inner_type, constants);
            ty::Declaration::Type(newtype)
        }
        ty::Declaration::Impl(mut impl_block) => {
            substitute_named_constants_in_block(&mut impl_block.callable.body, constants);
            ty::Declaration::Impl(impl_block)
        }
        other => other,
    }
}

fn substitute_named_constants_in_block(
    block: &mut ty_expression::Block,
    constants: &HashMap<String, u128>,
) {
    for statement in &mut block.statements {
        substitute_named_constants_in_statement(statement, constants);
    }
    if let Some(result) = &mut block.result {
        substitute_named_constants_in_expression(result, constants);
    }
}

fn substitute_named_constants_in_statement(
    statement: &mut ty_expression::Statement,
    constants: &HashMap<String, u128>,
) {
    use ty_expression::Statement as S;
    match statement {
        S::Expression(expression) | S::Return(Some(expression)) => {
            substitute_named_constants_in_expression(expression, constants);
        }
        S::Let {
            value,
            declared_type,
            ..
        } => {
            substitute_named_constants_in_expression(value, constants);
            if let Some(declared_type) = declared_type {
                substitute_named_constants_in_type(declared_type, constants);
            }
        }
        S::Assign(target, value) => {
            substitute_named_constants_in_expression(target, constants);
            substitute_named_constants_in_expression(value, constants);
        }
        S::While { condition, body } => {
            substitute_named_constants_in_expression(condition, constants);
            substitute_named_constants_in_block(body, constants);
        }
        S::Label {
            initial_arguments, ..
        } => {
            for argument in initial_arguments {
                substitute_named_constants_in_expression(argument, constants);
            }
        }
        S::Jump { arguments, .. } => {
            for argument in arguments {
                substitute_named_constants_in_expression(argument, constants);
            }
        }
        S::Defer(inner) => substitute_named_constants_in_statement(inner, constants),
        S::Return(None) => {}
    }
}

fn substitute_named_constants_in_expression(
    expression: &mut ty_expression::Expression,
    constants: &HashMap<String, u128>,
) {
    substitute_named_constants_in_type(&mut expression.resolved_type, constants);
    match &mut expression.kind {
        ty_expression::ExpressionKind::Specialize(inner, generic_arguments) => {
            substitute_named_constants_in_expression(inner, constants);
            for argument in generic_arguments {
                if let ty_types::GenericArgument::Const(constant) = argument {
                    substitute_named_constant(constant, constants);
                }
            }
        }
        ty_expression::ExpressionKind::Call(callee, arguments) => {
            substitute_named_constants_in_expression(callee, constants);
            for argument in arguments {
                substitute_named_constants_in_expression(argument, constants);
            }
        }
        ty_expression::ExpressionKind::TypeConstruction(_, generic_arguments, fields) => {
            for argument in generic_arguments {
                if let ty_types::GenericArgument::Const(constant) = argument {
                    substitute_named_constant(constant, constants);
                }
            }
            for (_, expression) in fields {
                substitute_named_constants_in_expression(expression, constants);
            }
        }
        ty_expression::ExpressionKind::BinaryOperation(_, left, right)
        | ty_expression::ExpressionKind::Replace(left, right) => {
            substitute_named_constants_in_expression(left, constants);
            substitute_named_constants_in_expression(right, constants);
        }
        ty_expression::ExpressionKind::UnaryOperation(_, operand)
        | ty_expression::ExpressionKind::Dereference(operand)
        | ty_expression::ExpressionKind::Convert(operand, _) => {
            substitute_named_constants_in_expression(operand, constants);
        }
        ty_expression::ExpressionKind::Field(object, _) => {
            substitute_named_constants_in_expression(object, constants);
        }
        ty_expression::ExpressionKind::Index(object, index) => {
            substitute_named_constants_in_expression(object, constants);
            substitute_named_constants_in_expression(index, constants);
        }
        ty_expression::ExpressionKind::Block(block) => {
            substitute_named_constants_in_block(block, constants);
        }
        ty_expression::ExpressionKind::If {
            condition,
            then_branch,
            else_branch,
        } => {
            substitute_named_constants_in_expression(condition, constants);
            substitute_named_constants_in_block(then_branch, constants);
            if let Some(else_branch) = else_branch {
                substitute_named_constants_in_block(else_branch, constants);
            }
        }
        ty_expression::ExpressionKind::ArrayLiteral(elements)
        | ty_expression::ExpressionKind::TupleLiteral(elements)
        | ty_expression::ExpressionKind::Print(elements) => {
            for element in elements {
                substitute_named_constants_in_expression(element, constants);
            }
        }
        ty_expression::ExpressionKind::Match { value, arms } => {
            substitute_named_constants_in_expression(value, constants);
            for arm in arms {
                substitute_named_constants_in_block(&mut arm.body, constants);
            }
        }
        _ => {}
    }
}

fn substitute_named_constants_in_type(
    source: &mut ty_types::Type,
    constants: &HashMap<String, u128>,
) {
    match source {
        ty_types::Type::Named(_, arguments) => {
            for argument in arguments {
                if let ty_types::GenericArgument::Const(constant) = argument {
                    substitute_named_constant(constant, constants);
                }
            }
        }
        ty_types::Type::Array(element, size) | ty_types::Type::Vector(element, size) => {
            substitute_named_constant(size, constants);
            substitute_named_constants_in_type(element, constants);
        }
        ty_types::Type::Pointer(_, _, inner)
        | ty_types::Type::Reference(_, _, inner)
        | ty_types::Type::Slice(_, inner) => {
            substitute_named_constants_in_type(inner, constants);
        }
        ty_types::Type::Tuple(fields) | ty_types::Type::Enum(fields) => {
            for field in fields {
                substitute_named_constants_in_type(field, constants);
            }
        }
        _ => {}
    }
}

fn substitute_named_constant(
    constant: &mut ty_types::ConstExpression,
    constants: &HashMap<String, u128>,
) {
    use ty_types::ConstExpression as C;
    match constant {
        C::Named(name) => {
            if let Some(value) = constants.get(name) {
                *constant = C::Integer(*value);
            }
        }
        C::Integer(_) => {}
        C::Add(l, r)
        | C::Subtract(l, r)
        | C::Multiply(l, r)
        | C::Divide(l, r)
        | C::Modulo(l, r)
        | C::BitwiseAnd(l, r)
        | C::BitwiseOr(l, r)
        | C::BitwiseXor(l, r)
        | C::ShiftLeft(l, r)
        | C::ShiftRight(l, r) => {
            substitute_named_constant(l, constants);
            substitute_named_constant(r, constants);
        }
        C::Call(_, arguments) => {
            for argument in arguments {
                substitute_named_constant(argument, constants);
            }
        }
        C::Field(inner, _) => substitute_named_constant(inner, constants),
        C::Construction(fields) => {
            for (_, value) in fields {
                substitute_named_constant(value, constants);
            }
        }
    }
}
