use std::{fmt, str::FromStr};

use common::prelude::Register;

use crate::AssembleError;

#[derive(Debug, Clone)]
pub struct Node {
    pub symbol: Option<Symbol>,
    pub opcode: Opcode,
    pub tokens: Vec<Token>,
}

#[macro_export]
#[macro_use]
macro_rules! node {
    ($symbol: expr, $opcode: expr, args: $tokens: expr) => {
        Node::new($symbol.clone(), $opcode.clone(), $tokens.clone())
    };

    ($symbol: expr, $opcode: expr, $($tokens: expr),+) => {
        Node::new($symbol.clone(), $opcode.clone(), vec![$($tokens.clone()),+])
    };
}

impl Node {
    pub fn new(symbol: Option<Symbol>, opcode: Opcode, tokens: Vec<Token>) -> Node {
        Node {
            symbol,
            opcode,
            tokens,
        }
    }

    pub fn label(&self) -> Option<Symbol> {
        self.symbol.clone()
    }

    pub fn opcode(&self) -> Opcode {
        self.opcode.clone()
    }

    pub fn args(&self) -> Vec<Token> {
        self.tokens.clone()
    }

    pub fn arg(&self, index: usize) -> Result<Token, AssembleError> {
        self.args()
            .get(index)
            .cloned()
            .ok_or(AssembleError::InvalidArg)
    }
}

impl fmt::Display for Node {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let symbol = match &self.label() {
            Some(symbol) => format!("{}:\n", symbol),
            None => "".to_string(),
        };

        write!(
            f,
            "\x1b[93m{} \t\x1b[94m{} \x1b[37m{:?} \x1b[0m",
            symbol,
            self.opcode(),
            self.args()
        )
    }
}

impl fmt::Display for Symbol {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} ( module: {})", self.name, self.module)
    }
}

impl fmt::Display for Module {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Module::Unresolved(name) => write!(f, "{name}"),
            Module::Resolved(name) => write!(f, "{name}"),
        }
    }
}

impl fmt::Display for Opcode {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Opcode::Nop => write!(f, "nop"),
            Opcode::Mov => write!(f, "mov"),
            Opcode::Movs => write!(f, "movs"),
            Opcode::Ldb => write!(f, "ldb"),
            Opcode::Ldbs => write!(f, "ldbs"),
            Opcode::Ldh => write!(f, "ldh"),
            Opcode::Ldhs => write!(f, "ldhs"),
            Opcode::Ldw => write!(f, "ldw"),
            Opcode::Stb => write!(f, "stb"),
            Opcode::Sth => write!(f, "sth"),
            Opcode::Stw => write!(f, "stw"),
            Opcode::Lli => write!(f, "lli"),
            Opcode::Lui => write!(f, "lui"),
            Opcode::Jmp => write!(f, "jmp"),
            Opcode::Jeq => write!(f, "jeq"),
            Opcode::Jne => write!(f, "jne"),
            Opcode::Jgt => write!(f, "jgt"),
            Opcode::Jge => write!(f, "jge"),
            Opcode::Jlt => write!(f, "jlt"),
            Opcode::Jle => write!(f, "jle"),
            Opcode::Cmp => write!(f, "cmp"),
            Opcode::Inc => write!(f, "inc"),
            Opcode::Dec => write!(f, "dec"),
            Opcode::Shl => write!(f, "shl"),
            Opcode::Shr => write!(f, "shr"),
            Opcode::Add => write!(f, "add"),
            Opcode::Sub => write!(f, "sub"),
            Opcode::And => write!(f, "and"),
            Opcode::Or => write!(f, "or"),
            Opcode::Not => write!(f, "not"),
            Opcode::Xor => write!(f, "xor"),
            Opcode::Nand => write!(f, "nand"),
            Opcode::Nor => write!(f, "nor"),
            Opcode::Xnor => write!(f, "xnor"),
            Opcode::Int => write!(f, "int"),
            Opcode::Irt => write!(f, "irt"),
            Opcode::Hlt => write!(f, "hlt"),
            Opcode::Iadd => write!(f, "iadd"),
            Opcode::Isub => write!(f, "isub"),
            Opcode::Db => write!(f, "db"),
            Opcode::Dh => write!(f, "dh"),
            Opcode::Dw => write!(f, "dw"),
            Opcode::Resb => write!(f, "resb"),
            Opcode::Resh => write!(f, "resh"),
            Opcode::Resw => write!(f, "resw"),
            Opcode::Push => write!(f, "push"),
            Opcode::Pop => write!(f, "pop"),
            Opcode::Lwi => write!(f, "lwi"),

            // utility - removed at compile time
            Opcode::Include => write!(f, "include"),

            // special - generated by assembler
            Opcode::Data => write!(f, "data"),
        }
    }
}

#[derive(Debug, Clone, Eq, Hash)]
pub struct Symbol {
    pub name: String,
    pub module: Module,
}

impl PartialEq for Symbol {
    fn eq(&self, other: &Self) -> bool {
        self.name == other.name && self.module == other.module
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Module {
    Resolved(u64),
    Unresolved(String),
}

#[derive(Debug, Clone)]
pub enum Token {
    Symbol(Symbol),
    Register(Register),
    Immediate(u32),
    StringLit(String),
    Opcode(Opcode),
}

#[derive(Debug, PartialEq, Copy, Clone)]
pub enum TokenType {
    Symbol,
    Register,
    Immediate,
    StringLit,
    Opcode,
}

impl TokenType {
    pub fn from_token(token: &Token) -> TokenType {
        match token {
            Token::Symbol(_) => TokenType::Symbol,
            Token::Register(_) => TokenType::Register,
            Token::Immediate(_) => TokenType::Immediate,
            Token::StringLit(_) => TokenType::StringLit,
            Token::Opcode(_) => TokenType::Opcode,
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum Opcode {
    // Real instructions (0x00-0x26)
    Nop,
    Mov,
    Movs,
    Ldb,
    Ldbs,
    Ldh,
    Ldhs,
    Ldw,
    Stb,
    Sth,
    Stw,
    Lli,
    Lui,
    Jmp,
    Jeq,
    Jne,
    Jgt,
    Jge,
    Jlt,
    Jle,
    Cmp,
    Inc,
    Dec,
    Shl,
    Shr,
    Add,
    Sub,
    And,
    Or,
    Not,
    Xor,
    Nand,
    Nor,
    Xnor,
    Int,
    Irt,
    Hlt,
    Iadd,
    Isub,
    // Pseudo-instructions
    Db,
    Dh,
    Dw,
    Resb,
    Resh,
    Resw,
    Push,
    Pop,
    Lwi,
    Include,

    // fake instructions (these aren't present in the binary as instructions)
    Data,
}

#[derive(Debug)]
pub enum OpcodeFromStrError {
    InvalidRegister(&'static str),
}

impl std::fmt::Display for OpcodeFromStrError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::InvalidRegister(reg) => write!(f, "register does not exist: {reg}"),
        }
    }
}

impl std::error::Error for OpcodeFromStrError {}

impl FromStr for Opcode {
    type Err = OpcodeFromStrError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "nop" => Ok(Self::Nop),
            "mov" => Ok(Self::Mov),
            "movs" => Ok(Self::Movs),
            "ldb" => Ok(Self::Ldb),
            "ldbs" => Ok(Self::Ldbs),
            "ldh" => Ok(Self::Ldh),
            "ldhs" => Ok(Self::Ldhs),
            "ldw" => Ok(Self::Ldw),
            "stb" => Ok(Self::Stb),
            "sth" => Ok(Self::Sth),
            "stw" => Ok(Self::Stw),
            "lli" => Ok(Self::Lli),
            "lui" => Ok(Self::Lui),
            "jmp" => Ok(Self::Jmp),
            "jeq" => Ok(Self::Jeq),
            "jne" => Ok(Self::Jne),
            "jgt" => Ok(Self::Jgt),
            "jge" => Ok(Self::Jge),
            "jlt" => Ok(Self::Jlt),
            "jle" => Ok(Self::Jle),
            "cmp" => Ok(Self::Cmp),
            "inc" => Ok(Self::Inc),
            "dec" => Ok(Self::Dec),
            "shl" => Ok(Self::Shl),
            "shr" => Ok(Self::Shr),
            "add" => Ok(Self::Add),
            "sub" => Ok(Self::Sub),
            "and" => Ok(Self::And),
            "or" => Ok(Self::Or),
            "not" => Ok(Self::Not),
            "xor" => Ok(Self::Xor),
            "nand" => Ok(Self::Nand),
            "nor" => Ok(Self::Nor),
            "xnor" => Ok(Self::Xnor),
            "int" => Ok(Self::Int),
            "irt" => Ok(Self::Irt),
            "hlt" => Ok(Self::Hlt),
            "iadd" => Ok(Self::Iadd),
            "isub" => Ok(Self::Isub),
            "db" => Ok(Self::Db),
            "dh" => Ok(Self::Dh),
            "dw" => Ok(Self::Dw),
            "resb" => Ok(Self::Resb),
            "resh" => Ok(Self::Resh),
            "resw" => Ok(Self::Resw),
            "push" => Ok(Self::Push),
            "pop" => Ok(Self::Pop),
            "lwi" => Ok(Self::Lwi),
            "include" => Ok(Self::Include),
            _ => Err(OpcodeFromStrError::InvalidRegister("unknown opcode")),
        }
    }
}

impl Opcode {
    pub const OPCODES: &[&str] = &[
        // Real instructions (0x00-0x26)
        "nop", "mov", "movs", "ldb", "ldbs", "ldh", "ldhs", "ldw", "stb", "sth", "stw",
        "lli", "lui", "jmp", "jeq", "jne", "jgt", "jge", "jlt", "jle", "cmp", "inc",
        "dec", "shl", "shr", "add", "sub", "and", "or", "not", "xor", "nand", "nor",
        "xnor", "int", "irt", "hlt", "iadd", "isub", // Pseudo-instructions
        "db", "dh", "dw", "resb", "resh", "resw", "push", "pop", "lwi", "include",
    ];

    pub fn to_opcode_value(&self) -> Option<u8> {
        match self {
            Self::Nop => Some(0x00),
            Self::Mov => Some(0x01),
            Self::Movs => Some(0x02),
            Self::Ldb => Some(0x03),
            Self::Ldbs => Some(0x04),
            Self::Ldh => Some(0x05),
            Self::Ldhs => Some(0x06),
            Self::Ldw => Some(0x07),
            Self::Stb => Some(0x08),
            Self::Sth => Some(0x09),
            Self::Stw => Some(0x0A),
            Self::Lli => Some(0x0B),
            Self::Lui => Some(0x0C),
            Self::Jmp => Some(0x0D),
            Self::Jeq => Some(0x0E),
            Self::Jne => Some(0x0F),
            Self::Jgt => Some(0x10),
            Self::Jge => Some(0x11),
            Self::Jlt => Some(0x12),
            Self::Jle => Some(0x13),
            Self::Cmp => Some(0x14),
            Self::Inc => Some(0x15),
            Self::Dec => Some(0x16),
            Self::Shl => Some(0x17),
            Self::Shr => Some(0x18),
            Self::Add => Some(0x19),
            Self::Sub => Some(0x1A),
            Self::And => Some(0x1B),
            Self::Or => Some(0x1C),
            Self::Not => Some(0x1D),
            Self::Xor => Some(0x1E),
            Self::Nand => Some(0x1F),
            Self::Nor => Some(0x20),
            Self::Xnor => Some(0x21),
            Self::Int => Some(0x22),
            Self::Irt => Some(0x23),
            Self::Hlt => Some(0x24),
            Self::Iadd => Some(0x25),
            Self::Isub => Some(0x26),
            // Pseudo-instructions don't have opcode values
            _ => None,
        }
    }

    pub fn is_pseudo_instruction(&self) -> bool {
        matches!(
            self,
            Self::Db
                | Self::Dh
                | Self::Dw
                | Self::Resb
                | Self::Resh
                | Self::Resw
                | Self::Push
                | Self::Pop
                | Self::Lwi
        )
    }
}