mod emit;
mod resolve;
mod rewrite;
mod substitute;
mod util;

use std::collections::{HashMap, HashSet};

use crate::const_eval::ConstEvaluator;
use crate::resolve::ResolveContext;
use ty_ree::declaration as ty;
use ty_ree::expression as ty_expression;
use ty_ree::types as ty_types;

use emit::{emit_impl_instances, emit_type_instances};
use resolve::resolve_const_calls_in_module;
use rewrite::{compute_nongeneric_field_renames, rewrite_block, rewrite_type};
use substitute::{substitute_block, substitute_named_constants_in_declaration, substitute_type};
use util::{
    Bindings, bindings_from_explicit_arguments, const_contains_unresolved,
    contains_unresolved_generic, mangle_name, unify,
};

pub fn monomorphize(module: ty::Module) -> Result<ty::Module, String> {
    let mut context = Context::new();

    let mut newtypes_for_const: HashMap<String, ty_types::Type> = HashMap::new();
    for declaration in &module {
        if let ty::Declaration::Type(newtype) = declaration {
            newtypes_for_const.insert(newtype.name.clone(), newtype.inner_type.clone());
        }
    }
    for declaration in &module {
        if let ty::Declaration::Impl(impl_block) = declaration
            && impl_block.const_fn
            && let ty_types::Type::Named(name, _) = &impl_block.target_type
        {
            context
                .const_evaluator
                .register_const_impl(name, &newtypes_for_const, impl_block);
        }
    }

    let mut named_constants: HashMap<String, u128> = HashMap::new();
    for declaration in &module {
        if let ty::Declaration::Constant(constant) = declaration {
            let value = match &constant.value.kind {
                ty_expression::ExpressionKind::Literal(ty_expression::Literal::Integer(value)) => {
                    Some(*value)
                }
                ty_expression::ExpressionKind::Call(callee, arguments) => {
                    if let ty_expression::ExpressionKind::Variable(name) = &callee.kind
                        && context.const_evaluator.has_function(name)
                    {
                        let evaluated_arguments: Result<Vec<_>, _> = arguments
                            .iter()
                            .map(|argument| match &argument.kind {
                                ty_expression::ExpressionKind::Literal(
                                    ty_expression::Literal::Integer(value),
                                ) => Ok(crate::const_eval::ConstValue::Integer(*value as i128)),
                                ty_expression::ExpressionKind::Variable(var_name) => {
                                    if let Some(value) = named_constants.get(var_name) {
                                        Ok(crate::const_eval::ConstValue::Integer(*value as i128))
                                    } else {
                                        Err(())
                                    }
                                }
                                _ => Err(()),
                            })
                            .collect();
                        if let Ok(arguments) = evaluated_arguments {
                            context
                                .const_evaluator
                                .call(name, arguments)
                                .ok()
                                .and_then(|result| result.as_integer().ok())
                                .map(|value| value as u128)
                        } else {
                            None
                        }
                    } else {
                        None
                    }
                }
                _ => None,
            };
            if let Some(value) = value {
                named_constants.insert(constant.name.clone(), value);
                context
                    .const_evaluator
                    .register_constant(&constant.name, value);
            }
        }
    }

    let module: ty::Module = if named_constants.is_empty() {
        module
    } else {
        module
            .into_iter()
            .map(|mut declaration| {
                if let ty::Declaration::Constant(ref mut constant) = declaration
                    && let Some(value) = named_constants.get(&constant.name)
                {
                    constant.value = ty_expression::Expression::new(
                        ty_expression::ExpressionKind::Literal(ty_expression::Literal::Integer(
                            *value,
                        )),
                        constant.value.resolved_type.clone(),
                    );
                }
                substitute_named_constants_in_declaration(declaration, &named_constants)
            })
            .collect()
    };

    for declaration in &module {
        if let ty::Declaration::Type(newtype) = declaration {
            context
                .type_definitions
                .insert(newtype.name.clone(), newtype.inner_type.clone());
        }
    }

    for declaration in &module {
        match declaration {
            ty::Declaration::Type(newtype) if !newtype.generic_parameters.is_empty() => {
                for parameter in &newtype.generic_parameters {
                    context
                        .generic_parameter_names
                        .insert(parameter.name.clone());
                }
                context
                    .generic_types
                    .insert(newtype.name.clone(), newtype.clone());
            }
            ty::Declaration::Impl(impl_block) if impl_block.const_fn => {
                if let ty_types::Type::Named(name, _) = &impl_block.target_type {
                    context.const_evaluator.register_const_impl(
                        name,
                        &newtypes_for_const,
                        impl_block,
                    );
                }
                if !impl_block.generic_parameters.is_empty() {
                    for parameter in &impl_block.generic_parameters {
                        context
                            .generic_parameter_names
                            .insert(parameter.name.clone());
                    }
                    if let ty_types::Type::Named(name, _) = &impl_block.target_type {
                        context
                            .generic_impls
                            .insert(name.clone(), impl_block.clone());
                    }
                }
            }
            ty::Declaration::Impl(impl_block) if !impl_block.generic_parameters.is_empty() => {
                for parameter in &impl_block.generic_parameters {
                    context
                        .generic_parameter_names
                        .insert(parameter.name.clone());
                }
                if let ty_types::Type::Named(name, _) = &impl_block.target_type {
                    context
                        .generic_impls
                        .insert(name.clone(), impl_block.clone());
                }
            }
            _ => {}
        }
    }

    if context.generic_types.is_empty() && context.generic_impls.is_empty() {
        return Ok(module);
    }

    for declaration in &module {
        match declaration {
            ty::Declaration::Impl(impl_block) if impl_block.generic_parameters.is_empty() => {
                context.scan_type(&impl_block.target_type);
                context.scan_type(&impl_block.callable.return_type);
                if let Some(parameters) = &impl_block.callable.parameters {
                    for parameter in parameters {
                        context.scan_type(&parameter.parameter_type);
                    }
                }
                context.scan_block(&impl_block.callable.body);
            }
            ty::Declaration::Type(newtype) if newtype.generic_parameters.is_empty() => {
                context.scan_type(&newtype.inner_type);
            }
            _ => {}
        }
    }

    context.process_pending_instances();

    context.compute_field_renames();

    for declaration in &module {
        if let ty::Declaration::Type(newtype) = declaration
            && newtype.generic_parameters.is_empty()
        {
            let mut rewritten = newtype.inner_type.clone();
            rewrite_type(&mut rewritten, &context);
            compute_nongeneric_field_renames(
                &newtype.name,
                &newtype.inner_type,
                &rewritten,
                &mut context,
            );
        }
    }

    let mut output = Vec::new();

    for declaration in module {
        match declaration {
            ty::Declaration::Type(newtype) if !newtype.generic_parameters.is_empty() => {
                emit_type_instances(&newtype, &context, &mut output);
            }
            ty::Declaration::Impl(impl_block) if !impl_block.generic_parameters.is_empty() => {
                emit_impl_instances(&impl_block, &context, &mut output)?;
            }
            ty::Declaration::Impl(mut impl_block) => {
                if impl_block.const_fn {
                    impl_block.const_fn = false;
                }
                rewrite_type(&mut impl_block.target_type, &context);
                rewrite_type(&mut impl_block.callable.return_type, &context);
                if let Some(parameters) = &mut impl_block.callable.parameters {
                    for parameter in parameters {
                        rewrite_type(&mut parameter.parameter_type, &context);
                    }
                }
                rewrite_block(&mut impl_block.callable.body, &context);
                output.push(ty::Declaration::Impl(impl_block));
            }
            ty::Declaration::Type(mut newtype) => {
                rewrite_type(&mut newtype.inner_type, &context);
                output.push(ty::Declaration::Type(newtype));
            }
            other => output.push(other),
        }
    }

    if context.const_evaluator.has_any() {
        resolve_const_calls_in_module(&mut output, &mut context.const_evaluator)?;
    }

    Ok(output)
}

struct Context {
    generic_types: HashMap<String, ty::Newtype>,
    generic_impls: HashMap<String, ty::ImplBlock>,
    generic_parameter_names: HashSet<String>,
    type_instances: HashMap<String, Vec<Bindings>>,
    pending_instances: Vec<(String, Bindings)>,
    field_renames: HashMap<String, HashMap<String, String>>,
    const_evaluator: ConstEvaluator,
    type_definitions: HashMap<String, ty_types::Type>,
}

impl Context {
    fn new() -> Self {
        Self {
            generic_types: HashMap::new(),
            generic_impls: HashMap::new(),
            generic_parameter_names: HashSet::new(),
            type_instances: HashMap::new(),
            pending_instances: Vec::new(),
            field_renames: HashMap::new(),
            const_evaluator: ConstEvaluator::new(),
            type_definitions: HashMap::new(),
        }
    }

    fn record_type_instance(&mut self, name: &str, bindings: Bindings) {
        let instances = self.type_instances.entry(name.into()).or_default();
        if instances.contains(&bindings) {
            return;
        }
        instances.push(bindings.clone());
        self.pending_instances.push((name.into(), bindings));
    }

    fn process_pending_instances(&mut self) {
        while let Some((name, bindings)) = self.pending_instances.pop() {
            if let Some(generic_type) = self.generic_types.get(&name).cloned() {
                let substituted_inner = substitute_type(&generic_type.inner_type, &bindings);
                self.scan_type(&substituted_inner);
            }

            if let Some(generic_impl) = self.generic_impls.get(&name).cloned() {
                let substituted_return =
                    substitute_type(&generic_impl.callable.return_type, &bindings);
                self.scan_type(&substituted_return);
                let substituted_body = substitute_block(&generic_impl.callable.body, &bindings);
                self.scan_block(&substituted_body);
            }
        }
    }

    fn compute_field_renames(&mut self) {
        for (type_name, instances) in &self.type_instances {
            let Some(generic_type) = self.generic_types.get(type_name) else {
                continue;
            };
            for bindings in instances {
                let mangled = mangle_name(type_name, bindings, &generic_type.generic_parameters);
                let substituted_inner = substitute_type(&generic_type.inner_type, bindings);
                let mut renames = HashMap::new();
                let fields_to_check: Vec<_> = match &substituted_inner {
                    ty_types::Type::Tuple(fields) => fields.iter().collect(),
                    ty_types::Type::Named(..) => vec![&substituted_inner],
                    _ => vec![],
                };
                let original_fields: Vec<_> = match &generic_type.inner_type {
                    ty_types::Type::Tuple(fields) => fields.iter().collect(),
                    ty_types::Type::Named(..) => vec![&generic_type.inner_type],
                    _ => vec![],
                };
                for (original, substituted) in original_fields.iter().zip(&fields_to_check) {
                    if let ty_types::Type::Named(field_name, generic_arguments) = substituted {
                        if !generic_arguments.is_empty()
                            && let Some(field_type_definition) = self.generic_types.get(field_name)
                            && let Some(field_bindings) = bindings_from_explicit_arguments(
                                &field_type_definition.generic_parameters,
                                generic_arguments,
                            )
                        {
                            let mangled_field = mangle_name(
                                field_name,
                                &field_bindings,
                                &field_type_definition.generic_parameters,
                            );
                            renames.insert(field_name.clone(), mangled_field);
                        }
                        if let ty_types::Type::Named(original_name, _) = original
                            && *original_name != *field_name
                            && self.generic_parameter_names.contains(original_name)
                        {
                            renames.insert(original_name.clone(), field_name.clone());
                        }
                    }
                }
                if !renames.is_empty() {
                    self.field_renames.insert(mangled, renames);
                }
            }
        }
    }

    fn scan_block(&mut self, block: &ty_expression::Block) {
        let mut stack = Vec::new();
        self.scan_block_to_stack(block, &mut stack);
        while let Some(current) = stack.pop() {
            self.scan_type(&current.resolved_type);
            self.collect_expression_children(&current.kind, &mut stack);
        }
    }

    fn collect_expression_children<'a>(
        &mut self,
        kind: &'a ty_expression::ExpressionKind,
        stack: &mut Vec<&'a ty_expression::Expression>,
    ) {
        use ty_expression::ExpressionKind as E;
        match kind {
            E::Call(callee, arguments) => {
                self.try_infer_function_call(callee, arguments);
                stack.push(callee);
                for argument in arguments {
                    stack.push(argument);
                }
            }
            E::Specialize(inner, generic_arguments) => {
                if let E::Variable(name) = &inner.kind {
                    let lookup_name = name.strip_prefix("impl_").unwrap_or(name);
                    if let Some(generic_type) = self.generic_types.get(lookup_name).cloned()
                        && let Some(bindings) = bindings_from_explicit_arguments(
                            &generic_type.generic_parameters,
                            generic_arguments,
                        )
                    {
                        self.record_type_instance(lookup_name, bindings);
                    }
                }
                stack.push(inner);
            }
            E::BinaryOperation(operator, left, right) => {
                self.try_instantiate_operator(*operator, &left.resolved_type, &right.resolved_type);
                stack.push(left);
                stack.push(right);
            }
            E::UnaryOperation(_, operand)
            | E::Field(operand, _)
            | E::Dereference(operand)
            | E::Eval(operand)
            | E::ArrayRepeat(operand, _) => stack.push(operand),
            E::Index(object, index) => {
                stack.push(object);
                stack.push(index);
            }
            E::Convert(inner, target) => {
                stack.push(inner);
                self.scan_type(target);
            }
            E::SizeOf(target) => self.scan_type(target),
            E::TypeConstruction(name, generic_arguments, fields) => {
                if !generic_arguments.is_empty()
                    && let Some(newtype) = self.generic_types.get(name).cloned()
                    && let Some(bindings) = bindings_from_explicit_arguments(
                        &newtype.generic_parameters,
                        generic_arguments,
                    )
                {
                    self.record_type_instance(name, bindings);
                }
                for (_, expression) in fields {
                    stack.push(expression);
                }
            }
            E::ArrayLiteral(elements) | E::TupleLiteral(elements) | E::Print(elements) => {
                for element in elements {
                    stack.push(element);
                }
            }
            E::Block(block) => self.scan_block_to_stack(block, stack),
            E::If {
                condition,
                then_branch,
                else_branch,
            } => {
                stack.push(condition);
                self.scan_block_to_stack(then_branch, stack);
                if let Some(else_branch) = else_branch {
                    self.scan_block_to_stack(else_branch, stack);
                }
            }
            E::Match { value, arms } => {
                stack.push(value);
                for arm in arms {
                    self.scan_block_to_stack(&arm.body, stack);
                }
            }
            E::Replace(target, value) => {
                stack.push(target);
                stack.push(value);
            }
            E::Literal(_) | E::Variable(_) => {}
        }
    }

    fn scan_block_to_stack<'a>(
        &mut self,
        block: &'a ty_expression::Block,
        stack: &mut Vec<&'a ty_expression::Expression>,
    ) {
        for statement in &block.statements {
            self.scan_statement_to_stack(statement, stack);
        }
        if let Some(result) = &block.result {
            stack.push(result);
        }
    }

    fn scan_statement_to_stack<'a>(
        &mut self,
        statement: &'a ty_expression::Statement,
        stack: &mut Vec<&'a ty_expression::Expression>,
    ) {
        use ty_expression::Statement as S;
        match statement {
            S::Expression(expression) => stack.push(expression),
            S::Let {
                value,
                declared_type,
                ..
            } => {
                if let Some(declared_type) = declared_type {
                    self.scan_type(declared_type);
                }
                stack.push(value);
            }
            S::Assign(target, value) => {
                stack.push(target);
                stack.push(value);
            }
            S::Return(expression) => {
                if let Some(expression) = expression {
                    stack.push(expression);
                }
            }
            S::Label {
                parameters,
                initial_arguments,
                ..
            } => {
                for parameter in parameters {
                    self.scan_type(&parameter.parameter_type);
                }
                for argument in initial_arguments {
                    stack.push(argument);
                }
            }
            S::Jump { arguments, .. } => {
                for argument in arguments {
                    stack.push(argument);
                }
            }
            S::While { condition, body } => {
                stack.push(condition);
                self.scan_block_to_stack(body, stack);
            }
            S::Defer(inner) => self.scan_statement_to_stack(inner, stack),
        }
    }

    fn scan_type(&mut self, scanned_type: &ty_types::Type) {
        let mut type_stack = vec![scanned_type.clone()];
        while let Some(current) = type_stack.pop() {
            if let ty_types::Type::Named(ref name, ref generic_arguments) = current
                && !generic_arguments.is_empty()
                && let Some(newtype) = self.generic_types.get(name).cloned()
                && let Some(bindings) =
                    bindings_from_explicit_arguments(&newtype.generic_parameters, generic_arguments)
            {
                let has_unresolved =
                    bindings.types.values().any(|value| {
                        contains_unresolved_generic(value, &self.generic_parameter_names)
                    }) || bindings.constants.values().any(|value| {
                        const_contains_unresolved(value, &self.generic_parameter_names)
                    });
                if !has_unresolved {
                    self.record_type_instance(name, bindings);
                }
            }
            match &current {
                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) => type_stack.push((**inner).clone()),
                ty_types::Type::Tuple(fields) | ty_types::Type::Enum(fields) => {
                    for field in fields {
                        type_stack.push(field.clone());
                    }
                }
                ty_types::Type::Function(signature) => {
                    for parameter in &signature.parameters {
                        type_stack.push(parameter.clone());
                    }
                    type_stack.push((*signature.return_type).clone());
                }
                _ => {}
            }
        }
    }

    fn try_infer_function_call(
        &mut self,
        callee: &ty_expression::Expression,
        arguments: &[ty_expression::Expression],
    ) {
        let name = match &callee.kind {
            ty_expression::ExpressionKind::Variable(name) => name.clone(),
            ty_expression::ExpressionKind::Specialize(inner, _) => {
                if let ty_expression::ExpressionKind::Variable(name) = &inner.kind {
                    name.clone()
                } else {
                    return;
                }
            }
            _ => return,
        };

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

        let generic_type = match self.generic_types.get(name) {
            Some(generic_type) => generic_type.clone(),
            None => return,
        };

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

        let mut bindings = if let ty_expression::ExpressionKind::Specialize(_, generic_arguments) =
            &callee.kind
        {
            bindings_from_explicit_arguments(&generic_type.generic_parameters, generic_arguments)
                .unwrap_or_else(Bindings::new)
        } else {
            Bindings::new()
        };

        let field_types: Vec<ty_types::Type> = match &generic_type.inner_type {
            ty_types::Type::Tuple(fields) => fields.clone(),
            other => vec![other.clone()],
        };

        for (field_type, argument) in field_types.iter().zip(arguments) {
            let _ = unify(
                field_type,
                &argument.resolved_type,
                &mut bindings,
                &generic_names,
            );
        }

        if generic_type
            .generic_parameters
            .iter()
            .all(|parameter| bindings.contains(&parameter.name))
        {
            self.record_type_instance(name, bindings);
        }
    }

    fn try_instantiate_operator(
        &mut self,
        operator: t_ree::operator::BinaryOperator,
        left_type: &ty_types::Type,
        right_type: &ty_types::Type,
    ) {
        let Some(wrapper_name) = ResolveContext::operator_wrapper_name(operator) else {
            return;
        };
        if !matches!(left_type, ty_types::Type::Named(..))
            && !matches!(right_type, ty_types::Type::Named(..))
        {
            return;
        }
        let Some(wrapper_type) = self.generic_types.get(wrapper_name).cloned() else {
            return;
        };
        if wrapper_type.generic_parameters.len() != 1 {
            return;
        }
        if contains_unresolved_generic(left_type, &self.generic_parameter_names) {
            return;
        }
        let param_name = &wrapper_type.generic_parameters[0].name;
        let mut wrapper_bindings = Bindings::new();
        wrapper_bindings
            .types
            .insert(param_name.clone(), left_type.clone());
        self.record_type_instance(wrapper_name, wrapper_bindings);
        let wrapper_concrete = ty_types::Type::Named(
            wrapper_name.into(),
            vec![ty_types::GenericArgument::Type(left_type.clone())],
        );
        for (name, generic_type) in self.generic_types.clone() {
            if generic_type.generic_parameters.len() != 1 {
                continue;
            }
            let param = &generic_type.generic_parameters[0].name;
            if let ty_types::Type::Tuple(fields) = &generic_type.inner_type
                && fields.len() == 2
            {
                let has_wrapper = fields.iter().any(|field| {
                    if let ty_types::Type::Named(field_name, args) = field
                        && args.len() == 1
                    {
                        field_name == wrapper_name
                    } else {
                        false
                    }
                });
                if has_wrapper {
                    let is_const_param = matches!(
                        generic_type.generic_parameters[0].kind,
                        ty_types::GenericParameterKind::Const(_)
                    );
                    if is_const_param {
                        self.try_instantiate_const_generic_operator(
                            &name,
                            &generic_type,
                            fields,
                            wrapper_name,
                            left_type,
                            right_type,
                        );
                    } else {
                        let mut bindings = Bindings::new();
                        bindings.types.insert(param.clone(), left_type.clone());
                        let substituted = substitute_type(&generic_type.inner_type, &bindings);
                        if let ty_types::Type::Tuple(sub_fields) = &substituted {
                            let matches = sub_fields.iter().any(|f| f == &wrapper_concrete)
                                && sub_fields.iter().any(|f| f == right_type);
                            if matches {
                                self.record_type_instance(&name, bindings);
                            }
                        }
                    }
                }
            }
        }
    }

    fn try_instantiate_const_generic_operator(
        &mut self,
        operator_type_name: &str,
        generic_type: &ty::Newtype,
        fields: &[ty_types::Type],
        wrapper_name: &str,
        left_type: &ty_types::Type,
        right_type: &ty_types::Type,
    ) {
        let generic_names: Vec<&str> = generic_type
            .generic_parameters
            .iter()
            .map(|p| p.name.as_str())
            .collect();

        let resolve = |operand: &ty_types::Type| -> ty_types::Type {
            if let ty_types::Type::Named(name, args) = operand
                && args.is_empty()
                && let Some(inner) = self.type_definitions.get(name)
            {
                inner.clone()
            } else {
                operand.clone()
            }
        };
        let resolved_left = resolve(left_type);
        let resolved_right = resolve(right_type);

        let mut bindings = Bindings::new();
        for field in fields {
            if let ty_types::Type::Named(field_name, field_args) = field
                && field_name == wrapper_name
                && field_args.len() == 1
                && let ty_types::GenericArgument::Type(inner_of_wrapper) = &field_args[0]
            {
                let _ = unify(
                    inner_of_wrapper,
                    &resolved_left,
                    &mut bindings,
                    &generic_names,
                );
            } else {
                let _ = unify(field, &resolved_right, &mut bindings, &generic_names);
            }
        }

        let all_bound = generic_type
            .generic_parameters
            .iter()
            .all(|p| bindings.contains(&p.name));
        if all_bound {
            self.record_type_instance(operator_type_name, bindings);

            if let Some(wrapper) = self.generic_types.get(wrapper_name).cloned() {
                let wrapper_param = &wrapper.generic_parameters[0].name;
                let mut wb = Bindings::new();
                wb.types
                    .insert(wrapper_param.clone(), resolved_left.clone());
                self.record_type_instance(wrapper_name, wb);
            }
        }
    }
}
