use std::collections::HashMap;
use std::sync::atomic::AtomicU32;
use std::time::SystemTime;

use chrono::{DateTime, Local};

use super::registers::RegisterAllocator;
use crate::backend::dsa::instruction::{InsBlock as IB, Instruction as I, Label};
use crate::backend::dsa::registers::Register;

use crate::model::{
    AssignmentOperator, BinaryOperator, Call, CompilerError, ConstExpr, Declaration,
    Dependency, Expression, Number, Program, Statement, TypeId, UnaryOperator, Variable,
};

pub struct CodeGenerator {
    ast: Program,
    imports: HashMap<String, I>,
    globals: HashMap<String, I>,
    functions: Vec<IB>,
    symbols: Vec<String>,
    allocator: RegisterAllocator,
    is_library: bool,
}

impl CodeGenerator {
    pub fn new(ast: Program, is_lib: bool) -> Self {
        CodeGenerator {
            ast,
            imports: HashMap::new(),
            globals: HashMap::new(),
            functions: Vec::new(),
            symbols: Vec::new(),
            allocator: RegisterAllocator::new(),
            is_library: is_lib,
        }
    }

    pub fn include(&mut self, name: impl Into<String>, path: impl Into<String>) {
        let name = name.into();
        self.imports.insert(name.clone(), I::include(name, path));
    }

    fn is_global(&self, name: &str) -> bool {
        // Check if this variable is in the globals list
        self.globals.contains_key(name)
    }

    pub fn generate(&mut self) -> Result<String, CompilerError> {
        // always include the print library for debugging!
        if !self.is_library {
            self.include("print", "./lib/print.dsa");
        }

        for block in self.ast.clone().declarations {
            match block {
                Declaration::Variable {
                    var: Variable { name, .. },
                    ..
                } => self.symbols.push(name),
                Declaration::Function { name, .. } => self.symbols.push(name),
                Declaration::Dependency(Dependency { name, .. }) => {
                    self.symbols.push(name)
                }
                Declaration::Struct { .. } => {} /* we can't do any code generation for
                                                  * a struct yet. we may need to later
                                                  * once these become class-like
                                                  * objects with implementations */
            }
        }

        for block in self.ast.clone().declarations {
            self.generate_block(block.clone())?;
        }

        let assembly = self.generate_layout()?;
        Ok(assembly
            .iter()
            .map(|i| i.to_string())
            .collect::<Vec<_>>()
            .join("\n"))
    }

    fn generate_layout(&mut self) -> Result<IB, CompilerError> {
        let datetime: DateTime<Local> = SystemTime::now().into();

        let mut block = IB::new();

        block.extend(vec![
            I::global_comment(format!(
                "GENERATED BY DSC COMPILER
                Generated at {}",
                datetime.format("%Y-%m-%d %H:%M:%S")
            )),
            I::Newline,
            I::global_comment("Imports"),
        ]);

        block.extend(self.imports.values().cloned().collect::<Vec<_>>());

        block.extend(vec![
            I::Newline,
            I::global_comment("Globals & Reserved Memory"),
        ]);

        block.extend(self.globals.values().cloned().collect::<Vec<_>>());

        if !self.is_library {
            block.extend(vec![
                I::Newline,
                I::global_comment("Entry Point"),
                I::db_word("stack", 0x10000),
                I::db_string("message", "Process Exited with code:"),
                // init function for stack setup.
                I::label("_init"),
                I::ldw_label("stack", Register::Bpr),
                I::mov(Register::Bpr, Register::Spr),
                I::push(Register::Zero),
                I::call("main"),
                I::call("print::print_newline"),
                I::lwi_label("message", Register::Rg0),
                I::push(Register::Rg0),
                I::call("print::print"),
                I::pop(Register::Zero),
                I::call("print::print_hex_word"),
                I::pop(Register::Zero),
                I::Hlt,
            ]);
        }

        block.extend(vec![
            I::Newline,
            // default return block boilerplate
            I::global_comment("Return"),
            I::label("_ret"),
            I::mov(Register::Bpr, Register::Spr),
            I::pop(Register::Bpr),
            I::Return,
        ]);

        for function in self.functions.iter() {
            block.extend(function.iter().cloned());
        }

        Ok(block)
    }

    fn generate_global(&mut self, name: &str, init: Option<ConstExpr>) {
        let init = init.unwrap_or(ConstExpr::Number(0));
        match init {
            ConstExpr::Number(value) => {
                self.globals
                    .insert(name.to_string(), I::db_word(name, value as u32));
            }
            ConstExpr::String(str) => {
                self.globals
                    .insert(name.to_string(), I::db_string(name, str));
            }
        }
    }

    fn generate_block(&mut self, block: Declaration) -> Result<(), CompilerError> {
        match block {
            Declaration::Variable { var, init, .. } => {
                self.generate_global(&var.name, init)
            }
            Declaration::Function {
                name,
                params,
                body,
                return_type,
            } => {
                let func = self.generate_function(&name, &params, &body, return_type);
                self.functions.push(func);
            }
            Declaration::Dependency(Dependency { name, path }) => {
                self.include(name, path);
            }
            Declaration::Struct { .. } => {} /* can't do any codegen for these yet,
                                              * they're just types. */
        };

        Ok(())
    }

    // Example: Generate code for a function
    fn generate_function(
        &mut self,
        name: &str,
        params: &[Variable],
        body: &[Statement],
        return_type: TypeId,
    ) -> IB {
        let mut code = IB::new();

        // Reset allocator for new function
        self.allocator.reset();

        let fmtparams = params
            .iter()
            .map(|p| format!("{}: {}", p.name, p.type_id))
            .collect::<Vec<String>>()
            .join(", ");

        code.extend(vec![
            I::global_comment(format!("fn {name}({fmtparams}) -> {return_type}")),
            I::label(name),
            I::push(Register::Bpr),
            I::mov(Register::Spr, Register::Bpr),
        ]);

        // Allocate parameters to registers or stack locations
        for (i, param) in params.iter().enumerate() {
            let offset = 8 + (i as i32 * 4); // Parameters start at bpr+8
            // Track that this parameter is at a stack location
            let (reg, load_code) = self.allocator.alloc_var(&param.name).unwrap();
            code.append(load_code);
            code.push(I::ldw_reg_offset(Register::Bpr, reg, offset));
        }

        // Generate code for function body
        for stmt in body {
            let stmt_code = self.generate_statement(stmt, &mut code).unwrap();
            code.append(stmt_code);
        }

        // automatically return at function end
        if let Some(x) = code.iter().last()
            && let I::Jmp { target: Label(val) } = x
            && val == "_ret"
        {
        } else {
            code.push(I::jmp("_ret"));
        }

        code.insert(0, I::Newline);

        code
    }

    // Example: Generate code for a statement
    fn generate_statement(
        &mut self,
        stmt: &Statement,
        func_body: &mut IB,
    ) -> Result<IB, CompilerError> {
        let mut code = IB::new();

        match stmt {
            Statement::Declaration { var, value } => {
                if let Some(expr) = value {
                    // Evaluate expression
                    let (result_reg, expr_code) =
                        self.generate_expression(expr, true, func_body)?;
                    code.append(expr_code);

                    // Store result in variable
                    let store_code = self.allocator.store_var(&var.name, &result_reg);
                    code.append(store_code);

                    // Free temporary register
                    self.allocator.free_temp(result_reg);
                } else {
                    // Just declaring variable without initialization
                    self.allocator.alloc_var(&var.name)?;
                }
            }

            Statement::Break => unimplemented!("need scope tracking first!"),
            Statement::Continue => unimplemented!("need scope tracking first!"),
            Statement::Defer(_func) => unimplemented!("we need scope tracking first!"),

            Statement::PtrWrite { ptr, value } => {
                let (result_reg, expr_code) =
                    self.generate_expression(value, true, func_body)?;
                code.append(expr_code);

                let (ptr_reg, ptr_code) =
                    self.generate_expression(ptr, true, func_body)?;
                code.append(ptr_code);

                code.push(I::stw_reg(result_reg, ptr_reg));

                self.allocator.free_temp(result_reg);
                self.allocator.free_temp(ptr_reg);
            }

            Statement::Assign {
                varname,
                value,
                operator,
            } => {
                // Evaluate expression
                let (result_reg, expr_code) =
                    self.generate_expression(value, true, func_body)?;
                code.append(expr_code);

                if *operator == AssignmentOperator::Assign {
                    // Check if this is a global variable
                    if self.is_global(varname) {
                        // Store to global label
                        code.push(I::stw_label(result_reg, varname.clone()))
                    } else {
                        // Store result in local variable
                        let store_code = self.allocator.store_var(varname, &result_reg);
                        code.append(store_code);
                    }

                    // Free temporary register
                    self.allocator.free_temp(result_reg);

                    return Ok(code);
                }

                // for more complex assignment cases we need an intermediate register.
                let (temp_reg, temp_code) = self.allocator.alloc_temp()?;
                code.append(temp_code);

                // fetch the value of the variable
                let var_reg = if self.is_global(varname) {
                    let instruction = I::ldw_label(varname.clone(), temp_reg);
                    code.push(instruction);
                    temp_reg
                } else {
                    let (rg, block) = self.allocator.load_var(varname)?;
                    code.append(block);
                    rg
                };

                let assign_code = match operator {
                    AssignmentOperator::Assign => {
                        unreachable!("assignment was already checked earlier.")
                    }
                    AssignmentOperator::AddAssign => {
                        I::add(var_reg, result_reg, temp_reg)
                    }
                    AssignmentOperator::SubAssign => {
                        I::sub(var_reg, result_reg, temp_reg)
                    }
                    AssignmentOperator::MulAssign => {
                        return Err(CompilerError::Unimplemented(
                            "TODO: implement multiplication for assignment".to_string(),
                        ));
                    }
                    AssignmentOperator::DivAssign => {
                        return Err(CompilerError::Unimplemented(
                            "TODO: write proper div function for DSA".to_string(),
                        ));
                    }
                    AssignmentOperator::ModAssign => {
                        return Err(CompilerError::Unimplemented(
                            "TODO: write proper mod function for DSA".to_string(),
                        ));
                    }
                    AssignmentOperator::AndAssign => {
                        I::and(var_reg, result_reg, temp_reg)
                    }
                    AssignmentOperator::OrAssign => I::or(var_reg, result_reg, temp_reg),
                    AssignmentOperator::XorAssign => {
                        I::xor(var_reg, result_reg, temp_reg)
                    }
                    AssignmentOperator::LeftShiftAssign => {
                        // this is only useful if we optimise out the register allocation
                        // inside value.
                        // if let Expression::Number { value, .. } = *value {
                        //     I::shl(var_reg, value, temp_reg)
                        // }
                        I::shl(var_reg, result_reg, 0, temp_reg)
                    }
                    AssignmentOperator::RightShiftAssign => {
                        // this is only useful if we optimise out the register allocation
                        // if let Expression::Number { value, .. } = *value {
                        //     I::shr(var_reg, value, temp_reg)
                        // }
                        I::shr(var_reg, result_reg, 0, temp_reg)
                    }
                };
                code.push(assign_code);

                // Check if this is a global variable
                if self.is_global(varname) {
                    // Store to global label
                    code.push(I::stw_label(temp_reg, varname.clone()))
                } else {
                    // Store result in local variable
                    let store_code = self.allocator.store_var(varname, &temp_reg);
                    code.append(store_code);
                }

                self.allocator.free_temp(result_reg);
                self.allocator.free_temp(temp_reg);
            }

            Statement::Return(expr) => {
                if let Some(e) = expr {
                    let (result_reg, expr_code) =
                        self.generate_expression(e, true, func_body)?;
                    code.append(expr_code);
                    code.push(I::stw_reg_offset(result_reg, Register::Bpr, 8));
                    code.push(I::jmp("_ret"));
                    self.allocator.free_temp(result_reg);
                }
            }

            Statement::If {
                condition,
                then_stmt,
                else_stmt,
            } => {
                // Generate condition
                let (cond_reg, cond_code) =
                    self.generate_expression(condition, true, func_body)?;
                code.append(cond_code);

                // Compare with zero
                code.push(I::cmp(cond_reg, Register::Zero));
                self.allocator.free_temp(cond_reg);

                // Generate unique labels
                let then_label = format!("_then_{}", self.get_unique_label());
                let else_label = format!("_else_{}", self.get_unique_label());
                let end_label = format!("_end_{}", self.get_unique_label());

                // Jump to else if condition is false (equal to zero)
                code.push(I::jeq(else_label.clone()));

                // Then block
                code.push(I::label(then_label));
                for s in then_stmt {
                    code.append(self.generate_statement(s, func_body)?);
                }

                if then_stmt.is_empty() {
                    code.push(I::Nop);
                }

                code.push(I::jmp(end_label.clone()));

                // Else block
                code.push(I::label(else_label));
                for s in else_stmt {
                    code.append(self.generate_statement(s, func_body)?);
                }

                if else_stmt.is_empty() {
                    code.push(I::Nop);
                }

                code.push(I::label(end_label));
            }

            Statement::While { condition, body } => {
                let loop_start = format!("_while_start_{}", self.get_unique_label());
                let loop_end = format!("_while_end_{}", self.get_unique_label());

                code.push(I::label(&loop_start));

                // Generate condition
                let (cond_reg, cond_code) =
                    self.generate_expression(condition, true, func_body)?;
                code.append(cond_code);

                code.push(I::cmp(cond_reg, Register::Zero));
                self.allocator.free_temp(cond_reg);

                code.push(I::jeq(loop_end.clone()));

                // Loop body
                for s in body {
                    code.append(self.generate_statement(s, func_body)?);
                }

                code.push(I::jmp(loop_start));
                code.push(I::label(loop_end));
            }

            Statement::Loop(body) => {
                let loop_start = format!("_loop_start_{}", self.get_unique_label());
                code.push(I::label(&loop_start));

                for s in body {
                    code.append(self.generate_statement(s, func_body)?);
                }

                code.push(I::jmp(loop_start));
            }

            Statement::Expression { expr } => {
                let (result_reg, expr_code) =
                    self.generate_expression(expr, false, func_body)?;
                code.append(expr_code);
                self.allocator.free_temp(result_reg);
            }

            Statement::Block(statements) => {
                for s in statements {
                    code.append(self.generate_statement(s, func_body)?);
                }
            }
        }

        Ok(code)
    }

    // Example: Generate code for an expression
    // Returns (register containing result, assembly code)
    fn generate_expression(
        &mut self,
        expr: &Expression,
        use_result: bool,
        func_body: &mut IB,
    ) -> Result<(Register, IB), CompilerError> {
        let mut code = IB::new();

        match expr {
            Expression::Empty => Ok((Register::Null, code)),

            Expression::Number(n) => match n {
                Number::Signed(value, _) => {
                    let (reg, alloc_code) = self.allocator.alloc_temp()?;
                    code.append(alloc_code);

                    // Load immediate value
                    code.push(I::lwi(*value as u32, reg));

                    Ok((reg, code))
                }
                Number::Unsigned(value, _) => {
                    let (reg, alloc_code) = self.allocator.alloc_temp()?;
                    code.append(alloc_code);

                    // Load immediate value
                    code.push(I::lwi(*value as u32, reg));

                    Ok((reg, code))
                }
            },

            Expression::CharLiteral(value) => {
                let (reg, alloc_code) = self.allocator.alloc_temp()?;
                code.append(alloc_code);

                // Load immediate value
                code.push(I::comment(format!("char literal '{value}'")));
                code.push(I::lwi(*value as u32, reg));

                Ok((reg, code))
            }

            Expression::StringLiteral(value) => {
                let (reg, alloc_code) = self.allocator.alloc_temp()?;
                code.append(alloc_code);

                // write string into memory
                let uuid = self.get_unique_label();

                func_body.insert(0, I::db_string(format!("str_{uuid}"), value));

                // Load pointer to string
                code.push(I::lwi_label(format!("str_{uuid}"), reg));

                Ok((reg, code))
            }

            Expression::ArrayLiteral { elements, type_id } => todo!(),
            Expression::StructLiteral {
                name,
                fields,
                type_id,
            } => todo!(),

            Expression::Variable { name, .. } => {
                if self.is_global(&name.name) {
                    // Allocate a temporary register for the global
                    let (reg, alloc_code) = self.allocator.alloc_temp()?;
                    code.append(alloc_code);

                    // Load from global label
                    code.push(I::ldw_label(name.name.clone(), reg));

                    Ok((reg, code))
                } else {
                    // Local variable - use existing allocator logic
                    let (reg, load_code) = self.allocator.load_var(&name.name)?;
                    code.append(load_code);
                    Ok((reg, code))
                }
            }

            Expression::Binary {
                op, left, right, ..
            } => {
                // Evaluate left operand
                let (left_reg, left_code) =
                    self.generate_expression(left, true, func_body)?;
                code.append(left_code);

                // Evaluate right operand
                let (right_reg, right_code) =
                    self.generate_expression(right, true, func_body)?;
                code.append(right_code);

                // Allocate result register
                let (result_reg, result_alloc) = self.allocator.alloc_temp()?;
                code.append(result_alloc);

                // Generate operation
                match op {
                    BinaryOperator::Add => {
                        code.push(I::add(left_reg, right_reg, result_reg));
                    }
                    BinaryOperator::Sub => {
                        code.push(I::sub(left_reg, right_reg, result_reg));
                    }
                    BinaryOperator::Mul => {
                        self.include("maths", "./lib/maths/core.dsa");
                        // Call multiply function
                        code.push(I::push(right_reg));
                        code.push(I::push(left_reg));
                        code.push(I::call("maths::multiply"));
                        code.push(I::pop(result_reg));
                        code.push(I::pop(Register::Zero));
                    }
                    BinaryOperator::Div => {
                        return Err(CompilerError::Unimplemented(
                            "TODO: write proper div function for DSA".to_string(),
                        ));
                        // self.include("maths", "./lib/maths/core.dsa");
                        // // Call divide function
                        // code.push(format!("\tpush {}", right_reg));
                        // code.push(format!("\tpush {}", left_reg));
                        // code.push("\tcall maths::divide".to_string());
                        // code.push(format!("\tpop {}", result_reg));
                        // code.push("\tpop zero".to_string());
                    }
                    BinaryOperator::Mod => {
                        return Err(CompilerError::Unimplemented(
                            "TODO: write proper mod function for DSA".to_string(),
                        ));
                        // self.include("maths", "./lib/maths/core.dsa");
                        // // Call modulo function
                        // code.push(format!("\tpush {}", right_reg));
                        // code.push(format!("\tpush {}", left_reg));
                        // code.push("\tcall maths::modulo".to_string());
                        // code.push(format!("\tpop {}", result_reg));
                        // code.push("\tpop zero".to_string());
                    }
                    BinaryOperator::BitwiseAnd => {
                        code.push(I::and(left_reg, right_reg, result_reg));
                    }
                    BinaryOperator::BitwiseOr => {
                        code.push(I::or(left_reg, right_reg, result_reg));
                    }
                    BinaryOperator::BitwiseXor => {
                        code.push(I::xor(left_reg, right_reg, result_reg));
                    }
                    BinaryOperator::LogicalAnd => {
                        return Err(CompilerError::Unimplemented(
                            "assembler/ISA does not yet support logical and!".to_string(),
                        ));
                    }
                    BinaryOperator::LogicalOr => {
                        return Err(CompilerError::Unimplemented(
                            "assembler/ISA does not yet support logical or!".to_string(),
                        ));
                    }
                    BinaryOperator::LeftShift => {
                        code.push(I::shl(left_reg, right_reg, 0, result_reg));
                    }
                    BinaryOperator::RightShift => {
                        code.push(I::shr(left_reg, right_reg, 0, result_reg));
                    }
                    // Comparison operators - return 1 (true) or 0 (false)
                    BinaryOperator::Equal => {
                        code.push(I::cmp(left_reg, right_reg));
                        code.push(I::lwi(1, result_reg));
                        let end_label = format!("_cmp_end_{}", self.get_unique_label());
                        code.push(I::jeq(end_label.clone()));
                        code.push(I::lwi(0, result_reg));
                        code.push(I::label(end_label));
                    }
                    BinaryOperator::NotEqual => {
                        code.push(I::cmp(left_reg, right_reg));
                        code.push(I::lwi(1, result_reg));
                        let end_label = format!("_cmp_end_{}", self.get_unique_label());
                        code.push(I::Jne {
                            target: Label(end_label.clone()),
                        });
                        code.push(I::lwi(0, result_reg));
                        code.push(I::label(&end_label));
                    }
                    BinaryOperator::LessThan => {
                        code.push(I::cmp(left_reg, right_reg));
                        code.push(I::lwi(1, result_reg));
                        let end_label = format!("_cmp_end_{}", self.get_unique_label());
                        code.push(I::Jlt {
                            target: Label(end_label.clone()),
                        });
                        code.push(I::lwi(0, result_reg));
                        code.push(I::label(&end_label));
                    }
                    BinaryOperator::LessOrEqual => {
                        code.push(I::cmp(left_reg, right_reg));
                        code.push(I::lwi(1, result_reg));
                        let end_label = format!("_cmp_end_{}", self.get_unique_label());
                        code.push(I::Jle {
                            target: Label(end_label.clone()),
                        });
                        code.push(I::lwi(0, result_reg));
                        code.push(I::label(&end_label));
                    }
                    BinaryOperator::GreaterThan => {
                        code.push(I::cmp(left_reg, right_reg));
                        code.push(I::lwi(1, result_reg));
                        let end_label = format!("_cmp_end_{}", self.get_unique_label());
                        code.push(I::Jgt {
                            target: Label(end_label.clone()),
                        });
                        code.push(I::lwi(0, result_reg));
                        code.push(I::label(&end_label));
                    }
                    BinaryOperator::GreaterOrEqual => {
                        code.push(I::cmp(left_reg, right_reg));
                        code.push(I::lwi(1, result_reg));
                        let end_label = format!("_cmp_end_{}", self.get_unique_label());
                        code.push(I::Jge {
                            target: Label(end_label.clone()),
                        });
                        code.push(I::lwi(0, result_reg));
                        code.push(I::label(&end_label));
                    } // _ => unimplemented!(),
                }

                // Free operand registers (allocator will protect variables)
                self.allocator.free_temp(left_reg);
                self.allocator.free_temp(right_reg);

                Ok((result_reg, code))
            }

            Expression::UnaryPostfix { op, operand, .. } => {
                let (operand_reg, operand_code) =
                    self.generate_expression(operand, true, func_body)?;
                code.append(operand_code);

                let (result_reg, result_alloc) = self.allocator.alloc_temp()?;
                code.append(result_alloc);

                match op {
                    UnaryOperator::Increment => {
                        // postfix increment - return old value
                        code.push(I::mov(operand_reg, result_reg));
                    }
                    UnaryOperator::Decrement => {
                        // postfix decrement - return old value
                        code.push(I::mov(operand_reg, result_reg));
                    }
                    _ => {
                        return Err(CompilerError::Generic(format!(
                            "{op} is prefix only!"
                        )));
                    }
                }

                self.allocator.free_temp(operand_reg);
                Ok((result_reg, code))
            }

            Expression::Unary { op, operand, .. } => {
                let (operand_reg, operand_code) =
                    self.generate_expression(operand, true, func_body)?;
                code.append(operand_code);

                let (result_reg, result_alloc) = self.allocator.alloc_temp()?;
                code.append(result_alloc);

                match op {
                    UnaryOperator::Minus => {
                        // Negate: result = 0 - operand
                        code.push(I::sub(Register::Zero, operand_reg, result_reg));
                    }
                    UnaryOperator::Plus => {
                        // Just move
                        code.push(I::mov(operand_reg, result_reg));
                    }
                    UnaryOperator::Dereference => {
                        code.push(I::ldw_reg(operand_reg, result_reg));
                    }
                    UnaryOperator::AddressOf => {
                        // ensure the referenced variable is on the stack and return its
                        // address.
                        let (offset, alloc_code) =
                            self.allocator.free_register(&operand_reg)?;
                        code.push(alloc_code);
                        code.push(I::iadd_dest(
                            Register::Spr,
                            offset - self.allocator.get_stack_offset(),
                            result_reg,
                        ));
                    }
                    UnaryOperator::SizeOf => {
                        if let Ok(id) = operand.type_id() {
                            let size = id.size();
                            code.push(I::lwi(size as u32, result_reg));
                        }
                    }
                    UnaryOperator::Increment => {
                        // prefix increment
                        code.push(I::mov(operand_reg, result_reg));
                        code.push(I::iadd_dest(operand_reg, 1, result_reg));
                    }
                    UnaryOperator::Decrement => {
                        // prefix decrement
                        code.push(I::mov(operand_reg, result_reg));
                        code.push(I::iadd_dest(operand_reg, -1, result_reg));
                    }
                    UnaryOperator::BitwiseNot => {
                        code.push(I::not(operand_reg, result_reg));
                    }
                    UnaryOperator::LogicalNot => {
                        return Err(CompilerError::Unimplemented(
                            "Assembler/ISA does not yet support logical not".to_string(),
                        ));
                    }
                    _ => {
                        return Err(CompilerError::Generic(format!(
                            "{op} is postfix only!"
                        )));
                    }
                }

                self.allocator.free_temp(operand_reg);
                Ok((result_reg, code))
            }

            Expression::Call {
                func: Call { name, args },
                ..
            } => {
                // first evaluate all the args we're going to need
                let mut arg_regs = Vec::new();
                for arg in args.iter().rev() {
                    let (arg_reg, arg_code) =
                        self.generate_expression(arg, true, func_body)?;
                    code.append(arg_code);
                    arg_regs.push(arg_reg);
                }

                // Save caller-saved registers and track which ones we saved
                let saved_regs = self.allocator.get_caller_saved_registers();
                for reg in &saved_regs {
                    // spill variables to stack
                    code.push(self.allocator.free_register(reg).unwrap().1);
                }

                // Evaluate and push arguments in reverse order
                for (i, arg_reg) in arg_regs.iter().enumerate() {
                    code.push(I::comment(format!("push arg {}", args.len() - 1 - i)));
                    code.push(I::push(*arg_reg));
                }

                if self.symbols.contains(&name.name) {
                    // Call local function
                    code.push(I::call(name.to_string()));
                } else if let Some(ns) = name.namespace.clone()
                    && self.imports.contains_key(&ns)
                {
                    code.push(I::call(name.to_string()));
                } else {
                    return Err(CompilerError::Undefined(name.clone()));
                }

                let result_reg: Register;

                if use_result {
                    let (temp_result_reg, result_alloc) = self.allocator.alloc_temp()?;
                    result_reg = temp_result_reg;

                    code.append(result_alloc);
                    code.push(I::pop(result_reg));

                    // Clean up arguments
                    if args.len() > 1 {
                        for _ in 0..(args.len() - 1) {
                            code.push(I::pop(Register::Zero));
                        }
                    }
                } else {
                    result_reg = Register::Zero;

                    // Clean up arguments
                    if args.len() > 0 {
                        for _ in 0..(args.len()) {
                            code.push(I::pop(Register::Zero));
                        }
                    }
                }

                // Free argument registers
                for reg in arg_regs {
                    self.allocator.free_temp(reg);
                }

                Ok((result_reg, code))
            }

            Expression::IndexAccess {
                expr,
                index,
                type_id,
            } => {
                let (expr_reg, expr_alloc) =
                    self.generate_expression(expr, true, func_body)?;
                code.append(expr_alloc);

                let (index_reg, index_alloc) =
                    self.generate_expression(index, true, func_body)?;
                code.append(index_alloc);

                let (result_reg, result_alloc) = self.allocator.alloc_temp()?;
                code.append(result_alloc);

                // add the expr pointer to the index to get the final address.
                code.push(I::add(expr_reg, index_reg, result_reg));
                // load the value at the address.
                code.push(I::ldw_reg(result_reg, result_reg));

                self.allocator.free_temp(expr_reg);
                self.allocator.free_temp(index_reg);

                Ok((result_reg, code))
            }
            Expression::MemberAccess {
                expr,
                field_name,
                type_id,
            } => Err(CompilerError::Unimplemented(
                "Structs are not yet implemented!".to_string(),
            )),

            Expression::TypeCast {
                expr,
                target_type,
                type_id,
            } => {
                let (expr_reg, expr_code) =
                    self.generate_expression(expr, true, func_body)?;

                // not sure if we actually need to do anything here.
                // for now we just return the previous expression.
                Ok((expr_reg, expr_code))
            }
        }
    }

    // Helper for generating unique labels
    fn get_unique_label(&mut self) -> String {
        // You'd implement a counter here
        static COUNTER: AtomicU32 = AtomicU32::new(0);

        let val = COUNTER.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        (val + 1).to_string()
    }
}
