use crate::ast::expression::Expression;
use crate::ast::ir_builder::IrBuilder;
use crate::diagnostic::Diagnostic;
use crate::error_codes::BINARY_INCOMPATIBLE_TYPES;
use crate::ir::ir_assign::IrAssign;
use crate::ir::ir_binary_operation::{IrBinaryOperation, IrBinaryOperator};
use crate::ir::ir_expression::IrExpression;
use crate::ir::ir_operation::IrOperation;
use crate::ir::ir_statement::IrStatement;
use crate::ir::ir_variable::IrVariable;
use crate::source_range::SourceRange;
use crate::symbol::class_symbol::ClassSymbol;
use crate::symbol_table::SymbolTable;
use crate::type_info::TypeInfo;
use crate::types_table::TypesTable;
use crate::{diagnostics_result, handle_diagnostic, handle_diagnostics, maybe_return_diagnostics};
use std::cell::RefCell;
use std::rc::Rc;

pub enum BinaryOperation {
    Multiply,
    Divide,
    Modulo,
    Add,
    Subtract,
    LeftShift,
    RightShift,
    BitwiseAnd,
    BitwiseXor,
    BitwiseOr,
}

pub struct BinaryExpression {
    lhs: Box<Expression>,
    rhs: Box<Expression>,
    op: BinaryOperation,
    source_range: SourceRange,
    type_info: Option<TypeInfo>,
}

impl BinaryExpression {
    pub fn new(
        lhs: Expression,
        rhs: Expression,
        op: BinaryOperation,
        source_range: SourceRange,
    ) -> Self {
        Self {
            lhs: lhs.into(),
            rhs: rhs.into(),
            op,
            source_range,
            type_info: None,
        }
    }

    pub fn lhs(&self) -> &Expression {
        &self.lhs
    }

    pub fn rhs(&self) -> &Expression {
        &self.rhs
    }

    pub fn op(&self) -> &BinaryOperation {
        &self.op
    }

    pub fn source_range(&self) -> &SourceRange {
        &self.source_range
    }

    pub fn type_info(&self) -> &TypeInfo {
        self.type_info.as_ref().unwrap()
    }

    pub fn init_scopes(&mut self, symbol_table: &mut SymbolTable, container_scope: usize) {
        self.lhs.init_scopes(symbol_table, container_scope);
        self.rhs.init_scopes(symbol_table, container_scope);
    }

    pub fn check_field_initializer_names(
        &self,
        symbol_table: &SymbolTable,
        class_symbol: &ClassSymbol,
    ) -> Vec<Diagnostic> {
        [
            self.lhs
                .check_field_initializer_names(symbol_table, class_symbol),
            self.rhs
                .check_field_initializer_names(symbol_table, class_symbol),
        ]
        .into_iter()
        .flatten()
        .collect()
    }

    pub fn check_constructor_local_names(
        &self,
        symbol_table: &SymbolTable,
        class_symbol: &ClassSymbol,
    ) -> Vec<Diagnostic> {
        [
            self.lhs
                .check_constructor_local_names(symbol_table, class_symbol),
            self.rhs
                .check_constructor_local_names(symbol_table, class_symbol),
        ]
        .into_iter()
        .flatten()
        .collect()
    }

    pub fn check_method_local_names(
        &self,
        symbol_table: &SymbolTable,
        class_symbol: &ClassSymbol,
    ) -> Vec<Diagnostic> {
        [
            self.lhs
                .check_method_local_names(symbol_table, class_symbol),
            self.rhs
                .check_method_local_names(symbol_table, class_symbol),
        ]
        .into_iter()
        .flatten()
        .collect()
    }

    pub fn check_static_fn_local_names(&self, symbol_table: &SymbolTable) -> Vec<Diagnostic> {
        [
            self.lhs.check_static_fn_local_names(symbol_table),
            self.rhs.check_static_fn_local_names(symbol_table),
        ]
        .into_iter()
        .flatten()
        .collect()
    }

    fn check_op(
        &mut self,
        symbol_table: &SymbolTable,
        types_table: &TypesTable,
        check: impl Fn(&TypeInfo, &TypeInfo) -> bool,
        op_result: impl Fn(&TypeInfo, &TypeInfo) -> TypeInfo,
        lazy_diagnostic_message: impl Fn(&TypeInfo, &TypeInfo) -> String,
    ) -> Result<(), Diagnostic> {
        let lhs_type_info = self.lhs.type_info(symbol_table, types_table);
        let rhs_type_info = self.rhs.type_info(symbol_table, types_table);

        if check(lhs_type_info, rhs_type_info) {
            self.type_info = Some(op_result(lhs_type_info, rhs_type_info));
            Ok(())
        } else {
            let diagnostic = Diagnostic::new(
                &lazy_diagnostic_message(lhs_type_info, rhs_type_info),
                self.source_range.start(),
                self.source_range.end(),
            )
            .with_primary_label_message("Incompatible types for addition.")
            .with_reporter(file!(), line!())
            .with_error_code(BINARY_INCOMPATIBLE_TYPES);
            Err(diagnostic)
        }
    }

    pub fn type_check(
        &mut self,
        symbol_table: &SymbolTable,
        types_table: &TypesTable,
    ) -> Result<(), Vec<Diagnostic>> {
        let mut diagnostics: Vec<Diagnostic> = vec![];

        handle_diagnostics!(self.lhs.type_check(symbol_table, types_table), diagnostics);
        handle_diagnostics!(self.rhs.type_check(symbol_table, types_table), diagnostics);

        maybe_return_diagnostics!(diagnostics);

        match &self.op {
            BinaryOperation::Multiply => {
                handle_diagnostic!(
                    self.check_op(
                        symbol_table,
                        types_table,
                        |lhs, rhs| lhs.can_multiply(rhs),
                        |lhs, rhs| lhs.multiply_result(rhs),
                        |lhs, rhs| format!(
                            "Incompatible types: cannot multiply {} by {}",
                            lhs, rhs
                        )
                    ),
                    diagnostics
                );
            }
            BinaryOperation::Divide => {
                handle_diagnostic!(
                    self.check_op(
                        symbol_table,
                        types_table,
                        |lhs, rhs| lhs.can_divide(rhs),
                        |lhs, rhs| lhs.divide_result(rhs),
                        |lhs, rhs| format!("Incompatible types: cannot divide {} by {}", lhs, rhs)
                    ),
                    diagnostics
                );
            }
            BinaryOperation::Modulo => {
                handle_diagnostic!(
                    self.check_op(
                        symbol_table,
                        types_table,
                        |lhs, rhs| lhs.can_modulo(rhs),
                        |lhs, rhs| lhs.modulo_result(rhs),
                        |lhs, rhs| format!("Incompatible types: cannot modulo {} by {}", lhs, rhs)
                    ),
                    diagnostics
                );
            }
            BinaryOperation::Add => {
                handle_diagnostic!(
                    self.check_op(
                        symbol_table,
                        types_table,
                        |lhs, rhs| lhs.can_add(rhs),
                        |lhs, rhs| lhs.add_result(&rhs),
                        |lhs, rhs| format!("Incompatible types: cannot add {} to {}.", rhs, lhs)
                    ),
                    diagnostics
                );
            }
            BinaryOperation::Subtract => {
                handle_diagnostic!(
                    self.check_op(
                        symbol_table,
                        types_table,
                        |lhs, rhs| lhs.can_subtract(rhs),
                        |lhs, rhs| lhs.subtract_result(rhs),
                        |lhs, rhs| format!(
                            "Incompatible types: cannot subtract {} from {}.",
                            rhs, lhs
                        )
                    ),
                    diagnostics
                )
            }
            BinaryOperation::LeftShift => {
                handle_diagnostic!(
                    self.check_op(
                        symbol_table,
                        types_table,
                        |lhs, rhs| lhs.can_left_shift(rhs),
                        |lhs, rhs| lhs.left_shift_result(rhs),
                        |lhs, rhs| format!(
                            "Incompatible types: cannot left shift {} by {}",
                            lhs, rhs
                        )
                    ),
                    diagnostics
                );
            }
            BinaryOperation::RightShift => {
                handle_diagnostic!(
                    self.check_op(
                        symbol_table,
                        types_table,
                        |lhs, rhs| lhs.can_right_shift(rhs),
                        |lhs, rhs| lhs.right_shift_result(rhs),
                        |lhs, rhs| format!(
                            "Incompatible types: cannot right shift {} by {}",
                            lhs, rhs
                        )
                    ),
                    diagnostics
                );
            }
            BinaryOperation::BitwiseAnd => {
                handle_diagnostic!(
                    self.check_op(
                        symbol_table,
                        types_table,
                        |lhs, rhs| lhs.can_bitwise_and(rhs),
                        |lhs, rhs| lhs.bitwise_and_result(rhs),
                        |lhs, rhs| format!(
                            "Incompatible types: cannot bitwise and {} by {}",
                            lhs, rhs
                        )
                    ),
                    diagnostics
                );
            }
            BinaryOperation::BitwiseXor => handle_diagnostic!(
                self.check_op(
                    symbol_table,
                    types_table,
                    |lhs, rhs| lhs.can_bitwise_xor(rhs),
                    |lhs, rhs| lhs.bitwise_xor_result(rhs),
                    |lhs, rhs| format!("Incompatible types: cannot bitwise xor {} by {}", lhs, rhs)
                ),
                diagnostics
            ),
            BinaryOperation::BitwiseOr => {
                handle_diagnostic!(
                    self.check_op(
                        symbol_table,
                        types_table,
                        |lhs, rhs| lhs.can_bitwise_or(rhs),
                        |lhs, rhs| lhs.bitwise_or_result(rhs),
                        |lhs, rhs| format!(
                            "Incompatible types: cannot bitwise or {} by {}",
                            lhs, rhs
                        )
                    ),
                    diagnostics
                );
            }
        }

        diagnostics_result!(diagnostics)
    }

    pub fn to_ir_operation(
        &self,
        builder: &mut IrBuilder,
        symbol_table: &SymbolTable,
        types_table: &TypesTable,
    ) -> IrOperation {
        let lhs = self
            .lhs
            .to_ir_expression(builder, symbol_table, types_table)
            .expect("Attempt to use a non-value expression in binary expression.");
        let rhs = self
            .rhs
            .to_ir_expression(builder, symbol_table, types_table)
            .expect("Attempt to use a non-value expression in binary expression.");
        let ir_binary_operation = match self.op {
            BinaryOperation::Multiply => {
                IrBinaryOperation::new(lhs, rhs, IrBinaryOperator::Multiply)
            }
            BinaryOperation::Divide => IrBinaryOperation::new(lhs, rhs, IrBinaryOperator::Divide),
            BinaryOperation::Modulo => IrBinaryOperation::new(lhs, rhs, IrBinaryOperator::Modulo),
            BinaryOperation::Add => IrBinaryOperation::new(lhs, rhs, IrBinaryOperator::Add),
            BinaryOperation::Subtract => {
                IrBinaryOperation::new(lhs, rhs, IrBinaryOperator::Subtract)
            }
            BinaryOperation::LeftShift => {
                IrBinaryOperation::new(lhs, rhs, IrBinaryOperator::LeftShift)
            }
            BinaryOperation::RightShift => {
                IrBinaryOperation::new(lhs, rhs, IrBinaryOperator::RightShift)
            }
            BinaryOperation::BitwiseAnd => {
                IrBinaryOperation::new(lhs, rhs, IrBinaryOperator::BitwiseAnd)
            }
            BinaryOperation::BitwiseXor => {
                IrBinaryOperation::new(lhs, rhs, IrBinaryOperator::BitwiseXor)
            }
            BinaryOperation::BitwiseOr => {
                IrBinaryOperation::new(lhs, rhs, IrBinaryOperator::BitwiseOr)
            }
        };
        IrOperation::Binary(ir_binary_operation)
    }

    pub fn to_ir_expression(
        &self,
        builder: &mut IrBuilder,
        symbol_table: &SymbolTable,
        types_table: &TypesTable,
    ) -> IrExpression {
        let ir_operation = self.to_ir_operation(builder, symbol_table, types_table);
        let t_var = IrVariable::new_vr(
            builder.new_t_var().into(),
            builder.current_block().id(),
            self.type_info(),
        );
        let as_rc = Rc::new(RefCell::new(t_var));
        let ir_assign = IrAssign::new(as_rc.clone(), ir_operation);
        builder
            .current_block_mut()
            .add_statement(IrStatement::Assign(ir_assign));
        IrExpression::Variable(as_rc)
    }
}