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

use chrono::{DateTime, Local};

use super::registers::RegisterAllocator;
use crate::{block, comment, dsa};

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

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

fn import(name: &str, path: &str) -> String {
    format!("include {name}: \"{}\"", path)
}

impl CodeGenerator {
    const RET: &'static str = "\tjmp _ret";

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

    pub fn include(&mut self, name: &str, path: &str) {
        self.imports.insert(name.to_string(), path.to_string());
    }

    fn is_global(&self, name: &str) -> bool {
        // Check if this variable is in the globals list
        self.globals
            .iter()
            .any(|g| g.contains(&format!("dw {}:", name)))
    }

    pub fn generate(&mut self) -> Result<String, CompilerError> {
        // always include the print library for debugging!
        self.include("print", "./lib/io/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)
                }
            }
        }

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

        self.generate_layout()
    }

    fn generate_layout(&mut self) -> Result<String, CompilerError> {
        let datetime: DateTime<Local> = SystemTime::now().into();
        Ok(dsa![
            "",
            comment!("GENERATED BY DSC COMPILER"),
            comment!(format!(
                "Generated at {}",
                datetime.format("%Y-%m-%d %H:%M:%S")
            )),
            "",
            // imports
            comment!("Imports"),
            self.imports
                .iter()
                .map(|(k, v)| import(k, v))
                .collect::<Vec<String>>()
                .join("\n"),
            "",
            // reserved memory
            comment!("Globals & Reserved Memory"),
            self.globals.join("\n"),
            "",
            // entry point
            comment!("Entry Point"),
            "dw stack: 0x10000",
            "db message: \"Process Exited with code:\"",
            block! [ "_init"
                dsa![ldw stack, bpr],
                dsa![mov bpr, spr],
                dsa![push zero],
                dsa![call main],
                dsa![call print::print_newline],
                dsa![lwi message, rg0],
                dsa![push rg0],
                dsa![call print::print],
                dsa![pop zero],
                dsa![call print::print_hex_word],
                dsa![pop zero],
                dsa![hlt]
            ],
            "",
            comment!("Return"),
            block! [ "_ret"
                dsa![mov bpr, spr],
                dsa![pop bpr],
                dsa![return]
            ],
            comment!("Compiled Code Starts..."),
            // block! [ "main"
            //     dsa![push bpr],
            //     dsa![mov spr, bpr],
            //     dsa![lwi 67, rg1],
            //     dsa![stw rg1, spr, 8],
            //     dsa![mov bpr, spr],
            //     dsa![pop bpr],
            //     dsa![return]
            // ],
            self.functions.join("\n"),
        ])
    }

    fn generate_global(&mut self, name: &str, init: Option<ConstExpr>) {
        self.globals.push(format!(
            "dw {}: {}",
            name,
            init.unwrap_or(ConstExpr::Number(0))
        ))
    }

    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, ..
            } => {
                let func = self.generate_function(&name, &params, &body).join("\n");

                self.functions.push(format!("{func}\n"));
            }
            Declaration::Dependency(Dependency { name, path }) => {
                self.imports.insert(name, path);
            }
        };

        Ok(())
    }

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

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

        // Function prologue
        code.push(format!("{}:", name));
        code.push("\tpush bpr".to_string());
        code.push("\tmov spr, bpr".to_string());
        code.push(String::new());

        // 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.extend(load_code);
            code.push(format!("\tldw bpr, {}, {}", reg, offset));
        }

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

        // automatically return at function end
        if let Some(x) = code.last()
            && x == Self::RET
        {
        } else {
            code.push(Self::RET.to_string());
        }

        code
    }

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

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

                    // Store result in variable
                    let store_code = self.allocator.store_var(&var.name, &result_reg);
                    code.extend(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!(),
            Statement::Continue => unimplemented!(),

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

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

                code.push(format!("\tstw {}, {}", result_reg, ptr_reg));

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

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

                // Check if this is a global variable
                if self.is_global(varname) {
                    // Store to global label
                    code.push(format!("\tstw {}, {}", result_reg, varname));
                } else {
                    // Store result in local variable
                    let store_code = self.allocator.store_var(varname, &result_reg);
                    code.extend(store_code);
                }

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

            Statement::Return(expr) => {
                if let Some(e) = expr {
                    let (result_reg, expr_code) = self.generate_expression(e, true)?;
                    code.extend(expr_code);
                    code.push(format!("\tstw {}, bpr, 8", result_reg));
                    code.push(format!("\tjmp _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)?;
                code.extend(cond_code);

                // Compare with zero
                code.push(format!("\tcmp {}, zero", cond_reg));
                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(format!("\tjeq {}", else_label));

                // Then block
                code.push(format!("{}:", then_label));
                for s in then_stmt {
                    code.extend(self.generate_statement(s)?);
                }

                if then_stmt.len() == 0 {
                    code.push("\tnop".to_string());
                }

                code.push(format!("\tjmp {}", end_label));

                // Else block
                code.push(format!("{}:", else_label));
                for s in else_stmt {
                    code.extend(self.generate_statement(s)?);
                }

                if else_stmt.len() == 0 {
                    code.push("\tnop".to_string());
                }

                code.push(format!("{}:", 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(format!("{}:", loop_start));

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

                code.push(format!("\tcmp {}, zero", cond_reg));
                self.allocator.free_temp(&cond_reg);

                code.push(format!("\tjeq {}", loop_end));

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

                code.push(format!("\tjmp {}", loop_start));
                code.push(format!("{}:", loop_end));
            }

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

                code.push(format!("{}:", loop_start));

                for s in body {
                    code.extend(self.generate_statement(s)?);
                }

                code.push(format!("\tjmp {}", loop_start));
            }

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

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

        Ok(code)
    }

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

        // optimisation to prevent generating dead code!
        if expr.is_pure() && !use_result {
            return Ok((String::new(), code));
        }

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

                // write string into memory
                let uuid = self.get_unique_label();
                code.push(format!("\tdb str_{uuid}: \"{value}\""));

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

                Ok((reg, code))
            }

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

                // Load immediate value
                code.push(format!("\tlli {}, {} // '{value}'", *value as u8, reg));

                Ok((reg, code))
            }

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

                // Load immediate value
                code.push(format!("\tlli {}, {}", value & 0xFFFF, reg));
                if *value > 0xFFFF || *value < 0 {
                    code.push(format!("\tlui {}, {}", (value >> 16) & 0xFFFF, reg));
                }

                Ok((reg, code))
            }

            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.extend(alloc_code);

                    // Load from global label
                    code.push(format!("\tldw {}, {}", name.name, reg));

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

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

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

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

                // Generate operation
                match op {
                    BinaryOperator::Add => {
                        code.push(format!(
                            "\tadd {}, {}, {}",
                            left_reg, right_reg, result_reg
                        ));
                    }
                    BinaryOperator::Sub => {
                        code.push(format!(
                            "\tsub {}, {}, {}",
                            left_reg, right_reg, result_reg
                        ));
                    }
                    BinaryOperator::Mul => {
                        self.include("maths", "./lib/maths/core.dsa");
                        // Call multiply function
                        code.push(format!("\tpush {}", right_reg));
                        code.push(format!("\tpush {}", left_reg));
                        code.push("\tcall maths::multiply".to_string());
                        code.push(format!("\tpop {}", result_reg));
                        code.push("\tpop zero".to_string());
                    }
                    // Comparison operators - return 1 (true) or 0 (false)
                    BinaryOperator::Eq => {
                        code.push(format!("\tcmp {}, {}", left_reg, right_reg));
                        code.push(format!("\tlli 0, {}", result_reg));
                        let end_label = format!("_cmp_end_{}", self.get_unique_label());
                        code.push(format!("\tjne {}", end_label)); // If not equal, skip setting to 1
                        code.push(format!("\tlli 1, {}", result_reg));
                        code.push(format!("{}:", end_label));
                    }
                    BinaryOperator::Ne => {
                        code.push(format!("\tcmp {}, {}", left_reg, right_reg));
                        code.push(format!("\tlli 0, {}", result_reg));
                        let end_label = format!("_cmp_end_{}", self.get_unique_label());
                        code.push(format!("\tjeq {}", end_label)); // If equal, skip setting to 1
                        code.push(format!("\tlli 1, {}", result_reg));
                        code.push(format!("{}:", end_label));
                    }
                    BinaryOperator::Lt => {
                        code.push(format!("\tcmp {}, {}", left_reg, right_reg));
                        code.push(format!("\tlli 0, {}", result_reg));
                        let end_label = format!("_cmp_end_{}", self.get_unique_label());
                        code.push(format!("\tjge {}", end_label)); // If greater or equal, skip setting to 1
                        code.push(format!("\tlli 1, {}", result_reg));
                        code.push(format!("{}:", end_label));
                    }
                    BinaryOperator::Le => {
                        code.push(format!("\tcmp {}, {}", left_reg, right_reg));
                        code.push(format!("\tlli 0, {}", result_reg));
                        let end_label = format!("_cmp_end_{}", self.get_unique_label());
                        code.push(format!("\tjgt {}", end_label)); // If greater than, skip setting to 1
                        code.push(format!("\tlli 1, {}", result_reg));
                        code.push(format!("{}:", end_label));
                    }
                    BinaryOperator::Gt => {
                        code.push(format!("\tcmp {}, {}", left_reg, right_reg));
                        code.push(format!("\tlli 0, {}", result_reg));
                        let end_label = format!("_cmp_end_{}", self.get_unique_label());
                        code.push(format!("\tjle {}", end_label)); // If less or equal, skip setting to 1
                        code.push(format!("\tlli 1, {}", result_reg));
                        code.push(format!("{}:", end_label));
                    }
                    BinaryOperator::Ge => {
                        code.push(format!("\tcmp {}, {}", left_reg, right_reg));
                        code.push(format!("\tlli 0, {}", result_reg));
                        let end_label = format!("_cmp_end_{}", self.get_unique_label());
                        code.push(format!("\tjlt {}", end_label)); // If less than, skip setting to 1
                        code.push(format!("\tlli 1, {}", result_reg));
                        code.push(format!("{}:", 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::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)?;
                    code.extend(arg_code);
                    arg_regs.push(arg_reg);
                }

                // Save caller-saved registers and track which ones we saved
                // old method, inefficient.
                // let saved_regs = self.allocator.get_caller_saved_registers();
                // for reg in &saved_regs {
                //     code.push(format!("\tpush {}", 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.extend(self.allocator.spill_register(reg).unwrap());
                }

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

                // if GLOBAL_METHODS.contains_key(name.name.as_str()) {
                //     code.push(format!("\tcall {}",
                // GLOBAL_METHODS[name.name.as_str()])); } else
                if self.symbols.contains(&name.name) {
                    // Call local function
                    code.push(format!("\tcall {}", name));
                } else if let Some(ns) = name.namespace.clone()
                    && self.imports.contains_key(&ns)
                {
                    code.push(format!("\tcall {}", name));
                } else {
                    return Err(CompilerError::Undefined(name.clone()));
                }

                let result_reg: String;

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

                    code.extend(result_alloc);
                    code.push(format!("\tpop {}", result_reg));

                    // Clean up arguments
                    if args.len() > 1 {
                        for _ in 0..(args.len() - 1) {
                            code.push("\tpop zero".to_string());
                        }
                    }
                } else {
                    result_reg = "zero".to_string();

                    // Clean up arguments
                    if args.len() > 0 {
                        for _ in 0..(args.len()) {
                            code.push("\tpop zero".to_string());
                        }
                    }
                }

                // Restore caller-saved registers in reverse order (LIFO)
                // for reg in saved_regs.iter().rev() {
                //     code.push(format!("\tpop {}", reg));
                // }

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

                Ok((result_reg, code))
            }

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

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

                match op {
                    UnaryOperator::Minus => {
                        // Negate: result = 0 - operand
                        code.push(format!("\tsub zero, {}, {}", operand_reg, result_reg));
                    }
                    UnaryOperator::Plus => {
                        // Just move
                        code.push(format!("\tmov {}, {}", operand_reg, result_reg));
                    }
                    UnaryOperator::Dereference => {
                        code.push(format!("\tldw {}, {}", operand_reg, result_reg));
                    }
                    UnaryOperator::Reference => {
                        code.extend(self.allocator.spill_register(&operand_reg)?);
                        code.push(format!(
                            "\tsubi bpr {} {}",
                            -(4 + self.allocator.get_stack_offset()),
                            result_reg
                        ))
                    }
                }

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

            Expression::Empty => Ok(("zero".to_string(), 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()
    }
}

/// Build a single string from any number of arguments.
/// Each argument must implement `Display` or be convertible to a string.
#[macro_export]
macro_rules! dsa {
    ($($arg:expr),* $(,)?) => {{
        // Start with an empty String – we’ll grow it as we go.
        use std::fmt::Write;
        let mut s = ::std::string::String::new();
        $(
            // `write!` is cheaper than `format!` for each element
            // because it re‑uses the same buffer.

            write!(s, "{}\n", $arg).expect("write to String failed");
        )*
        s
    }};
}

// ────────────────────────  dsa!  ────────────────────────
// A tiny helper that just turns its token‑stream into a string.
// The trailing comma is kept – it’s part of the syntax you want.
#[macro_export]
macro_rules! cmd {
    ($($tokens:tt)*) => {{
        // We’ll just stringify the tokens and return a String.
        format!("{}", concat!(stringify!($tokens), "\n"))
    }};
}

// ────────────────────────  block!  ────────────────────────
// Usage:
//
//   let asm = block![ "name"
//       dsa![mov rg0, rg1],
//       dsa![add rg1, rg1]
//   ];
//
// `asm` is a `&'static str` containing:
//
//   name:
//       mov rg0, rg1
//       add rg1, rg1
//
#[macro_export]
macro_rules! block {
    // The first token must be a string literal – that’s the label.
    ($label:literal $(dsa![$($ins:tt)*]),* ) => {{
        // Build a single string at compile time.
        const CODE: &str = concat!(
            $label, ":\n",
            // Each `dsa!` call yields a string like `"mov rg0, rg1"`.
            // We add a newline after each one to get the desired layout.
            $(concat!("\t", stringify!($($ins)*), "\n")),*
        );
        CODE
    }};
}

#[macro_export]
macro_rules! comment {
    ($text:expr) => {{ format!("// {}", $text) }};
}