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

use t_ree::operator::{ArithmeticOperator, BinaryOperator, ComparisonOperator};
use ty_ree::declaration as ty;
use ty_ree::expression as ty_expression;
use ty_ree::types as ty_types;

struct OperatorMatch {
    impl_name: String,
    left_wrap: Option<String>,
    right_wrap: Option<String>,
}

impl OperatorMatch {
    fn wrap_operand(
        expression: ty_expression::Expression,
        wrapper: Option<&String>,
    ) -> ty_expression::Expression {
        match wrapper {
            None => expression,
            Some(wrapper_type) => {
                let wrapped_type = ty_types::Type::Named(wrapper_type.clone(), vec![]);
                ty_expression::Expression::new(
                    ty_expression::ExpressionKind::Convert(
                        Box::new(expression),
                        wrapped_type.clone(),
                    ),
                    wrapped_type,
                )
            }
        }
    }
}

pub struct ResolveContext {
    pub impl_targets: HashSet<String>,
    pub extern_targets: HashSet<String>,
    pub all_callables: HashSet<String>,
    pub newtypes: HashMap<String, ty_types::Type>,
}

impl ResolveContext {
    pub fn operator_wrapper_name(operator: BinaryOperator) -> Option<&'static str> {
        match operator {
            BinaryOperator::Arithmetic(ArithmeticOperator::Add) => Some("add"),
            BinaryOperator::Arithmetic(ArithmeticOperator::Subtract) => Some("sub"),
            BinaryOperator::Arithmetic(ArithmeticOperator::Multiply) => Some("mul"),
            BinaryOperator::Arithmetic(ArithmeticOperator::Divide) => Some("div"),
            BinaryOperator::Arithmetic(ArithmeticOperator::Remainder) => Some("rem"),
            BinaryOperator::Comparison(ComparisonOperator::Equal) => Some("equal"),
            BinaryOperator::Comparison(ComparisonOperator::NotEqual) => Some("not_equal"),
            BinaryOperator::Comparison(ComparisonOperator::Less) => Some("less"),
            BinaryOperator::Comparison(ComparisonOperator::LessEqual) => Some("less_equal"),
            BinaryOperator::Comparison(ComparisonOperator::Greater) => Some("greater"),
            BinaryOperator::Comparison(ComparisonOperator::GreaterEqual) => Some("greater_equal"),
            _ => None,
        }
    }

    fn find_operator(
        &self,
        operator: BinaryOperator,
        left_type: &ty_types::Type,
        right_type: &ty_types::Type,
    ) -> Option<OperatorMatch> {
        let wrapper_name = Self::operator_wrapper_name(operator)?;
        for impl_target in &self.impl_targets {
            let Some(inner) = self.newtypes.get(impl_target) else {
                continue;
            };
            let (field_left, field_right) = match inner {
                ty_types::Type::Tuple(fields) if fields.len() == 2 => (&fields[0], &fields[1]),
                ty_types::Type::Array(element, _) => {
                    let base = impl_target.split('$').next().unwrap_or(impl_target);
                    if !base.starts_with(wrapper_name) {
                        continue;
                    }
                    (element.as_ref(), element.as_ref())
                }
                _ => continue,
            };
            let left_match = self.match_operator_field(field_left, left_type, wrapper_name);
            let right_match = self.match_operator_field(field_right, right_type, wrapper_name);
            if let (Some(left_wrap), Some(right_wrap)) = (left_match, right_match) {
                return Some(OperatorMatch {
                    impl_name: format!("impl_{impl_target}"),
                    left_wrap,
                    right_wrap,
                });
            }
        }
        None
    }

    #[allow(clippy::option_option)]
    fn match_operator_field(
        &self,
        field_type: &ty_types::Type,
        operand_type: &ty_types::Type,
        wrapper_name: &str,
    ) -> Option<Option<String>> {
        if field_type == operand_type {
            return Some(None);
        }
        if self.resolves_to(operand_type, field_type) {
            return Some(None);
        }
        if let ty_types::Type::Named(name, args) = field_type
            && args.is_empty()
            && name.split('$').next().unwrap_or(name) == wrapper_name
            && let Some(inner) = self.newtypes.get(name)
            && (self.resolves_to(inner, operand_type) || self.resolves_to(operand_type, inner))
        {
            return Some(Some(name.clone()));
        }
        None
    }

    fn resolves_to(&self, expected: &ty_types::Type, actual: &ty_types::Type) -> bool {
        if expected == actual {
            return true;
        }
        if let ty_types::Type::Named(name, args) = expected
            && args.is_empty()
            && let Some(inner) = self.newtypes.get(name)
        {
            return self.resolves_to(inner, actual);
        }
        false
    }
}

pub fn resolve_callable_constructions(module: ty::Module, context: &ResolveContext) -> ty::Module {
    module
        .into_iter()
        .map(|declaration| resolve_callable_in_declaration(declaration, context))
        .collect()
}

fn has_named_operand(left: &ty_types::Type, right: &ty_types::Type) -> bool {
    fn is_named(resolved_type: &ty_types::Type) -> bool {
        matches!(resolved_type, ty_types::Type::Named(..))
    }
    is_named(left) || is_named(right)
}

fn resolve_callable_in_declaration(
    declaration: ty::Declaration,
    context: &ResolveContext,
) -> ty::Declaration {
    match declaration {
        ty::Declaration::Impl(mut impl_block) => {
            let impl_name = if let ty_types::Type::Named(name, _) = &impl_block.target_type {
                Some(name.as_str())
            } else {
                None
            };
            resolve_callable_in_block(&mut impl_block.callable.body, context, impl_name);
            ty::Declaration::Impl(impl_block)
        }
        other => other,
    }
}

fn resolve_callable_in_block(
    block: &mut ty_expression::Block,
    context: &ResolveContext,
    current_impl: Option<&str>,
) {
    for statement in &mut block.statements {
        resolve_callable_in_statement(statement, context, current_impl);
    }
    if let Some(result) = &mut block.result {
        resolve_callable_in_expression(result, context, current_impl);
    }
}

fn resolve_callable_in_statement(
    statement: &mut ty_expression::Statement,
    context: &ResolveContext,
    current_impl: Option<&str>,
) {
    use ty_expression::Statement;
    match statement {
        Statement::Expression(expression) | Statement::Return(Some(expression)) => {
            resolve_callable_in_expression(expression, context, current_impl);
        }
        Statement::Let { value, .. } | Statement::Assign(_, value) => {
            resolve_callable_in_expression(value, context, current_impl);
        }
        Statement::Label {
            initial_arguments, ..
        } => {
            for value in initial_arguments {
                resolve_callable_in_expression(value, context, current_impl);
            }
        }
        Statement::Jump { arguments, .. } => {
            for argument in arguments {
                resolve_callable_in_expression(argument, context, current_impl);
            }
        }
        Statement::While { condition, body } => {
            resolve_callable_in_expression(condition, context, current_impl);
            resolve_callable_in_block(body, context, current_impl);
        }
        Statement::Defer(inner) => {
            resolve_callable_in_statement(inner, context, current_impl);
        }
        Statement::Return(None) => {}
    }
}

fn resolve_callable_in_expression(
    expression: &mut ty_expression::Expression,
    context: &ResolveContext,
    current_impl: Option<&str>,
) {
    use ty_expression::ExpressionKind as E;

    let kind = std::mem::replace(
        &mut expression.kind,
        E::Literal(t_ree::expression::Literal::Bool(false)),
    );

    expression.kind = match kind {
        E::Eval(inner) => match inner.kind {
            E::TypeConstruction(name, generic_arguments, fields)
                if generic_arguments.is_empty() && context.all_callables.contains(&name) =>
            {
                let resolved_name = if context.extern_targets.contains(&name) {
                    name
                } else {
                    format!("impl_{name}")
                };
                let callee = ty_expression::Expression::new(
                    E::Variable(resolved_name),
                    ty_types::Type::Inferred,
                );
                let arguments = fields
                    .into_iter()
                    .map(|(_, expression)| expression)
                    .collect();
                E::Call(Box::new(callee), arguments)
            }
            other => E::Eval(Box::new(ty_expression::Expression::new(
                other,
                inner.resolved_type,
            ))),
        },
        E::Call(callee, arguments) if matches!(&callee.kind, E::Variable(name) if context.all_callables.contains(name)) =>
        {
            let E::Variable(name) = callee.kind else {
                unreachable!()
            };
            let resolved_name = if context.extern_targets.contains(&name) {
                name
            } else {
                format!("impl_{name}")
            };
            let callee = ty_expression::Expression::new(
                E::Variable(resolved_name),
                ty_types::Type::Inferred,
            );
            E::Call(Box::new(callee), arguments)
        }
        E::Call(callee, mut arguments) if matches!(&callee.kind, E::Field(_, name) if context.all_callables.contains(name)) =>
        {
            let E::Field(object, name) = callee.kind else {
                unreachable!()
            };
            arguments.insert(0, *object);
            let resolved_name = if context.impl_targets.contains(&name) {
                format!("impl_{name}")
            } else {
                name
            };
            let callee = ty_expression::Expression::new(
                E::Variable(resolved_name),
                ty_types::Type::Inferred,
            );
            E::Call(Box::new(callee), arguments)
        }
        E::BinaryOperation(operator, left, right)
            if has_named_operand(&left.resolved_type, &right.resolved_type) =>
        {
            match context.find_operator(operator, &left.resolved_type, &right.resolved_type) {
                Some(op_match)
                    if current_impl.is_none_or(|name| !op_match.impl_name.ends_with(name)) =>
                {
                    let callee = ty_expression::Expression::new(
                        E::Variable(op_match.impl_name.clone()),
                        ty_types::Type::Inferred,
                    );
                    let wrapped_left =
                        OperatorMatch::wrap_operand(*left, op_match.left_wrap.as_ref());
                    let wrapped_right =
                        OperatorMatch::wrap_operand(*right, op_match.right_wrap.as_ref());
                    E::Call(Box::new(callee), vec![wrapped_left, wrapped_right])
                }
                _ => E::BinaryOperation(operator, left, right),
            }
        }
        other => other,
    };

    match &mut expression.kind {
        E::Call(callee, arguments) => {
            resolve_callable_in_expression(callee, context, current_impl);
            for argument in arguments {
                resolve_callable_in_expression(argument, context, current_impl);
            }
        }
        E::BinaryOperation(_, left, right) | E::Replace(left, right) => {
            resolve_callable_in_expression(left, context, current_impl);
            resolve_callable_in_expression(right, context, current_impl);
        }
        E::UnaryOperation(_, operand)
        | E::Dereference(operand)
        | E::Eval(operand)
        | E::Convert(operand, _)
        | E::ArrayRepeat(operand, _) => {
            resolve_callable_in_expression(operand, context, current_impl);
        }
        E::Field(object, _) => {
            resolve_callable_in_expression(object, context, current_impl);
        }
        E::Index(object, index) => {
            resolve_callable_in_expression(object, context, current_impl);
            resolve_callable_in_expression(index, context, current_impl);
        }
        E::TypeConstruction(_, _, fields) => {
            for (_, expression) in fields {
                resolve_callable_in_expression(expression, context, current_impl);
            }
        }
        E::ArrayLiteral(elements) | E::TupleLiteral(elements) | E::Print(elements) => {
            for element in elements {
                resolve_callable_in_expression(element, context, current_impl);
            }
        }
        E::Block(block) => {
            resolve_callable_in_block(block, context, current_impl);
        }
        E::If {
            condition,
            then_branch,
            else_branch,
        } => {
            resolve_callable_in_expression(condition, context, current_impl);
            resolve_callable_in_block(then_branch, context, current_impl);
            if let Some(else_branch) = else_branch {
                resolve_callable_in_block(else_branch, context, current_impl);
            }
        }
        E::Match { value, arms } => {
            resolve_callable_in_expression(value, context, current_impl);
            for arm in arms {
                resolve_callable_in_block(&mut arm.body, context, current_impl);
            }
        }
        E::Literal(_) | E::Variable(_) | E::Specialize(_, _) | E::SizeOf(_) => {}
    }
}
