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

use common::prelude::Register;

use crate::assembler::AssembleError;

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

impl Node {
    #[must_use]
    pub const fn new(symbol: Option<Symbol>, opcode: Opcode, tokens: Vec<Token>) -> Self {
        Self {
            symbol,
            opcode,
            tokens,
        }
    }

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

    #[must_use]
    pub const fn opcode(&self) -> Opcode {
        self.opcode
    }

    #[must_use]
    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 = self
            .label()
            .as_ref()
            .map_or_else(String::new, |symbol| format!("{symbol}:\n"));

        let args = self
            .args()
            .into_iter()
            .map(|arg| arg.to_string())
            .collect::<Vec<_>>()
            .join(" ");

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

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

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

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

            // meta instructions
            Self::Include => write!(f, "include"),
            Self::Data => write!(f, "data"),
            Self::Segment => write!(f, "[SEGMENT]"),
        }
    }
}

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

impl std::hash::Hash for Symbol {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.name.hash(state);
        self.module.hash(state);
    }
}

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),
}

impl fmt::Display for Token {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Symbol(symbol) => write!(f, "{symbol}"),
            Self::Register(register) => write!(f, "{register}",),
            Self::Immediate(immediate) => write!(f, "{immediate}",),
            Self::StringLit(string_lit) => write!(f, "{string_lit}",),
            Self::Opcode(opcode) => write!(f, "{opcode}",),
        }
    }
}

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

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

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
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,
    AddI,
    SubI,

    // Pseudo-instructions
    Db,
    Dh,
    Dw,
    Resb,
    Resh,
    Resw,
    Push,
    Pop,
    Pusha,
    Popa,
    Lwi,
    Call,
    Return,

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

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

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}"),
            Self::InvalidOpcode(op) => write!(f, "instruction does not exist: {op}"),
        }
    }
}

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),
            "addi" => Ok(Self::AddI),
            "subi" => Ok(Self::SubI),
            "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),
            "call" => Ok(Self::Call),
            "return" => Ok(Self::Return),
            "pusha" => Ok(Self::Pusha),
            "popa" => Ok(Self::Popa),
            _ => Err(OpcodeFromStrError::InvalidOpcode(s.to_string())),
        }
    }
}

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", "addi", "subi", // Pseudo-instructions
        "db", "dh", "dw", "resb", "resh", "resw", "push", "pop", "lwi", "call", "return",
        "pusha", "popa", // meta instructions
        "include",
    ];

    #[must_use]
    pub const 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::AddI => Some(0x25),
            Self::SubI => Some(0x26),
            Self::Segment => Some(0x27),

            // Pseudo-instructions don't have opcode values
            _ => None,
        }
    }

    #[must_use]
    pub const 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
        )
    }
}
