mod util;

use crate::ast::NodeId;
use crate::ast::assign_statement::AssignStatement;
use crate::ast::binary_expression::BinaryOperation;
use crate::ast::call::Call;
use crate::ast::compilation_unit::CompilationUnit;
use crate::ast::expression::Expression;
use crate::ast::expression_statement::ExpressionStatement;
use crate::ast::function::Function;
use crate::ast::let_statement::LetStatement;
use crate::ast::statement::Statement;
use crate::ir::ir_assign::IrAssign;
use crate::ir::ir_binary_operation::{IrBinaryOperation, IrBinaryOperator};
use crate::ir::ir_block::{IrBlock, IrBlockId};
use crate::ir::ir_call::IrCall;
use crate::ir::ir_expression::IrExpression;
use crate::ir::ir_function::IrFunction;
use crate::ir::ir_operation::IrOperation;
use crate::ir::ir_parameter::{IrParameter, IrParameterId};
use crate::ir::ir_return::IrReturn;
use crate::ir::ir_statement::IrStatement;
use crate::ir::ir_type_info::IrTypeInfo;
use crate::ir::ir_variable::{IrVariable, IrVariableId};
use crate::lowering::util::{return_type_info_to_ir_type_info, to_ir_type_info};
use crate::semantic_analysis::symbol::{Symbol, SymbolId};
use crate::semantic_analysis::type_info::{TypeInfo, TypeInfoId};
use std::collections::HashMap;
use std::ops::Neg;

pub struct LowerToIrResult {
    pub functions: Vec<IrFunction>,
}

struct LowerToIrContext<'a> {
    symbols: &'a [Symbol],
    nodes_to_symbols: &'a HashMap<NodeId, SymbolId>,
    type_infos: &'a [TypeInfo],
    symbols_to_type_infos: &'a HashMap<SymbolId, TypeInfoId>,
    nodes_to_type_infos: &'a HashMap<NodeId, TypeInfoId>,
    ir_functions: Vec<IrFunction>,
}

pub fn lower_to_ir(
    compilation_unit: &CompilationUnit,
    symbols: &[Symbol],
    nodes_to_symbols: &HashMap<NodeId, SymbolId>,
    type_infos: &[TypeInfo],
    symbols_to_type_infos: &HashMap<SymbolId, TypeInfoId>,
    nodes_to_type_infos: &HashMap<NodeId, TypeInfoId>,
) -> LowerToIrResult {
    let mut ctx = LowerToIrContext {
        symbols,
        nodes_to_symbols,
        type_infos,
        symbols_to_type_infos,
        nodes_to_type_infos,
        ir_functions: Vec::new(),
    };

    for function in compilation_unit.functions() {
        lower_to_ir_function(function, &mut ctx);
    }

    LowerToIrResult {
        functions: ctx.ir_functions,
    }
}

struct LowerToIrFunctionContext {
    blocks: Vec<IrBlock>,
    ir_variables: Vec<IrVariable>,
    symbols_to_variables: HashMap<SymbolId, IrVariableId>,
    ir_parameters: Vec<IrParameter>,
    symbols_to_parameters: HashMap<SymbolId, IrParameterId>,
    current_block_statements: Vec<IrStatement>,
    t_var_counter: usize,
}

impl LowerToIrFunctionContext {
    fn new() -> Self {
        Self {
            blocks: Vec::new(),
            ir_variables: Vec::new(),
            symbols_to_variables: HashMap::new(),
            ir_parameters: Vec::new(),
            symbols_to_parameters: HashMap::new(),
            current_block_statements: Vec::new(),
            t_var_counter: 0,
        }
    }

    fn current_block_id(&self) -> IrBlockId {
        self.blocks.len()
    }

    fn finish_block(&mut self) {
        let current_block_id = self.current_block_id();
        let ir_block = IrBlock::new(
            current_block_id,
            &format!("#{}", current_block_id),
            std::mem::take(&mut self.current_block_statements),
        );
        self.blocks.push(ir_block);
    }

    fn insert_ir_parameter(
        &mut self,
        ir_parameter: IrParameter,
        maybe_associated_symbol_id: Option<SymbolId>,
    ) -> IrParameterId {
        self.ir_parameters.push(ir_parameter);
        let ir_parameter_id = self.ir_parameters.len() - 1;
        if let Some(associated_symbol_id) = maybe_associated_symbol_id {
            self.symbols_to_parameters
                .insert(associated_symbol_id, ir_parameter_id);
        }
        ir_parameter_id
    }

    fn insert_ir_variable(
        &mut self,
        ir_variable: IrVariable,
        maybe_associated_symbol_id: Option<SymbolId>,
    ) -> IrVariableId {
        self.ir_variables.push(ir_variable);
        let ir_variable_id = self.ir_variables.len() - 1;
        if let Some(associated_symbol_id) = maybe_associated_symbol_id {
            self.symbols_to_variables
                .insert(associated_symbol_id, ir_variable_id);
        }
        ir_variable_id
    }

    fn make_t_var(&mut self, ir_type_info: IrTypeInfo) -> IrVariableId {
        let t_var_id = self.t_var_counter;
        self.t_var_counter += 1;
        let t_var_ir_variable = IrVariable::new(&format!("t_{}", t_var_id), ir_type_info);
        self.insert_ir_variable(t_var_ir_variable, None)
    }
}

fn lower_to_ir_function(function: &Function, ctx: &mut LowerToIrContext) {
    let mut fn_ctx = LowerToIrFunctionContext::new();

    lower_to_ir_parameters(function, ctx, &mut fn_ctx);

    let n_statements = function.statements().len();
    for (i, statement) in function.statements().iter().enumerate() {
        let can_return_value = i == n_statements - 1;
        lower_to_ir_statement(statement, ctx, &mut fn_ctx, can_return_value);
    }

    fn_ctx.finish_block();

    let function_symbol_id = ctx.nodes_to_symbols[&function.node_id()];
    let function_symbol = match &ctx.symbols[function_symbol_id] {
        Symbol::Function(function_symbol) => function_symbol,
        _ => panic!("Expected function symbol"),
    };

    let function_type_info_id = ctx.symbols_to_type_infos[&function_symbol_id];
    let function_type_info = match &ctx.type_infos[function_type_info_id] {
        TypeInfo::Function(function_type_info) => function_type_info,
        _ => panic!("Expected FunctionTypeInfo"),
    };
    let return_type_info = &ctx.type_infos[function_type_info.return_type_id()];

    let ir_function = IrFunction::new(
        function_symbol.fqn_owned(),
        fn_ctx.ir_parameters,
        fn_ctx.ir_variables,
        return_type_info_to_ir_type_info(return_type_info),
        fn_ctx.blocks,
    );

    ctx.ir_functions.push(ir_function);
}

fn lower_to_ir_parameters(
    function: &Function,
    ctx: &LowerToIrContext,
    fn_ctx: &mut LowerToIrFunctionContext,
) {
    let n_parameters = function.parameters().len() as isize;
    for (i, parameter) in function.parameters().iter().enumerate() {
        let parameter_type_info_id = ctx.nodes_to_type_infos[&parameter.node_id()];
        let parameter_type_info = &ctx.type_infos[parameter_type_info_id];
        let ir_parameter = IrParameter::new(
            parameter.declared_name(),
            to_ir_type_info(parameter_type_info),
            n_parameters.neg() + (i as isize),
        );
        fn_ctx.insert_ir_parameter(
            ir_parameter,
            Some(ctx.nodes_to_symbols[&parameter.node_id()]),
        );
    }
}

fn lower_to_ir_statement(
    statement: &Statement,
    ctx: &LowerToIrContext,
    fn_ctx: &mut LowerToIrFunctionContext,
    can_return_value: bool,
) {
    match statement {
        Statement::Let(let_statement) => {
            lower_to_ir_let_statement(let_statement, ctx, fn_ctx);
        }
        Statement::Expression(expression_statement) => {
            lower_to_ir_expression_statement(expression_statement, ctx, fn_ctx, can_return_value);
        }
        Statement::Assign(assign_statement) => {
            lower_to_ir_assign_statement(assign_statement, ctx, fn_ctx);
        }
    }
}

fn lower_binary_operator(binary_operation: &BinaryOperation) -> IrBinaryOperator {
    match binary_operation {
        BinaryOperation::Multiply => IrBinaryOperator::Multiply,
        BinaryOperation::Divide => IrBinaryOperator::Divide,
        BinaryOperation::Modulo => IrBinaryOperator::Modulo,
        BinaryOperation::Add => IrBinaryOperator::Add,
        BinaryOperation::Subtract => IrBinaryOperator::Subtract,
        BinaryOperation::LeftShift => IrBinaryOperator::LeftShift,
        BinaryOperation::RightShift => IrBinaryOperator::RightShift,
        BinaryOperation::BitwiseAnd => IrBinaryOperator::BitwiseAnd,
        BinaryOperation::BitwiseXor => IrBinaryOperator::BitwiseXor,
        BinaryOperation::BitwiseOr => IrBinaryOperator::BitwiseOr,
    }
}

fn lower_to_ir_let_statement(
    let_statement: &LetStatement,
    ctx: &LowerToIrContext,
    fn_ctx: &mut LowerToIrFunctionContext,
) {
    let symbol_id = ctx.nodes_to_symbols[&let_statement.node_id()];

    let type_info_id = ctx.nodes_to_type_infos[&let_statement.node_id()];
    let type_info = &ctx.type_infos[type_info_id];

    let ir_variable = IrVariable::new(let_statement.declared_name(), to_ir_type_info(type_info));
    let ir_variable_id = fn_ctx.insert_ir_variable(ir_variable, Some(symbol_id));

    let initializer_ir_operation =
        lower_expression_to_ir_operation(let_statement.initializer(), ctx, fn_ctx);

    let ir_assign = IrAssign::new(ir_variable_id, initializer_ir_operation);
    let ir_statement = IrStatement::Assign(ir_assign);
    fn_ctx.current_block_statements.push(ir_statement);
}

fn lower_to_ir_expression_statement(
    expression_statement: &ExpressionStatement,
    ctx: &LowerToIrContext,
    fn_ctx: &mut LowerToIrFunctionContext,
    returns_value: bool,
) {
    if returns_value {
        let ir_expression =
            lower_expression_to_ir_expression(expression_statement.expression(), ctx, fn_ctx);
        let ir_statement = IrStatement::Return(IrReturn::new(Some(ir_expression)));
        fn_ctx.current_block_statements.push(ir_statement);
    } else {
        // we could try to filter out useless code, but it's not guaranteed that there aren't
        // intended side effects from the lowering, so, for now, let's throw the result in a t_var.
        // in the future we may want to somehow enforce with the type system a way that we can't
        // create useless, dead ir code
        let ir_operation =
            lower_expression_to_ir_operation(expression_statement.expression(), ctx, fn_ctx);

        let result_type_info_id =
            ctx.nodes_to_type_infos[&expression_statement.expression().node_id()];
        let result_type_info = &ctx.type_infos[result_type_info_id];
        let t_var_ir_variable_id = fn_ctx.make_t_var(to_ir_type_info(result_type_info));

        let ir_statement = IrStatement::Assign(IrAssign::new(t_var_ir_variable_id, ir_operation));
        fn_ctx.current_block_statements.push(ir_statement);
    }
}

fn lower_to_ir_assign_statement(
    assign_statement: &AssignStatement,
    ctx: &LowerToIrContext,
    fn_ctx: &mut LowerToIrFunctionContext,
) {
    match assign_statement.destination() {
        Expression::Identifier(identifier) => {
            let destination_symbol_id = ctx.nodes_to_symbols[&identifier.node_id()];
            let destination_ir_variable_id = fn_ctx.symbols_to_variables[&destination_symbol_id];

            let ir_operation =
                lower_expression_to_ir_operation(assign_statement.value(), ctx, fn_ctx);
            let ir_assign = IrAssign::new(destination_ir_variable_id, ir_operation);
            let ir_statement = IrStatement::Assign(ir_assign);
            fn_ctx.current_block_statements.push(ir_statement);
        }
        _ => panic!("Assigning to non-Identifiers is not yet supported."),
    }
}

fn lower_expression_to_ir_operation(
    expression: &Expression,
    ctx: &LowerToIrContext,
    fn_ctx: &mut LowerToIrFunctionContext,
) -> IrOperation {
    match expression {
        Expression::Binary(binary_expression) => {
            let lhs = lower_expression_to_ir_expression(binary_expression.lhs(), ctx, fn_ctx);
            let rhs = lower_expression_to_ir_expression(binary_expression.rhs(), ctx, fn_ctx);
            IrOperation::Binary(IrBinaryOperation::new(
                lhs,
                rhs,
                lower_binary_operator(binary_expression.op()),
            ))
        }
        Expression::Negative(negative_expression) => {
            let operand =
                lower_expression_to_ir_expression(negative_expression.operand(), ctx, fn_ctx);
            IrOperation::Binary(IrBinaryOperation::new(
                operand,
                IrExpression::Int(-1),
                IrBinaryOperator::Multiply,
            ))
        }
        Expression::Call(call) => IrOperation::Call(lower_to_ir_call(call, ctx, fn_ctx)),
        Expression::Identifier(identifier) => {
            let identifier_symbol_id = ctx.nodes_to_symbols[&identifier.node_id()];
            let identifier_ir_variable_id = fn_ctx.symbols_to_variables[&identifier_symbol_id];
            let ir_expression = IrExpression::Variable(identifier_ir_variable_id);
            IrOperation::Load(ir_expression)
        }
        Expression::Integer(integer_literal) => {
            IrOperation::Load(IrExpression::Int(integer_literal.value()))
        }
        Expression::Double(double_literal) => {
            IrOperation::Load(IrExpression::Double(double_literal.value()))
        }
        Expression::String(string_literal) => {
            IrOperation::Load(IrExpression::String(string_literal.content().into()))
        }
    }
}

fn lower_expression_to_ir_expression(
    expression: &Expression,
    ctx: &LowerToIrContext,
    fn_ctx: &mut LowerToIrFunctionContext,
) -> IrExpression {
    match expression {
        Expression::Binary(binary_expression) => {
            // make binary operation
            let lhs = lower_expression_to_ir_expression(binary_expression.lhs(), ctx, fn_ctx);
            let rhs = lower_expression_to_ir_expression(binary_expression.rhs(), ctx, fn_ctx);
            let ir_operation = IrOperation::Binary(IrBinaryOperation::new(
                lhs,
                rhs,
                lower_binary_operator(binary_expression.op()),
            ));

            // make destination temp var
            let result_type_info_id = ctx.nodes_to_type_infos[&binary_expression.node_id()];
            let result_type_info = &ctx.type_infos[result_type_info_id];
            let destination_ir_variable_id = fn_ctx.make_t_var(to_ir_type_info(result_type_info));

            // make assign statement to destination temp var
            let ir_assign = IrAssign::new(destination_ir_variable_id, ir_operation);
            fn_ctx
                .current_block_statements
                .push(IrStatement::Assign(ir_assign));

            // return location of temp var
            IrExpression::Variable(destination_ir_variable_id)
        }
        Expression::Negative(negative_expression) => {
            let operand =
                lower_expression_to_ir_expression(negative_expression.operand(), ctx, fn_ctx);
            let negative_one = IrExpression::Int(-1);
            let ir_operation = IrOperation::Binary(IrBinaryOperation::new(
                operand,
                negative_one,
                IrBinaryOperator::Multiply,
            ));
            let result_type_info_id = ctx.nodes_to_type_infos[&negative_expression.node_id()];
            let result_type_info = &ctx.type_infos[result_type_info_id];
            let destination_ir_variable_id = fn_ctx.make_t_var(to_ir_type_info(result_type_info));
            let ir_assign = IrAssign::new(destination_ir_variable_id, ir_operation);

            // push the statement which does the multiply by negative one
            fn_ctx
                .current_block_statements
                .push(IrStatement::Assign(ir_assign));

            IrExpression::Variable(destination_ir_variable_id)
        }
        Expression::Call(call) => {
            let ir_call = lower_to_ir_call(call, ctx, fn_ctx);

            // make temp var
            let return_type_info_id = ctx.nodes_to_type_infos[&call.node_id()];
            let return_type_info = &ctx.type_infos[return_type_info_id];
            let return_ir_type_info = to_ir_type_info(return_type_info);
            let t_var_ir_variable_id = fn_ctx.make_t_var(return_ir_type_info);

            // assign call to temp var, return temp var expression
            let ir_operation = IrOperation::Call(ir_call);
            let ir_assign = IrAssign::new(t_var_ir_variable_id, ir_operation);
            fn_ctx
                .current_block_statements
                .push(IrStatement::Assign(ir_assign));

            // return an expression referencing the temp var
            IrExpression::Variable(t_var_ir_variable_id)
        }
        Expression::Identifier(identifier) => {
            let rhs_symbol_id = ctx.nodes_to_symbols[&identifier.node_id()];
            if let Some(rhs_ir_variable_id) = fn_ctx.symbols_to_variables.get(&rhs_symbol_id) {
                IrExpression::Variable(*rhs_ir_variable_id)
            } else if let Some(rhs_ir_parameter_id) =
                fn_ctx.symbols_to_parameters.get(&rhs_symbol_id)
            {
                IrExpression::Parameter(*rhs_ir_parameter_id)
            } else {
                panic!(
                    "Could not find parameter or variable for symbol_id {}",
                    rhs_symbol_id
                );
            }
        }
        Expression::Integer(integer_literal) => IrExpression::Int(integer_literal.value()),
        Expression::Double(double_literal) => IrExpression::Double(double_literal.value()),
        Expression::String(string_literal) => IrExpression::String(string_literal.content().into()),
    }
}

fn lower_to_ir_call(
    call: &Call,
    ctx: &LowerToIrContext,
    fn_ctx: &mut LowerToIrFunctionContext,
) -> IrCall {
    let callee_symbol_id = ctx.nodes_to_symbols[&call.callee().node_id()];
    let callee_symbol = &ctx.symbols[callee_symbol_id];
    match callee_symbol {
        Symbol::Function(function_symbol) => {
            let arguments = call
                .arguments()
                .iter()
                .map(|e| lower_expression_to_ir_expression(e, ctx, fn_ctx))
                .collect();
            IrCall::new(
                function_symbol.fqn_owned(),
                arguments,
                function_symbol.is_extern(),
            )
        }
        _ => panic!("Expected function symbol"),
    }
}