use crate::asm::{
    AsmBinaryOperator, AsmFunction, AsmInstruction, AsmOperand, ControlUnit, ControlUnitId,
    VirtualRegisterId,
};
use crate::ir::{
    IrAssign, IrBinaryOperation, IrBinaryOperator, IrCall, IrCallType, IrExpression, IrFunction,
    IrLiteral, IrReturn, IrStatement,
};
use std::collections::HashMap;
use std::rc::Rc;

struct AssemblyContext {
    control_unit_counter: ControlUnitId,
    control_units: Vec<ControlUnit>,
    virtual_register_counter: VirtualRegisterId,
    virtual_registers: HashMap<Rc<str>, VirtualRegisterId>,
}

impl AssemblyContext {
    pub fn new() -> Self {
        Self {
            control_unit_counter: 0,
            control_units: Vec::new(),
            virtual_register_counter: 0,
            virtual_registers: HashMap::new(),
        }
    }

    fn next_control_unit_id(&mut self) -> ControlUnitId {
        let id = self.control_unit_counter;
        self.control_unit_counter += 1;
        id
    }

    pub fn new_control_unit(&mut self) -> ControlUnitId {
        let new_id = self.next_control_unit_id();
        self.control_units.push(ControlUnit::new(new_id));
        new_id
    }

    pub fn current_control_unit(&mut self) -> &ControlUnit {
        if self.control_units.is_empty() {
            self.new_control_unit();
        }
        self.control_units.last().unwrap()
    }

    pub fn current_control_unit_mut(&mut self) -> &mut ControlUnit {
        if self.control_units.is_empty() {
            self.new_control_unit();
        }
        self.control_units.last_mut().unwrap()
    }

    pub fn into_control_units(self) -> Vec<ControlUnit> {
        self.control_units
    }

    pub fn get_virtual_register(&self, name: &str) -> VirtualRegisterId {
        *self
            .virtual_registers
            .get(name)
            .expect(&format!("Did not set a virtual register for name {}", name))
    }

    pub fn next_virtual_register(&mut self, name: &str) -> VirtualRegisterId {
        if self.virtual_registers.contains_key(name) {
            panic!(
                "Should not already have a virtual register with name {}",
                name
            );
        }
        let virtual_register_id = self.virtual_register_counter;
        self.virtual_register_counter += 1;
        self.virtual_registers
            .insert(name.into(), virtual_register_id);
        virtual_register_id
    }
}

pub fn assemble_ir_function(ir_function: &IrFunction) -> AsmFunction {
    let mut context = AssemblyContext::new();

    // body
    for statement in ir_function.statements() {
        assemble_ir_statement(statement, &mut context);
    }

    AsmFunction::new(ir_function.fqn().into(), context.into_control_units())
}

fn assemble_ir_statement(ir_statement: &IrStatement, context: &mut AssemblyContext) {
    match ir_statement {
        IrStatement::Allocate(_) => {
            todo!()
        }
        IrStatement::Call(ir_call) => {
            assemble_ir_call(ir_call, context);
        }
        IrStatement::Return(ir_return) => {
            assemble_ir_return(ir_return, context);
        }
        IrStatement::Assign(ir_assign) => {
            assemble_ir_assign(ir_assign, context);
        }
        IrStatement::MakeClosure(_) => {
            todo!()
        }
        IrStatement::BinaryOperation(ir_binary_operation) => {
            assemble_ir_binary_operation(ir_binary_operation, context)
        }
    }
}

fn assemble_ir_assign(ir_assign: &IrAssign, context: &mut AssemblyContext) {
    let source_operand = match ir_assign.rhs() {
        IrExpression::Variable(ir_variable) => {
            let virtual_register_id = context.get_virtual_register(ir_variable.name());
            AsmOperand::VirtualRegister(virtual_register_id)
        }
        IrExpression::Literal(ir_literal) => ir_literal_to_asm_operand(ir_literal),
    };
    let destination = context.next_virtual_register(ir_assign.lhs().name());
    context
        .current_control_unit_mut()
        .append_instruction(AsmInstruction::Move {
            source: source_operand,
            destination,
        })
}

fn ir_literal_to_asm_operand(ir_literal: &IrLiteral) -> AsmOperand {
    match ir_literal {
        IrLiteral::I8(i) => AsmOperand::I8(*i),
        IrLiteral::I16(i) => AsmOperand::I16(*i),
        IrLiteral::I32(i) => AsmOperand::I32(*i),
        IrLiteral::I64(i) => AsmOperand::I64(*i),
        IrLiteral::I128(i) => AsmOperand::I128(*i),
        IrLiteral::ISize(i) => AsmOperand::ISize(*i),
        IrLiteral::U8(u) => AsmOperand::U8(*u),
        IrLiteral::U16(u) => AsmOperand::U16(*u),
        IrLiteral::U32(u) => AsmOperand::U32(*u),
        IrLiteral::U64(u) => AsmOperand::U64(*u),
        IrLiteral::U128(u) => AsmOperand::U128(*u),
        IrLiteral::USize(u) => AsmOperand::USize(*u),
        IrLiteral::Boolean(b) => AsmOperand::Boolean(*b),
        IrLiteral::Float(f) => AsmOperand::Float(*f),
        IrLiteral::Double(d) => AsmOperand::Double(*d),
        IrLiteral::String(s) => AsmOperand::String(s.clone()),
    }
}

fn assemble_ir_call(ir_call: &IrCall, context: &mut AssemblyContext) {
    // push args
    for argument in ir_call.arguments() {
        match argument {
            IrExpression::Variable(ir_variable) => {
                let variable_register_id = context.get_virtual_register(ir_variable.name());
                context
                    .current_control_unit_mut()
                    .append_instruction(AsmInstruction::Push {
                        source: AsmOperand::VirtualRegister(variable_register_id),
                    })
            }
            IrExpression::Literal(ir_literal) => context
                .current_control_unit_mut()
                .append_instruction(AsmInstruction::Push {
                    source: ir_literal_to_asm_operand(ir_literal),
                }),
        }
    }
    // issue call
    match ir_call.call_type() {
        IrCallType::Static => {
            context
                .current_control_unit_mut()
                .append_instruction(AsmInstruction::InvokeStatic {
                    fqn: ir_call.function_name_owned(),
                });
        }
        IrCallType::Object => {
            context
                .current_control_unit_mut()
                .append_instruction(AsmInstruction::InvokeObject {
                    fqn: ir_call.function_name_owned(),
                })
        }
        IrCallType::PlatformStatic => {}
        IrCallType::PlatformObject => {}
    }
}

fn assemble_ir_return(ir_return: &IrReturn, context: &mut AssemblyContext) {
    let operand = ir_return.expression().map(|expression| match expression {
        IrExpression::Variable(ir_variable) => {
            let variable_virtual_register_id = context.get_virtual_register(ir_variable.name());
            AsmOperand::VirtualRegister(variable_virtual_register_id)
        }
        IrExpression::Literal(ir_literal) => ir_literal_to_asm_operand(ir_literal),
    });
    context
        .current_control_unit_mut()
        .append_instruction(AsmInstruction::Return { operand });
}

fn assemble_ir_binary_operation(
    ir_binary_operation: &IrBinaryOperation,
    context: &mut AssemblyContext,
) {
    let left = match ir_binary_operation.left() {
        IrExpression::Variable(ir_variable) => {
            AsmOperand::VirtualRegister(context.get_virtual_register(ir_variable.name()))
        }
        IrExpression::Literal(ir_literal) => ir_literal_to_asm_operand(ir_literal),
    };
    let right = match ir_binary_operation.right() {
        IrExpression::Variable(ir_variable) => {
            AsmOperand::VirtualRegister(context.get_virtual_register(ir_variable.name()))
        }
        IrExpression::Literal(ir_literal) => ir_literal_to_asm_operand(ir_literal),
    };
    let destination = context.next_virtual_register(ir_binary_operation.destination().name());
    let operator = match ir_binary_operation.operator() {
        IrBinaryOperator::Add => AsmBinaryOperator::Add,
        IrBinaryOperator::Subtract => AsmBinaryOperator::Subtract,
        IrBinaryOperator::Multiply => AsmBinaryOperator::Multiply,
        IrBinaryOperator::Divide => AsmBinaryOperator::Divide,
        IrBinaryOperator::Exponent => {
            todo!()
        }
    };

    let instruction = AsmInstruction::BinaryOperation {
        destination,
        left,
        right,
        operator,
    };
    context
        .current_control_unit_mut()
        .append_instruction(instruction);
}