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

use ecow::EcoString;

use crate::cranelift::mir::{
    ast::{CompleteType, Expression, Function, FunctionParameter, IncompleteType, Module, Var},
    visit::Visit,
};

pub fn lower(modules: Vec<Module<IncompleteType>>) -> Vec<Module<CompleteType>> {
    let (generics, functions): (HashMap<_, _>, HashMap<_, _>) = modules
        .into_iter()
        .flat_map(|module| {
            module
                .functions
                .into_iter()
                .map(|function| ((module.name.clone(), function.name.clone()), function))
                .collect::<Vec<_>>()
        })
        .partition(|(_, function)| is_generic(function));

    let mut functions = functions
        .into_iter()
        .map(|((module, name), function)| {
            let function = Function::<CompleteType>::try_from(function)?;

            let instantiation = Instantiation {
                argument_types: function
                    .parameters
                    .iter()
                    .map(|parameter| parameter.type_.clone())
                    .collect(),
                return_type: function.return_type.clone(),
            };

            Ok(((module, name, instantiation), function))
        })
        .collect::<Result<
            HashMap<(EcoString, EcoString, Instantiation), Function<CompleteType>>,
            NonGenericError,
        >>()
        .expect("expected all functions to be non-generic");

    let mut monomorphizer = Monomorphizer::new(&generics);
    for (_, function) in functions.iter_mut() {
        monomorphizer.visit_function(function);
    }

    let mut queue: VecDeque<_> = monomorphizer.instantiations.into_iter().collect();

    while let Some((module, name, instantiation)) = queue.pop_front() {
        if let Some(function) = generics.get(&(module.clone(), name.clone())).cloned() {
            let mut function = FunctionMonomorphizer::default()
                .translate_function(function, instantiation.clone());

            _ = functions.insert((module, name, instantiation), function.clone());

            let mut monomorphizer = Monomorphizer::new(&generics);
            monomorphizer.visit_function(&mut function);

            for instantiation in monomorphizer.instantiations {
                if !functions.contains_key(&instantiation) {
                    queue.push_back(instantiation);
                }
            }
        }
    }

    let mut module_map: HashMap<EcoString, Vec<Function<CompleteType>>> = HashMap::new();
    for ((module, _, _), function) in functions {
        module_map
            .entry(module)
            .or_insert_with(Vec::new)
            .push(function);
    }

    module_map
        .into_iter()
        .map(|(name, functions)| Module { name, functions })
        .collect()
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct Instantiation {
    argument_types: Vec<CompleteType>,
    return_type: CompleteType,
}

#[derive(Default)]
struct FunctionMonomorphizer {
    generic_id_map: HashMap<u64, CompleteType>,
}

impl FunctionMonomorphizer {
    fn collect_type_ids(&mut self, incomplete: &IncompleteType, complete: &CompleteType) {
        match incomplete {
            IncompleteType::Generic { id } => {
                _ = self.generic_id_map.insert(*id, complete.clone());
            }

            IncompleteType::Func { arguments, returns } => {
                let CompleteType::Func {
                    arguments: complete_arguments,
                    returns: complete_returns,
                } = complete
                else {
                    unreachable!("types should not be diverged")
                };

                for (argument, complete) in arguments.iter().zip(complete_arguments.iter()) {
                    self.collect_type_ids(argument, complete);
                }

                self.collect_type_ids(returns, complete_returns);
            }
            IncompleteType::Tuple { elements } => {
                let CompleteType::Tuple { elements: complete } = complete else {
                    unreachable!("types should not be diverged")
                };

                for (element, complete) in elements.iter().zip(complete.iter()) {
                    self.collect_type_ids(element, complete);
                }
            }
            IncompleteType::Struct { elements } => {
                let CompleteType::Struct { elements: complete } = complete else {
                    unreachable!("types should not be diverged")
                };

                for (element, complete) in elements.iter().zip(complete.iter()) {
                    self.collect_type_ids(element, complete);
                }
            }
            IncompleteType::List(element_type) => {
                let CompleteType::List(complete) = complete else {
                    unreachable!("types should not be diverged")
                };

                self.collect_type_ids(element_type, complete);
            }

            IncompleteType::Int
            | IncompleteType::Float
            | IncompleteType::Bool
            | IncompleteType::String => {}
        }
    }

    fn translate_function(
        mut self,
        function: Function<IncompleteType>,
        instantiation: Instantiation,
    ) -> Function<CompleteType> {
        for (parameter, argument_type) in function
            .parameters
            .iter()
            .zip(instantiation.argument_types.iter())
        {
            self.collect_type_ids(&parameter.type_, argument_type);
        }

        self.collect_type_ids(&function.return_type, &instantiation.return_type);

        Function {
            name: function.name,
            return_type: instantiation.return_type,
            parameters: function
                .parameters
                .into_iter()
                .zip(instantiation.argument_types.into_iter())
                .map(|(parameter, argument_type)| FunctionParameter {
                    name: parameter.name,
                    type_: argument_type,
                })
                .collect(),
            body: self.translate_expression(function.body),
        }
    }

    fn translate_expression(
        &mut self,
        expression: Expression<IncompleteType>,
    ) -> Expression<CompleteType> {
        match expression {
            Expression::Block(expressions) => Expression::Block(
                expressions
                    .into_iter()
                    .map(|expression| self.translate_expression(expression))
                    .collect(),
            ),
            Expression::FunctionRef {
                module,
                name,
                arity,
                type_,
            } => Expression::FunctionRef {
                module,
                name,
                arity,
                type_: self.translate_type(type_),
            },
            Expression::Var(var) => Expression::Var(Var {
                name: var.name,
                type_: self.translate_type(var.type_),
            }),
            Expression::Int { value } => Expression::Int { value },
            Expression::Float { value } => Expression::Float { value },
            Expression::Bool { value } => Expression::Bool { value },
            Expression::String { value } => Expression::String { value },
            Expression::Equals { lhs, rhs } => Expression::Equals {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::NotEquals { lhs, rhs } => Expression::NotEquals {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::IntGt { lhs, rhs } => Expression::IntGt {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::IntGtEq { lhs, rhs } => Expression::IntGtEq {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::IntLt { lhs, rhs } => Expression::IntLt {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::IntLtEq { lhs, rhs } => Expression::IntLtEq {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::IntAdd { lhs, rhs } => Expression::IntAdd {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::IntSub { lhs, rhs } => Expression::IntSub {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::IntMul { lhs, rhs } => Expression::IntMul {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::IntDiv { lhs, rhs } => Expression::IntDiv {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::IntRem { lhs, rhs } => Expression::IntRem {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::FloatGt { lhs, rhs } => Expression::FloatGt {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::FloatGtEq { lhs, rhs } => Expression::FloatGtEq {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::FloatLt { lhs, rhs } => Expression::FloatLt {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::FloatLtEq { lhs, rhs } => Expression::FloatLtEq {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::FloatAdd { lhs, rhs } => Expression::FloatAdd {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::FloatSub { lhs, rhs } => Expression::FloatSub {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::FloatMul { lhs, rhs } => Expression::FloatMul {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::FloatDiv { lhs, rhs } => Expression::FloatDiv {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::StringConcat { lhs, rhs } => Expression::StringConcat {
                lhs: Box::new(self.translate_expression(*lhs)),
                rhs: Box::new(self.translate_expression(*rhs)),
            },
            Expression::List { items, tail, type_ } => Expression::List {
                items: items
                    .into_iter()
                    .map(|item| self.translate_expression(item))
                    .collect(),
                tail: tail.map(|t| Box::new(self.translate_expression(*t))),
                type_: self.translate_type(type_),
            },
            Expression::Tuple { items, type_ } => Expression::Tuple {
                items: items
                    .into_iter()
                    .map(|item| self.translate_expression(item))
                    .collect(),
                type_: self.translate_type(type_),
            },
            Expression::TupleAccess {
                value,
                index,
                type_,
            } => Expression::TupleAccess {
                value: Box::new(self.translate_expression(*value)),
                index,
                type_: self.translate_type(type_),
            },
            Expression::Struct { tag, items, type_ } => Expression::Struct {
                tag,
                items: items
                    .into_iter()
                    .map(|item| self.translate_expression(item))
                    .collect(),
                type_: self.translate_type(type_),
            },
            Expression::StructTag { value } => Expression::StructTag {
                value: Box::new(self.translate_expression(*value)),
            },
            Expression::StructAccess {
                value,
                index,
                type_,
            } => Expression::StructAccess {
                value: Box::new(self.translate_expression(*value)),
                index,
                type_: self.translate_type(type_),
            },
            Expression::Set { name, value } => Expression::Set {
                name: Var {
                    name: name.name,
                    type_: self.translate_type(name.type_),
                },
                value: Box::new(self.translate_expression(*value)),
            },
            Expression::If { cond, then, else_ } => Expression::If {
                cond: Box::new(self.translate_expression(*cond)),
                then: Box::new(self.translate_expression(*then)),
                else_: Box::new(self.translate_expression(*else_)),
            },
            Expression::Call {
                target,
                args,
                type_,
            } => Expression::Call {
                target: Box::new(self.translate_expression(*target)),
                args: args
                    .into_iter()
                    .map(|arg| self.translate_expression(arg))
                    .collect(),
                type_: self.translate_type(type_),
            },
            Expression::Panic { type_ } => Expression::Panic {
                type_: self.translate_type(type_),
            },
        }
    }

    fn translate_type(&mut self, type_: IncompleteType) -> CompleteType {
        match type_ {
            IncompleteType::Int => CompleteType::Int,
            IncompleteType::Float => CompleteType::Float,
            IncompleteType::Bool => CompleteType::Bool,
            IncompleteType::String => CompleteType::String,
            IncompleteType::Func { arguments, returns } => CompleteType::Func {
                arguments: arguments
                    .into_iter()
                    .map(|argument| self.translate_type(argument))
                    .collect(),
                returns: Box::new(self.translate_type(*returns)),
            },
            IncompleteType::Tuple { elements } => CompleteType::Tuple {
                elements: elements
                    .into_iter()
                    .map(|element| self.translate_type(element))
                    .collect(),
            },
            IncompleteType::Struct { elements } => CompleteType::Struct {
                elements: elements
                    .into_iter()
                    .map(|element| self.translate_type(element))
                    .collect(),
            },
            IncompleteType::List(element_type) => {
                CompleteType::List(Box::new(self.translate_type(*element_type)))
            }
            IncompleteType::Generic { id } => self
                .generic_id_map
                .get(&id)
                .cloned()
                .expect("expected generic id to have a matching type"),
        }
    }
}

struct Monomorphizer<'g> {
    generics: &'g HashMap<(EcoString, EcoString), Function<IncompleteType>>,
    instantiations: HashSet<(EcoString, EcoString, Instantiation)>,
}

impl<'g> Monomorphizer<'g> {
    fn new(generics: &'g HashMap<(EcoString, EcoString), Function<IncompleteType>>) -> Self {
        Self {
            generics,
            instantiations: Default::default(),
        }
    }
}

impl<'mir, 'g> Visit<'mir, CompleteType> for Monomorphizer<'g> {
    fn visit_expression_function_ref(
        &mut self,
        module: &mut EcoString,
        name: &mut EcoString,
        _arity: &mut usize,
        type_: &mut CompleteType,
    ) {
        if !self.generics.contains_key(&(module.clone(), name.clone())) {
            return;
        }

        let CompleteType::Func { arguments, returns } = type_ else {
            unreachable!("a function ref should always have a func type")
        };

        let instantiation = Instantiation {
            argument_types: arguments.clone(),
            return_type: *returns.clone(),
        };

        _ = self
            .instantiations
            .insert((module.clone(), name.clone(), instantiation));
    }
}

fn is_generic(function: &Function<IncompleteType>) -> bool {
    let has_generic_parameter = function
        .parameters
        .iter()
        .any(|parameter| matches!(parameter.type_, IncompleteType::Generic { .. }));

    has_generic_parameter || matches!(function.return_type, IncompleteType::Generic { .. })
}

#[derive(Debug, Clone, Copy)]
pub struct NonGenericError;

impl TryFrom<Function<IncompleteType>> for Function<CompleteType> {
    type Error = NonGenericError;

    fn try_from(value: Function<IncompleteType>) -> Result<Self, Self::Error> {
        Ok(Function {
            name: value.name,
            return_type: value.return_type.try_into()?,
            parameters: value
                .parameters
                .into_iter()
                .map(|parameter| {
                    Ok(FunctionParameter {
                        name: parameter.name,
                        type_: CompleteType::try_from(parameter.type_)?,
                    })
                })
                .collect::<Result<Vec<_>, _>>()?,
            body: Expression::<CompleteType>::try_from(value.body)?,
        })
    }
}

impl TryFrom<Expression<IncompleteType>> for Expression<CompleteType> {
    type Error = NonGenericError;

    fn try_from(value: Expression<IncompleteType>) -> Result<Self, Self::Error> {
        match value {
            Expression::Block(expressions) => Ok(Expression::Block(
                expressions
                    .into_iter()
                    .map(Expression::try_from)
                    .collect::<Result<Vec<_>, _>>()?,
            )),
            Expression::FunctionRef {
                module,
                name,
                arity,
                type_,
            } => Ok(Expression::FunctionRef {
                module,
                name,
                arity,
                type_: type_.try_into()?,
            }),
            Expression::Var(var) => Ok(Expression::Var(Var {
                name: var.name,
                type_: var.type_.try_into()?,
            })),
            Expression::Int { value } => Ok(Expression::Int { value }),
            Expression::Float { value } => Ok(Expression::Float { value }),
            Expression::Bool { value } => Ok(Expression::Bool { value }),
            Expression::String { value } => Ok(Expression::String { value }),
            Expression::Equals { lhs, rhs } => Ok(Expression::Equals {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::NotEquals { lhs, rhs } => Ok(Expression::NotEquals {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::IntGt { lhs, rhs } => Ok(Expression::IntGt {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::IntGtEq { lhs, rhs } => Ok(Expression::IntGtEq {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::IntLt { lhs, rhs } => Ok(Expression::IntLt {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::IntLtEq { lhs, rhs } => Ok(Expression::IntLtEq {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::IntAdd { lhs, rhs } => Ok(Expression::IntAdd {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::IntSub { lhs, rhs } => Ok(Expression::IntSub {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::IntMul { lhs, rhs } => Ok(Expression::IntMul {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::IntDiv { lhs, rhs } => Ok(Expression::IntDiv {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::IntRem { lhs, rhs } => Ok(Expression::IntRem {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::FloatGt { lhs, rhs } => Ok(Expression::FloatGt {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::FloatGtEq { lhs, rhs } => Ok(Expression::FloatGtEq {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::FloatLt { lhs, rhs } => Ok(Expression::FloatLt {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::FloatLtEq { lhs, rhs } => Ok(Expression::FloatLtEq {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::FloatAdd { lhs, rhs } => Ok(Expression::FloatAdd {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::FloatSub { lhs, rhs } => Ok(Expression::FloatSub {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::FloatMul { lhs, rhs } => Ok(Expression::FloatMul {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::FloatDiv { lhs, rhs } => Ok(Expression::FloatDiv {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::StringConcat { lhs, rhs } => Ok(Expression::StringConcat {
                lhs: Box::new(Expression::try_from(*lhs)?),
                rhs: Box::new(Expression::try_from(*rhs)?),
            }),
            Expression::List { items, tail, type_ } => Ok(Expression::List {
                items: items
                    .into_iter()
                    .map(Expression::try_from)
                    .collect::<Result<Vec<_>, _>>()?,
                tail: tail
                    .map(|t| Ok(Box::new(Expression::try_from(*t)?)))
                    .transpose()?,
                type_: type_.try_into()?,
            }),
            Expression::Tuple { items, type_ } => Ok(Expression::Tuple {
                items: items
                    .into_iter()
                    .map(Expression::try_from)
                    .collect::<Result<Vec<_>, _>>()?,
                type_: type_.try_into()?,
            }),
            Expression::TupleAccess {
                value,
                index,
                type_,
            } => Ok(Expression::TupleAccess {
                value: Box::new(Expression::try_from(*value)?),
                index,
                type_: type_.try_into()?,
            }),
            Expression::Struct { tag, items, type_ } => Ok(Expression::Struct {
                tag,
                items: items
                    .into_iter()
                    .map(Expression::try_from)
                    .collect::<Result<Vec<_>, _>>()?,
                type_: type_.try_into()?,
            }),
            Expression::StructTag { value } => Ok(Expression::StructTag {
                value: Box::new(Expression::try_from(*value)?),
            }),
            Expression::StructAccess {
                value,
                index,
                type_,
            } => Ok(Expression::StructAccess {
                value: Box::new(Expression::try_from(*value)?),
                index,
                type_: type_.try_into()?,
            }),
            Expression::Set { name, value } => Ok(Expression::Set {
                name: Var {
                    name: name.name,
                    type_: name.type_.try_into()?,
                },
                value: Box::new(Expression::try_from(*value)?),
            }),
            Expression::If { cond, then, else_ } => Ok(Expression::If {
                cond: Box::new(Expression::try_from(*cond)?),
                then: Box::new(Expression::try_from(*then)?),
                else_: Box::new(Expression::try_from(*else_)?),
            }),
            Expression::Call {
                target,
                args,
                type_,
            } => Ok(Expression::Call {
                target: Box::new(Expression::try_from(*target)?),
                args: args
                    .into_iter()
                    .map(Expression::try_from)
                    .collect::<Result<Vec<_>, _>>()?,
                type_: type_.try_into()?,
            }),
            Expression::Panic { type_ } => Ok(Expression::Panic {
                type_: type_.try_into()?,
            }),
        }
    }
}

impl TryFrom<IncompleteType> for CompleteType {
    type Error = NonGenericError;

    fn try_from(value: IncompleteType) -> Result<Self, Self::Error> {
        match value {
            IncompleteType::Int => Ok(Self::Int),
            IncompleteType::Float => Ok(Self::Float),
            IncompleteType::Bool => Ok(Self::Bool),
            IncompleteType::String => Ok(Self::String),
            IncompleteType::Func { arguments, returns } => Ok(Self::Func {
                arguments: arguments
                    .into_iter()
                    .map(Self::try_from)
                    .collect::<Result<Vec<_>, _>>()?,
                returns: Box::new(Self::try_from(*returns)?),
            }),
            IncompleteType::Tuple { elements } => Ok(Self::Tuple {
                elements: elements
                    .into_iter()
                    .map(Self::try_from)
                    .collect::<Result<Vec<_>, _>>()?,
            }),
            IncompleteType::Struct { elements } => Ok(Self::Struct {
                elements: elements
                    .into_iter()
                    .map(Self::try_from)
                    .collect::<Result<Vec<_>, _>>()?,
            }),
            IncompleteType::List(element_type) => {
                Ok(Self::List(Box::new(Self::try_from(*element_type)?)))
            }
            IncompleteType::Generic { id: _ } => Err(NonGenericError),
        }
    }
}