use common::prelude::Register;

use crate::assembler::model::{Node, Opcode, Token};
use crate::{assembler::AssembleError, expect_token, expect_type, node};

pub fn expand_pseudo_ops(
    mut nodes: Vec<Node>,
    module: u64,
) -> Result<Vec<Node>, AssembleError> {
    let mut result = Vec::<Node>::with_capacity(nodes.len());

    for node in &mut nodes {
        if try_expand(node.clone(), &mut result, module).is_err() {
            result.push(node.clone());
        }
    }

    Ok(result)
}

fn try_expand(
    node: Node,
    result: &mut Vec<Node>,
    _module: u64,
) -> Result<(), AssembleError> {
    match node.opcode() {
        Opcode::Push => expand_push(&node, result)?,
        Opcode::Pop => expand_pop(&node, result)?,
        Opcode::Pusha => expand_pusha(&node, result)?,
        Opcode::Popa => expand_popa(&node, result)?,
        Opcode::Call => expand_call(&node, result)?,
        Opcode::Return => expand_return(&node, result),
        Opcode::Ldb | Opcode::Ldbs | Opcode::Ldh | Opcode::Ldhs | Opcode::Ldw => {
            expand_ldx(&node, result)?;
        }
        Opcode::Stb | Opcode::Sth | Opcode::Stw => expand_stx(&node, result)?,

        Opcode::Lwi => expand_lwi(&node, result)?,
        Opcode::Resb | Opcode::Resh | Opcode::Resw => expand_resx(&node, result)?,
        Opcode::Db | Opcode::Dh | Opcode::Dw => expand_dx(&node, result)?,
        _ => result.push(node),
    }
    Ok(())
}

fn expand_push(current: &Node, nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let label = current.label();
    let Ok(arg0) = current.arg(0) else {
        return Err(AssembleError::Generic);
    };

    let reg = expect_type!(arg0, Register)?;
    let spr = Token::Register(Register::Spr);

    nodes.extend(vec![
        node!(label, Opcode::SubI, spr, 4, spr),
        node!(None, Opcode::Stw, reg, spr, 0),
    ]);

    Ok(())
}

fn expand_pusha(current: &Node, nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let label = current.label();
    let Ok(arg0) = current.arg(0) else {
        return Err(AssembleError::Generic);
    };

    let count = expect_token!(arg0, Immediate)?;
    let spr = Token::Register(Register::Spr);
    let registers: Vec<Register> = Register::general();

    nodes.push(node!(
        label,
        Opcode::SubI,
        spr,
        Token::Immediate(count * 4),
        spr
    ));

    nodes.extend((0..count).rev().map(|i| {
        node!(
            None,
            Opcode::Stw,
            Token::Register(registers[i as usize]),
            spr,
            Token::Immediate(i * 4)
        )
    }));

    Ok(())
}

fn expand_popa(current: &Node, nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let label = current.label();

    let Ok(arg0) = current.arg(0) else {
        return Err(AssembleError::Generic);
    };

    let count = expect_token!(arg0, Immediate)?;
    let spr = Token::Register(Register::Spr);
    let registers: Vec<Register> = Register::general();

    nodes.extend((0..count).rev().map(|i| {
        node!(
            { if i == 0 { label.clone() } else { None } },
            Opcode::Ldw,
            spr,
            Token::Register(registers[i as usize]),
            Token::Immediate(i * 4)
        )
    }));

    nodes.push(node!(
        None,
        Opcode::AddI,
        spr,
        Token::Immediate(count * 4),
        spr
    ));

    Ok(())
}

fn expand_call(current: &Node, nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let label = current.label();

    let Ok(arg0) = current.arg(0) else {
        return Err(AssembleError::Generic);
    };

    let addr = expect_type!(arg0, Symbol)?;
    let spr = Token::Register(Register::Spr);
    let pcx = Token::Register(Register::Pcx);
    let zero = Token::Register(Register::Zero);

    nodes.extend(vec![
        node!(label, Opcode::SubI, spr, 4, spr),
        node!(None, Opcode::Stw, pcx, spr, 0),
        node!(None, Opcode::Jmp, addr, zero),
    ]);

    Ok(())
}

fn expand_return(current: &Node, nodes: &mut Vec<Node>) {
    let label = current.label();
    let spr = Token::Register(Register::Spr);
    let ret = Token::Register(Register::Ret);

    nodes.extend(vec![
        node!(label, Opcode::Ldw, spr, ret, 0),
        node!(None, Opcode::AddI, spr, 4, spr),
        node!(None, Opcode::Jmp, 4, ret),
    ]);
}

fn expand_pop(current: &Node, nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let label = current.label();

    let Ok(arg0) = current.arg(0) else {
        return Err(AssembleError::Generic);
    };

    let reg = expect_type!(arg0, Register)?;
    let spr = Token::Register(Register::Spr);

    nodes.extend(vec![
        node!(label, Opcode::Ldw, spr, reg, 0),
        node!(None, Opcode::AddI, spr, 4, spr),
    ]);

    Ok(())
}

fn expand_ldx(current: &Node, nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let opcode = current.opcode();
    let args: Vec<Token> = current.args().into_iter().take(3).collect();

    let Some(name) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };

    let Some(reg) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };

    let Some(offset) = args.get(2) else {
        return Err(AssembleError::MissingArgument(2));
    };

    let name = expect_type!(name, Symbol)?;
    let reg = expect_type!(reg, Register)?;
    let offset = expect_type!(offset, Immediate)?;

    nodes.extend(vec![
        node!(current.label(), Opcode::Lli, name, reg),
        node!(None, Opcode::Lui, name, reg),
        node!(None, opcode, reg, reg, offset),
    ]);

    Ok(())
}

fn expand_stx(current: &Node, nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let opcode = current.opcode();

    let args: Vec<Token> = current.args().into_iter().take(3).collect();

    let Some(base) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };

    let Some(dest) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };

    let Some(offset) = args.get(2) else {
        return Err(AssembleError::MissingArgument(2));
    };

    let base = expect_type!(base, Register)?;
    let dest = expect_type!(dest, Symbol)?;
    let offset = expect_type!(offset, Immediate)?;
    let temp = Token::Register(Register::Acc);

    nodes.extend(vec![
        node!(current.label(), Opcode::Lli, dest, temp),
        node!(None, Opcode::Lui, dest, temp),
        node!(None, opcode, base, temp, offset),
    ]);

    Ok(())
}

fn expand_lwi(current: &Node, nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let Ok(val) = current.arg(0) else {
        return Err(AssembleError::MissingArgument(0));
    };

    let Ok(reg) = current.arg(1) else {
        return Err(AssembleError::MissingArgument(1));
    };

    let val = expect_type!(val, Symbol, Immediate)?;
    let reg = expect_type!(reg, Register)?;

    nodes.extend(vec![
        node!(current.label(), Opcode::Lli, val, reg),
        node!(None, Opcode::Lui, val, reg),
    ]);

    Ok(())
}

fn expand_resx(current: &Node, nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let Ok(region_label) = current.arg(0) else {
        return Err(AssembleError::MissingArgument(0));
    };

    let Ok(size) = current.arg(1) else {
        return Err(AssembleError::MissingArgument(1));
    };

    let region_label = expect_token!(region_label, Symbol)?;
    let size = expect_token!(size, Immediate)?;

    let units_per = match current.opcode() {
        Opcode::Resb => 4,
        Opcode::Resh => 2,
        Opcode::Resw => 1,
        _ => unreachable!(),
    };

    let mut buffer = vec![];
    // push the inital node with the label
    for _ in 0..size.div_ceil(units_per) {
        // push the rest of the nodes
        buffer.push(node!(None, Opcode::Data, 0));
    }
    buffer[0].symbol = Some(region_label);
    nodes.extend(buffer);

    Ok(())
}

fn expand_dx(current: &Node, nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let Ok(region_label) = current.arg(0) else {
        return Err(AssembleError::MissingArgument(0));
    };

    let region_label = expect_token!(region_label, Symbol)?;
    let size = match current.opcode() {
        Opcode::Db => 4,
        Opcode::Dh => 2,
        Opcode::Dw => 1,
        _ => unreachable!(),
    };

    let mut buffer = vec![];

    let mut args = current.args();
    let _label = args.remove(0);

    for word in process_dx_data(args, size)? {
        buffer.push(node!(None, Opcode::Data, Token::Immediate(word)));
    }
    buffer[0].symbol = Some(region_label);

    nodes.extend(buffer);
    Ok(())
}

fn process_dx_data(args: Vec<Token>, size: usize) -> Result<Vec<u32>, AssembleError> {
    assert!(matches!(size, 1 | 2 | 4));

    let mut buffer = Vec::<u8>::new();

    // Process each token
    for token in args {
        match token {
            Token::StringLit(mut s) => {
                s.push('\0');
                // Split string into chars and write as bytes
                for ch in s.chars() {
                    // Convert char to bytes (UTF-8 encoding)
                    let mut char_buf = [0u8; 4];
                    let char_bytes = ch.encode_utf8(&mut char_buf);
                    buffer.extend_from_slice(char_bytes.as_bytes());
                }
            }
            Token::Immediate(value) => {
                // Split u32 into bytes (little-endian)
                buffer.extend_from_slice(&value.to_be_bytes());
            }
            _ => {
                return Err(AssembleError::Generic);
            }
        }

        // Pad buffer to alignment boundary with zeros
        let remainder = buffer.len() % size;
        if remainder != 0 {
            let padding = size - remainder;
            buffer.resize(buffer.len() + padding, 0);
        }
    }

    // Convert byte buffer to u32 chunks
    // Pad final buffer to u32 boundary if needed
    let remainder = buffer.len() % 4;
    if remainder != 0 {
        let padding = 4 - remainder;
        buffer.resize(buffer.len() + padding, 0);
    }

    // Convert bytes to u32s efficiently using chunks_exact
    let result = buffer
        .chunks_exact(4)
        .map(|chunk| {
            // Convert 4 bytes to u32 (little-endian)
            u32::from_be_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])
        })
        .collect();

    Ok(result)
}
