use common::prelude::*;

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

fn log(message: &str) {
    println!("\x1b[32mINFO:\x1b[0m {message}");
}

pub fn codegen(nodes: Vec<Node>) -> Result<Vec<Instruction>, AssembleError> {
    let mut instructions = vec![];

    for node in nodes {
        println!("{:?}", node);
        instructions.push(build_instruction(&node)?);
    }

    log("Assembly Successful ✅");

    Ok(instructions)
}

fn build_instruction(node: &Node) -> Result<Instruction, AssembleError> {
    let opcode = node.opcode();
    let args = node.args();

    match opcode {
        Opcode::Nop => Ok(Instruction::nop()),
        Opcode::Mov => ins_mov(opcode, &args),
        Opcode::CMov => ins_cmov(opcode, &args),
        Opcode::Ldb
        | Opcode::Ldw
        | Opcode::Ldh
        | Opcode::Ldbs
        | Opcode::Ldhs
        | Opcode::Stb
        | Opcode::Stw
        | Opcode::Sth => ins_ldx_stx(opcode, &args),
        Opcode::Lli | Opcode::Lui => ins_load_imm(opcode, &args),
        Opcode::Push | Opcode::Pop => ins_stack(opcode, &args),
        Opcode::Ieq | Opcode::Ine | Opcode::Igt | Opcode::Ige | Opcode::Ile | Opcode::Ilt => {
            ins_comparison(opcode, &args)
        }
        Opcode::Jmp | Opcode::Call | Opcode::Jic | Opcode::Jnc => {
            ins_jump_unconditional(opcode, &args)
        }
        Opcode::Jez | Opcode::Jnz => ins_jump_conditional(opcode, &args),
        Opcode::Shl | Opcode::Shr => ins_bitshift(opcode, &args),
        Opcode::Add
        | Opcode::Sub
        | Opcode::And
        | Opcode::Or
        | Opcode::Xor
        | Opcode::Nand
        | Opcode::Nor
        | Opcode::Xnor => ins_arithmetic(opcode, &args),
        Opcode::AddI | Opcode::SubI => ins_imm_arithmetic(opcode, &args),
        Opcode::Not => ins_not(&args),
        Opcode::Int => ins_interrupt(&args),
        Opcode::Ret => Ok(Instruction::ret()),
        Opcode::IRet => Ok(Instruction::irt()),
        Opcode::Hlt => Ok(Instruction::hlt()),
        Opcode::Data => build_data_instruction(&args),
        Opcode::Segment => build_segment_instruction(&args),
        Opcode::Db
        | Opcode::Dh
        | Opcode::Dw
        | Opcode::Resb
        | Opcode::Resh
        | Opcode::Resw
        | Opcode::Lwi
        | Opcode::Include
        | Opcode::Func
        | Opcode::Return => Err(AssembleError::InvalidArg),
    }
}

fn ins_mov(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(src_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(dest_token) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };

    let src = expect_token!(src_token, Register)?;
    let dest = expect_token!(dest_token, Register)?;

    match opcode {
        Opcode::Mov => Ok(Instruction::mov(src, dest)),
        _ => unreachable!(),
    }
}

fn ins_cmov(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(src_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(dest_token) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };
    let Some(cmp_token) = args.get(2) else {
        return Err(AssembleError::MissingArgument(1));
    };

    let src = expect_token!(src_token, Register)?;
    let dest = expect_token!(dest_token, Register)?;
    let cmp = expect_token!(cmp_token, Register)?;

    match opcode {
        Opcode::CMov => Ok(Instruction::cmov(src, dest, cmp)),
        _ => unreachable!(),
    }
}

fn ins_ldx_stx(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(src_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(dest_token) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };
    let Some(offset_token) = args.get(2) else {
        return Err(AssembleError::MissingArgument(2));
    };

    let src = expect_token!(src_token, Register)?;
    let dest = expect_token!(dest_token, Register)?;
    let offset = expect_token!(offset_token, Immediate)?;

    match opcode {
        Opcode::Ldb => Ok(Instruction::ldb(src, dest, offset as u16)),
        Opcode::Ldbs => Ok(Instruction::ldbs(src, dest, offset as u16)),
        Opcode::Ldh => Ok(Instruction::ldh(src, dest, offset as u16)),
        Opcode::Ldhs => Ok(Instruction::ldhs(src, dest, offset as u16)),
        Opcode::Ldw => Ok(Instruction::ldw(src, dest, offset as u16)),
        Opcode::Stb => Ok(Instruction::stb(src, dest, offset as u16)),
        Opcode::Sth => Ok(Instruction::sth(src, dest, offset as u16)),
        Opcode::Stw => Ok(Instruction::stw(src, dest, offset as u16)),
        _ => unreachable!(),
    }
}

fn ins_stack(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(reg_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };

    let reg = expect_token!(reg_token, Register)?;
    match opcode {
        Opcode::Push => Ok(Instruction::push(reg)),
        Opcode::Pop => Ok(Instruction::pop(reg)),
        _ => unreachable!(),
    }
}

fn ins_load_imm(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(value_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(dest_token) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };

    let value = expect_token!(value_token, Immediate)?;
    let dest = expect_token!(dest_token, Register)?;

    match opcode {
        Opcode::Lli => Ok(Instruction::lli(dest, value as u16)),
        Opcode::Lui => Ok(Instruction::lui(dest, (value >> 16) as u16)),
        _ => unreachable!(),
    }
}

fn ins_jump_unconditional(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(addr_imm_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    // addr_reg is optional; default to Register::Zero if missing
    let addr_reg_token = args.get(1).unwrap_or(&Token::Register(Register::Zero));

    let addr_imm = expect_token!(addr_imm_token, Immediate)?;
    let addr_reg = expect_token!(addr_reg_token, Register)?;

    match opcode {
        Opcode::Jmp => Ok(Instruction::jmp(addr_reg, addr_imm as u16)),
        Opcode::Jic => Ok(Instruction::jic(addr_reg, addr_imm as u16)),
        Opcode::Jnc => Ok(Instruction::jnc(addr_reg, addr_imm as u16)),
        Opcode::Call => Ok(Instruction::call(addr_reg, addr_imm as u16)),
        _ => unreachable!(),
    }
}

fn ins_jump_conditional(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(condition_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(addr_imm_token) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };
    let addr_reg_token = if let Some(token) = args.get(2) {
        token
    } else {
        &Token::Register(Register::Zero)
    };

    let condition = expect_token!(condition_token, Register)?;
    let addr_imm = expect_token!(addr_imm_token, Immediate)?;
    let addr_reg = expect_token!(addr_reg_token, Register)?;

    Ok(match opcode {
        Opcode::Jez => Instruction::jez(condition, addr_reg, addr_imm as u16),
        Opcode::Jnz => Instruction::jnz(condition, addr_reg, addr_imm as u16),
        _ => unreachable!(),
    })
}

fn ins_comparison(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(left_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(right_token) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };
    let Some(dest_token) = args.get(2) else {
        return Err(AssembleError::MissingArgument(2));
    };

    let left = expect_token!(left_token, Register)?;
    let right = expect_token!(right_token, Register)?;
    let dest = expect_token!(dest_token, Register)?;
    Ok(match opcode {
        Opcode::Ieq => Instruction::ieq(left, right, dest),
        Opcode::Ine => Instruction::ine(left, right, dest),
        Opcode::Igt => Instruction::igt(left, right, dest),
        Opcode::Ige => Instruction::ige(left, right, dest),
        Opcode::Ile => Instruction::ile(left, right, dest),
        Opcode::Ilt => Instruction::ilt(left, right, dest),
        _ => unreachable!(),
    })
}

fn ins_bitshift(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(src_reg) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(r_shamt) = args.get(1) else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(i_shamt) = args.get(2) else {
        return Err(AssembleError::MissingArgument(1));
    };
    let Some(dest_reg) = args.get(3) else {
        return Err(AssembleError::MissingArgument(1));
    };

    let src = expect_token!(src_reg, Register)?;
    let r_shamt = expect_token!(r_shamt, Register)?;
    let i_shamt = expect_token!(i_shamt, Immediate)? as u8;
    let dest = expect_token!(dest_reg, Register)?;

    Ok(match opcode {
        Opcode::Shl => Instruction::shl(src, r_shamt, dest, i_shamt),
        Opcode::Shr => Instruction::shr(src, r_shamt, dest, i_shamt),
        _ => unreachable!(),
    })
}

fn ins_arithmetic(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(left_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(right_token) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };
    let Some(dest_token) = args.get(2) else {
        return Err(AssembleError::MissingArgument(2));
    };

    let left = expect_token!(left_token, Register)?;
    let right = expect_token!(right_token, Register)?;
    let dest = expect_token!(dest_token, Register)?;

    Ok(match opcode {
        Opcode::Add => Instruction::add(left, right, dest),
        Opcode::Sub => Instruction::sub(left, right, dest),
        Opcode::And => Instruction::and(left, right, dest),
        Opcode::Or => Instruction::or(left, right, dest),
        Opcode::Xor => Instruction::xor(left, right, dest),
        Opcode::Nand => Instruction::nand(left, right, dest),
        Opcode::Nor => Instruction::nor(left, right, dest),
        Opcode::Xnor => Instruction::xnor(left, right, dest),
        _ => unreachable!(),
    })
}

fn ins_imm_arithmetic(
    opcode: Opcode,
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(reg_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(immediate_token) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };
    let Some(dest_token) = args.get(2) else {
        return Err(AssembleError::MissingArgument(2));
    };

    let reg = expect_token!(reg_token, Register)?;
    let immediate = expect_token!(immediate_token, Immediate)? as u16;
    let dest = expect_token!(dest_token, Register)?;

    Ok(match opcode {
        Opcode::AddI => Instruction::addi(reg, dest, immediate),
        Opcode::SubI => Instruction::subi(reg, dest, immediate),
        _ => unreachable!(),
    })
}

fn ins_not(args: &[crate::assembler::model::Token]) -> Result<Instruction, AssembleError> {
    let Some(reg_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };
    let Some(dest_token) = args.get(1) else {
        return Err(AssembleError::MissingArgument(1));
    };

    let src = expect_token!(reg_token, Register)?;
    let dest = expect_token!(dest_token, Register)?;
    Ok(Instruction::not(src, dest))
}

fn ins_interrupt(args: &[crate::assembler::model::Token]) -> Result<Instruction, AssembleError> {
    let Some(code_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };

    let code = expect_token!(code_token, Immediate)? as u8;
    Ok(Instruction::int(code))
}

fn build_data_instruction(
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(immediate_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };

    let immediate = expect_token!(immediate_token, Immediate)?;
    Ok(Instruction::data(immediate))
}

fn build_segment_instruction(
    args: &[crate::assembler::model::Token],
) -> Result<Instruction, AssembleError> {
    let Some(immediate_token) = args.first() else {
        return Err(AssembleError::MissingArgument(0));
    };

    let immediate = expect_token!(immediate_token, Immediate)?;
    Ok(Instruction::data(
        // mask the first 6 bits for the opcode
        // this is really deprecated tbh. may remove in future.
        Opcode::Segment as u32 | (immediate & 0x05FFFFFF),
    ))
}