use crate::constants_table::ConstantsTable;
use crate::ir::ir_assign::IrAssign;
use crate::ir::ir_binary_operation::IrBinaryOperator;
use crate::ir::ir_block::IrBlock;
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;
use crate::ir::ir_return::IrReturn;
use crate::ir::ir_statement::IrStatement;
use crate::ir::ir_type_info::{IrTypeInfo, IrTypeInfoId};
use crate::ir::ir_variable::IrVariableId;
use crate::ir::variable_locations::{VariableLocation, VariableLocations};
use dvm_lib::instruction::{
    AddOperand, Instruction, Location, LocationOrInteger, LocationOrNumber, MoveOperand,
    MultiplyOperand, PushOperand, ReturnOperand, SubtractOperand,
};
use std::collections::HashMap;

struct FunctionAssemblyContext<'a> {
    type_infos: &'a Vec<IrTypeInfo>,
    variables_to_type_infos: &'a HashMap<IrVariableId, IrTypeInfoId>,
    variable_locations: &'a VariableLocations,
    parameters: &'a [IrParameter],
    constants_table: &'a mut ConstantsTable,
    instructions: Vec<Instruction>,
}

pub fn assemble_ir_function(
    ir_function: &IrFunction,
    type_infos: &Vec<IrTypeInfo>,
    variables_to_type_infos: &HashMap<IrVariableId, IrTypeInfoId>,
    variable_locations: &VariableLocations,
    constants_table: &mut ConstantsTable,
) -> Vec<Instruction> {
    let mut ctx = FunctionAssemblyContext {
        type_infos,
        variables_to_type_infos,
        variable_locations,
        constants_table,
        parameters: ir_function.parameters(),
        instructions: Vec::new(),
    };
    for block in ir_function.blocks() {
        assemble_ir_block(block, &mut ctx);
    }
    ctx.instructions
}

fn assemble_ir_block(block: &IrBlock, ctx: &mut FunctionAssemblyContext) {
    for statement in block.statements() {
        assemble_ir_statement(statement, ctx);
    }
}

fn assemble_ir_statement(statement: &IrStatement, ctx: &mut FunctionAssemblyContext) {
    match statement {
        IrStatement::Assign(ir_assign) => {
            assemble_ir_assign(ir_assign, ctx);
        }
        IrStatement::Call(ir_call) => {
            assemble_ir_call(ir_call, ctx);
        }
        IrStatement::Return(ir_return) => {
            assemble_ir_return(ir_return, ctx);
        }
        IrStatement::SetField(ir_set_field) => {
            todo!()
        }
    }
}

fn assemble_ir_assign(ir_assign: &IrAssign, ctx: &mut FunctionAssemblyContext) {
    let destination_location = to_location(
        ctx.variable_locations
            .get_variable_location(&ir_assign.destination()),
    );
    match ir_assign.initializer() {
        IrOperation::GetFieldRef(_) => {}
        IrOperation::GetFieldRefMut(_) => {}
        IrOperation::ReadField(_) => {}
        IrOperation::Load(ir_expression) => {
            let move_operand = to_move_operand(ir_expression, ctx);
            let instruction = Instruction::Move(move_operand, destination_location);
            ctx.instructions.push(instruction);
        }
        IrOperation::Binary(ir_binary_operation) => {
            let instruction = match ir_binary_operation.op() {
                IrBinaryOperator::Multiply => Instruction::Multiply(
                    to_multiply_operand(ir_binary_operation.left(), ctx),
                    to_multiply_operand(ir_binary_operation.right(), ctx),
                    destination_location,
                ),
                IrBinaryOperator::Divide => Instruction::Divide(
                    to_location_or_number(ir_binary_operation.left(), ctx),
                    to_location_or_number(ir_binary_operation.right(), ctx),
                    destination_location,
                ),
                IrBinaryOperator::Modulo => Instruction::Modulo(
                    to_location_or_number(ir_binary_operation.left(), ctx),
                    to_location_or_number(ir_binary_operation.right(), ctx),
                    destination_location,
                ),
                IrBinaryOperator::Add => Instruction::Add(
                    to_add_operand(ir_binary_operation.left(), ctx),
                    to_add_operand(ir_binary_operation.right(), ctx),
                    destination_location,
                ),
                IrBinaryOperator::Subtract => Instruction::Subtract(
                    to_subtract_operand(ir_binary_operation.left(), ctx),
                    to_subtract_operand(ir_binary_operation.right(), ctx),
                    destination_location,
                ),
                IrBinaryOperator::LeftShift => Instruction::LeftShift(
                    to_location_or_integer(ir_binary_operation.left(), ctx),
                    to_location_or_integer(ir_binary_operation.right(), ctx),
                    destination_location,
                ),
                IrBinaryOperator::RightShift => Instruction::RightShift(
                    to_location_or_integer(ir_binary_operation.left(), ctx),
                    to_location_or_integer(ir_binary_operation.right(), ctx),
                    destination_location,
                ),
                IrBinaryOperator::BitwiseAnd => Instruction::BitwiseAnd(
                    to_location_or_integer(ir_binary_operation.left(), ctx),
                    to_location_or_integer(ir_binary_operation.right(), ctx),
                    destination_location,
                ),
                IrBinaryOperator::BitwiseXor => Instruction::BitwiseXor(
                    to_location_or_integer(ir_binary_operation.left(), ctx),
                    to_location_or_integer(ir_binary_operation.right(), ctx),
                    destination_location,
                ),
                IrBinaryOperator::BitwiseOr => Instruction::BitwiseOr(
                    to_location_or_integer(ir_binary_operation.left(), ctx),
                    to_location_or_integer(ir_binary_operation.right(), ctx),
                    destination_location,
                ),
            };

            ctx.instructions.push(instruction);
        }
        IrOperation::Call(_) => {
            todo!()
        }
        IrOperation::Allocate(_) => {
            todo!()
        }
    }
}

fn assemble_ir_call(ir_call: &IrCall, ctx: &mut FunctionAssemblyContext) {
    // push all args
    let args_as_push_operands = ir_call
        .arguments()
        .iter()
        .map(|ir_expression| to_push_operand(ir_expression, ctx))
        .collect::<Vec<_>>();
    for push_operand in args_as_push_operands {
        ctx.instructions.push(Instruction::Push(push_operand));
    }

    // push invoke instruction
    if ir_call.is_extern() {
        ctx.instructions.push(Instruction::InvokePlatformStatic(
            ir_call.function_name_owned(),
            ir_call.arguments().len(),
        ));
    } else {
        ctx.instructions.push(Instruction::InvokeStatic(
            ir_call.function_name_owned(),
            ir_call.arguments().len(),
        ));
    }
}

fn assemble_ir_return(ir_return: &IrReturn, ctx: &mut FunctionAssemblyContext) {
    if let Some(ir_expression) = ir_return.value() {
        let return_operand = to_return_operand(ir_expression, ctx);
        ctx.instructions
            .push(Instruction::SetReturnValue(return_operand));
    }
    ctx.instructions.push(Instruction::Return);
}

fn to_move_operand(ir_expression: &IrExpression, ctx: &mut FunctionAssemblyContext) -> MoveOperand {
    match ir_expression {
        IrExpression::Parameter(ir_parameter_id) => MoveOperand::Location(
            Location::StackFrameOffset(ctx.parameters[*ir_parameter_id].stack_offset()),
        ),
        IrExpression::Variable(ir_variable_id) => MoveOperand::Location(to_location(
            ctx.variable_locations.get_variable_location(ir_variable_id),
        )),
        IrExpression::Int(i) => MoveOperand::Int(*i),
        IrExpression::Double(d) => MoveOperand::Double(*d),
        IrExpression::String(s) => {
            let constant_name = ctx.constants_table.get_or_insert(s);
            MoveOperand::String(constant_name)
        }
    }
}

fn to_multiply_operand(
    ir_expression: &IrExpression,
    ctx: &mut FunctionAssemblyContext,
) -> MultiplyOperand {
    match ir_expression {
        IrExpression::Parameter(ir_parameter_id) => MultiplyOperand::Location(
            Location::StackFrameOffset(ctx.parameters[*ir_parameter_id].stack_offset()),
        ),
        IrExpression::Variable(ir_variable_id) => {
            let ir_type_info_id = ctx.variables_to_type_infos[ir_variable_id];
            let ir_type_info = &ctx.type_infos[ir_type_info_id];
            match ir_type_info {
                IrTypeInfo::Int | IrTypeInfo::Double => {
                    let location =
                        to_location(ctx.variable_locations.get_variable_location(ir_variable_id));
                    MultiplyOperand::Location(location)
                }
                _ => panic!("Attempt to multiply non-number (found {})", ir_type_info),
            }
        }
        IrExpression::Int(i) => MultiplyOperand::Int(*i),
        IrExpression::Double(d) => MultiplyOperand::Double(*d),
        IrExpression::String(_) => {
            panic!("Attempt to multiply with a string");
        }
    }
}

fn to_add_operand(ir_expression: &IrExpression, ctx: &mut FunctionAssemblyContext) -> AddOperand {
    match ir_expression {
        IrExpression::Parameter(ir_parameter_id) => {
            let ir_parameter = &ctx.parameters[*ir_parameter_id];
            match ir_parameter.type_info() {
                IrTypeInfo::Int | IrTypeInfo::Double | IrTypeInfo::String => {
                    AddOperand::Location(Location::StackFrameOffset(ir_parameter.stack_offset()))
                }
            }
        }
        IrExpression::Variable(ir_variable_id) => {
            let ir_type_info_id = ctx.variables_to_type_infos[ir_variable_id];
            let ir_type_info = &ctx.type_infos[ir_type_info_id];
            match ir_type_info {
                IrTypeInfo::Int | IrTypeInfo::Double | IrTypeInfo::String => AddOperand::Location(
                    to_location(ctx.variable_locations.get_variable_location(ir_variable_id)),
                ),
            }
        }
        IrExpression::Int(i) => AddOperand::Int(*i),
        IrExpression::Double(d) => AddOperand::Double(*d),
        IrExpression::String(s) => {
            let constant_name = ctx.constants_table.get_or_insert(s);
            AddOperand::String(constant_name)
        }
    }
}

fn to_subtract_operand(
    ir_expression: &IrExpression,
    ctx: &mut FunctionAssemblyContext,
) -> SubtractOperand {
    match ir_expression {
        IrExpression::Parameter(ir_parameter_id) => {
            let ir_parameter = &ctx.parameters[*ir_parameter_id];
            match ir_parameter.type_info() {
                IrTypeInfo::Int | IrTypeInfo::Double => SubtractOperand::Location(
                    Location::StackFrameOffset(ir_parameter.stack_offset()),
                ),
                _ => panic!(
                    "Attempt to subtract with non-integer type (found {})",
                    ir_parameter.type_info()
                ),
            }
        }
        IrExpression::Variable(ir_variable) => {
            let ir_type_info_id = ctx.variables_to_type_infos[ir_variable];
            let ir_type_info = &ctx.type_infos[ir_type_info_id];
            match ir_type_info {
                IrTypeInfo::Int | IrTypeInfo::Double => SubtractOperand::Location(to_location(
                    ctx.variable_locations.get_variable_location(ir_variable),
                )),
                _ => panic!(
                    "Attempt to subtract with non-number type (found {})",
                    ir_type_info
                ),
            }
        }
        IrExpression::Int(i) => SubtractOperand::Int(*i),
        IrExpression::Double(d) => SubtractOperand::Double(*d),
        IrExpression::String(_) => {
            panic!("Attempt to subtract with a string");
        }
    }
}

fn to_location_or_number(
    ir_expression: &IrExpression,
    ctx: &mut FunctionAssemblyContext,
) -> LocationOrNumber {
    match ir_expression {
        IrExpression::Parameter(ir_parameter_id) => {
            let ir_parameter = &ctx.parameters[*ir_parameter_id];
            LocationOrNumber::Location(Location::StackFrameOffset(ir_parameter.stack_offset()))
        }
        IrExpression::Variable(ir_variable_id) => LocationOrNumber::Location(to_location(
            ctx.variable_locations.get_variable_location(ir_variable_id),
        )),
        IrExpression::Int(i) => LocationOrNumber::Int(*i),
        IrExpression::Double(d) => LocationOrNumber::Double(*d),
        _ => panic!(
            "Attempt to convert {} to a location or number",
            ir_expression
        ),
    }
}

fn to_push_operand(ir_expression: &IrExpression, ctx: &mut FunctionAssemblyContext) -> PushOperand {
    match ir_expression {
        IrExpression::Parameter(ir_parameter_id) => PushOperand::Location(
            Location::StackFrameOffset(ctx.parameters[*ir_parameter_id].stack_offset()),
        ),
        IrExpression::Variable(ir_variable_id) => PushOperand::Location(to_location(
            ctx.variable_locations.get_variable_location(ir_variable_id),
        )),
        IrExpression::Int(i) => PushOperand::Int(*i),
        IrExpression::Double(d) => PushOperand::Double(*d),
        IrExpression::String(s) => {
            let constant_name = ctx.constants_table.get_or_insert(s);
            PushOperand::String(constant_name)
        }
    }
}

fn to_return_operand(
    ir_expression: &IrExpression,
    ctx: &mut FunctionAssemblyContext,
) -> ReturnOperand {
    match ir_expression {
        IrExpression::Parameter(ir_parameter_id) => ReturnOperand::Location(
            Location::StackFrameOffset(ctx.parameters[*ir_parameter_id].stack_offset()),
        ),
        IrExpression::Variable(ir_variable) => ReturnOperand::Location(to_location(
            ctx.variable_locations.get_variable_location(ir_variable),
        )),
        IrExpression::Int(i) => ReturnOperand::Int(*i),
        IrExpression::Double(d) => ReturnOperand::Double(*d),
        IrExpression::String(s) => {
            let constant_name = ctx.constants_table.get_or_insert(s);
            ReturnOperand::String(constant_name)
        }
    }
}

fn to_location_or_integer(
    ir_expression: &IrExpression,
    ctx: &mut FunctionAssemblyContext,
) -> LocationOrInteger {
    match ir_expression {
        IrExpression::Parameter(ir_parameter_id) => LocationOrInteger::Location(
            Location::StackFrameOffset(ctx.parameters[*ir_parameter_id].stack_offset()),
        ),
        IrExpression::Variable(ir_variable_id) => LocationOrInteger::Location(to_location(
            ctx.variable_locations.get_variable_location(ir_variable_id),
        )),
        IrExpression::Int(i) => LocationOrInteger::Int(*i),
        _ => panic!(
            "Attempt to convert {} to a location or integer",
            ir_expression
        ),
    }
}

fn to_location(variable_location: VariableLocation) -> Location {
    match variable_location {
        VariableLocation::Register(register_assignment) => Location::Register(register_assignment),
        VariableLocation::Stack(stack_frame_offset) => {
            Location::StackFrameOffset(stack_frame_offset)
        }
    }
}