use std::fmt;

use crate::backend::dsa::registers::Register;

pub struct InsBlock {
    instructions: Vec<Instruction>,
}

impl InsBlock {
    pub fn new() -> Self {
        Self {
            instructions: vec![],
        }
    }

    pub fn insert(&mut self, index: usize, instr: Instruction) {
        self.instructions.insert(index, instr);
    }

    pub fn push(&mut self, instr: Instruction) {
        self.instructions.push(instr);
    }

    pub fn append(&mut self, mut other: Self) {
        self.instructions.append(&mut other.instructions);
    }

    pub fn extend(&mut self, instrs: impl IntoIterator<Item = Instruction>) {
        self.instructions.extend(instrs);
    }

    pub fn is_empty(&self) -> bool {
        self.instructions.is_empty()
    }

    pub fn len(&self) -> usize {
        self.instructions.len()
    }

    pub fn iter(&self) -> impl Iterator<Item = &Instruction> {
        self.instructions.iter()
    }
}

impl From<Vec<Instruction>> for InsBlock {
    fn from(instructions: Vec<Instruction>) -> Self {
        Self { instructions }
    }
}

impl From<Instruction> for InsBlock {
    fn from(instr: Instruction) -> Self {
        Self {
            instructions: vec![instr],
        }
    }
}

#[derive(Debug, Clone)]
pub enum Instruction {
    // Labels and comments
    Label(Label),
    Comment {
        text: String,
        top_level: bool,
    },
    Newline,

    // Data Directives
    Db {
        label: String,
        data: Vec<u8>,
    },
    Dh {
        label: String,
        data: Vec<u16>,
    },
    Dw {
        label: String,
        data: Vec<u32>,
    },
    DString {
        // alias for db.
        label: String,
        data: String,
    },

    Resx {
        label: String,
        size: u32,
    },

    // Include
    Include {
        name: String,
        path: String,
    },

    // Data movement
    Mov {
        src: Register,
        dest: Register,
    },
    Movs {
        src: Register,
        dest: Register,
    },

    // Memory operations
    Ldb {
        src: MemOperand,
        dest: Register,
    },
    Ldh {
        src: MemOperand,
        dest: Register,
    },
    Ldw {
        src: MemOperand,
        dest: Register,
    },
    Stb {
        src: Register,
        dest: MemOperand,
    },
    Sth {
        src: Register,
        dest: MemOperand,
    },
    Stw {
        src: Register,
        dest: MemOperand,
    },

    // Immediate loads
    Lli {
        imm: Imm,
        dest: Register,
    },
    Lui {
        imm: Imm,
        dest: Register,
    },
    Lwi {
        imm: Imm,
        dest: Register,
    },
    LwiLabel {
        label: String,
        dest: Register,
    },

    // Arithmetic
    Add {
        src1: Register,
        src2: Register,
        dest: Register,
    },
    Sub {
        src1: Register,
        src2: Register,
        dest: Register,
    },
    IAdd {
        src: Register,
        imm: Imm,
        dest: Option<Register>,
    },
    ISub {
        src: Register,
        imm: Imm,
        dest: Option<Register>,
    },
    Inc {
        reg: Register,
    },
    Dec {
        reg: Register,
    },

    // Bitwise
    And {
        src1: Register,
        src2: Register,
        dest: Register,
    },
    Or {
        src1: Register,
        src2: Register,
        dest: Register,
    },
    Xor {
        src1: Register,
        src2: Register,
        dest: Register,
    },
    Not {
        src: Register,
        dest: Register,
    },
    Nand {
        src1: Register,
        src2: Register,
        dest: Register,
    },
    Nor {
        src1: Register,
        src2: Register,
        dest: Register,
    },
    Xnor {
        src1: Register,
        src2: Register,
        dest: Register,
    },

    // Shifts
    Shl {
        src1: Register,
        r_shamt: Register,
        i_shamt: u16,
        dest: Register,
    },
    Shr {
        src1: Register,
        r_shamt: Register,
        i_shamt: u16,
        dest: Register,
    },

    // Comparison
    Cmp {
        reg1: Register,
        reg2: Register,
    },

    // Jumps
    Jmp {
        target: Label,
    },
    Jeq {
        target: Label,
    },
    Jne {
        target: Label,
    },
    Jgt {
        target: Label,
    },
    Jge {
        target: Label,
    },
    Jlt {
        target: Label,
    },
    Jle {
        target: Label,
    },

    // Stack
    Push {
        reg: Register,
    },
    Pop {
        reg: Register,
    },

    // Function calls
    Call {
        target: String,
    }, // namespace::function
    Return,

    // System
    Hlt,
    Nop,
    Int {
        code: u8,
    },
}

pub enum DataDirective {
    U8(Vec<u8>),
    U16(Vec<u16>),
    U32(Vec<u32>),
    String(String),
    Char(char),
}

impl fmt::Display for Instruction {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Self::Label(l) => write!(f, "{}:", l),
            Self::Newline => write!(f, ""), /* empty string as newlines are inserted */
            // automatically.
            Self::Comment { text, top_level } => write!(
                f,
                "{}",
                text.lines()
                    .map(|line| format!(
                        "{}// {}",
                        if *top_level { "" } else { "    " },
                        line.trim(),
                    ))
                    .collect::<Vec<String>>()
                    .join("\n")
            ),

            Self::Include { name, path } => write!(f, "include {name}: \"{}\"", path),

            Self::Db { label, data } => write!(
                f,
                "db {}: {}",
                label,
                data.iter()
                    .map(|&b| format!("{:#04X}", b))
                    .collect::<Vec<String>>()
                    .join(", ")
            ),
            Self::Dh { label, data } => write!(
                f,
                "dh {}: {}",
                label,
                data.iter()
                    .map(|&b| format!("{:#06X}", b))
                    .collect::<Vec<String>>()
                    .join(", ")
            ),
            Self::Dw { label, data } => write!(
                f,
                "dw {}: {}",
                label,
                data.iter()
                    .map(|&b| format!("{:#08X}", b))
                    .collect::<Vec<String>>()
                    .join(", ")
            ),
            Self::DString { label, data } => write!(f, "db {}: \"{}\"", label, data),

            Self::Resx { label, size } => write!(f, "resx {}: {}", label, size),

            Self::Mov { src, dest } => write!(f, "    mov {}, {}", src, dest),
            Self::Movs { src, dest } => write!(f, "    movs {}, {}", src, dest),

            Self::Ldb { src: addr, dest } => {
                let (reg, offset) = reg_and_offset(addr);
                write!(f, "    ldb {}, {}, {}", reg, dest, offset)
            }
            Self::Ldh { src: addr, dest } => {
                let (reg, offset) = reg_and_offset(addr);
                write!(f, "    ldh {}, {}, {}", reg, dest, offset)
            }
            Self::Ldw { src, dest } => {
                let (reg, offset) = reg_and_offset(src);
                write!(f, "    ldw {}, {}, {}", reg, dest, offset)
            }
            // Self::Ldbs { addr, dest } => {
            //     write!(f, "    ldbs {}, {}", format_mem_operand(addr), dest)
            // }
            // Self::Ldhs { addr, dest } => {
            //     write!(f, "    ldhs {}, {}", format_mem_operand(addr), dest)
            // }
            // Self::Ldws { addr, dest } => {
            //     write!(f, "    ldws {}, {}", format_mem_operand(addr), dest)
            // }
            Self::Stb { src, dest: addr } => {
                let (reg, offset) = reg_and_offset(addr);
                write!(f, "    stb {}, {}, {}", src, reg, offset)
            }
            Self::Sth { src, dest: addr } => {
                let (reg, offset) = reg_and_offset(addr);
                write!(f, "    sth {}, {}, {}", src, reg, offset)
            }
            Self::Stw { src, dest: addr } => {
                let (reg, offset) = reg_and_offset(addr);
                write!(f, "    stw {}, {}, {}", src, reg, offset)
            }

            Self::Lli { imm, dest } => write!(f, "    lli {}, {}", imm, dest),
            Self::Lui { imm, dest } => write!(f, "    lui {}, {}", imm, dest),
            Self::Lwi { imm, dest } => write!(f, "    lwi {}, {}", imm, dest),
            Self::LwiLabel { label, dest } => write!(f, "    lwi {}, {}", label, dest),

            // arithmetic
            Self::Add { src1, src2, dest } => {
                write!(f, "    add {}, {}, {}", src1, src2, dest)
            }
            Self::Sub { src1, src2, dest } => {
                write!(f, "    sub {}, {}, {}", src1, src2, dest)
            }
            Self::And { src1, src2, dest } => {
                write!(f, "    and {}, {}, {}", src1, src2, dest)
            }
            Self::Or { src1, src2, dest } => {
                write!(f, "    or {}, {}, {}", src1, src2, dest)
            }
            Self::Nand { src1, src2, dest } => {
                write!(f, "    nand {}, {}, {}", src1, src2, dest)
            }
            Self::Xor { src1, src2, dest } => {
                write!(f, "    xor {}, {}, {}", src1, src2, dest)
            }
            Self::Nor { src1, src2, dest } => {
                write!(f, "    nor {}, {}, {}", src1, src2, dest)
            }
            Self::Not { src, dest } => {
                write!(f, "    not {} {}", src, dest)
            }
            Self::Xnor { src1, src2, dest } => {
                write!(f, "    xnor {}, {}, {}", src1, src2, dest)
            }
            Self::IAdd { src, imm, dest } => {
                if let Some(d) = dest {
                    write!(f, "    addi {}, {}, {}", src, imm, d)
                } else {
                    write!(f, "    addi {}, {}", src, imm)
                }
            }
            Self::ISub { src, imm, dest } => {
                if let Some(d) = dest {
                    write!(f, "    subi {}, {}, {}", src, imm, d)
                } else {
                    write!(f, "    subi {}, {}", src, imm)
                }
            }

            // shift instructions
            Self::Shl {
                src1,
                r_shamt,
                i_shamt,
                dest,
            } => {
                write!(f, "    shl {}, {}, {}, {}", src1, r_shamt, i_shamt, dest)
            }
            Self::Shr {
                src1,
                r_shamt,
                i_shamt,
                dest,
            } => {
                write!(f, "    shl {}, {}, {}, {}", src1, r_shamt, i_shamt, dest)
            }

            // increment instructions
            Self::Inc { reg } => write!(f, "    inc {}", reg),
            Self::Dec { reg } => write!(f, "    dec {}", reg),

            Self::Cmp { reg1, reg2 } => write!(f, "    cmp {}, {}", reg1, reg2),

            // jump instructions
            Self::Jmp { target } => write!(f, "    jmp {}", target),
            Self::Jeq { target } => write!(f, "    jeq {}", target),
            Self::Jne { target } => write!(f, "    jne {}", target),
            Self::Jgt { target } => write!(f, "    jgt {}", target),
            Self::Jge { target } => write!(f, "    jge {}", target),
            Self::Jlt { target } => write!(f, "    jlt {}", target),
            Self::Jle { target } => write!(f, "    jle {}", target),

            // stack pseudoinstructions
            Self::Push { reg } => write!(f, "    push {}", reg),
            Self::Pop { reg } => write!(f, "    pop {}", reg),

            // call & return pseudoinstructions
            Self::Call { target } => write!(f, "    call {}", target),
            Self::Return => write!(f, "    return"),

            // misc instructions
            Self::Int { code } => write!(f, "    int {}", code),
            Self::Hlt => write!(f, "    hlt"),
            Self::Nop => write!(f, "    nop"),
        }
    }
}

impl Instruction {
    // data directives
    pub fn db_string(label: impl Into<String>, data: impl Into<String>) -> Self {
        Self::DString {
            label: label.into(),
            data: data.into(),
        }
    }

    pub fn db_word(label: impl Into<String>, data: u32) -> Self {
        Self::Dw {
            label: label.into(),
            data: vec![data],
        }
    }

    pub fn db_bytes(label: impl Into<String>, data: &[u8]) -> Self {
        Self::Db {
            label: label.into(),
            data: data.to_vec(),
        }
    }

    // Movement
    pub fn mov<R1, R2>(src: R1, dest: R2) -> Self
    where
        R1: Into<Register>,
        R2: Into<Register>,
    {
        Self::Mov {
            src: src.into(),
            dest: dest.into(),
        }
    }

    // Memory loads
    pub fn ldw_reg<R>(base: R, dest: Register) -> Self
    where
        R: Into<Register>,
    {
        Self::Ldw {
            src: MemOperand::RegIndirect(base.into()),
            dest,
        }
    }

    pub fn ldw_reg_offset<R>(base: R, dest: Register, offset: i32) -> Self
    where
        R: Into<Register>,
    {
        Self::Ldw {
            src: MemOperand::RegOffset(base.into(), offset),
            dest,
        }
    }

    pub fn ldw_label(label: impl Into<Label>, dest: Register) -> Self {
        Self::Ldw {
            src: MemOperand::Label(label.into()),
            dest,
        }
    }

    // Memory stores
    pub fn stw_reg<R>(src: Register, base: R) -> Self
    where
        R: Into<Register>,
    {
        Self::Stw {
            src,
            dest: MemOperand::RegIndirect(base.into()),
        }
    }

    pub fn stw_reg_offset<R>(src: Register, base: R, offset: i32) -> Self
    where
        R: Into<Register>,
    {
        Self::Stw {
            src,
            dest: MemOperand::RegOffset(base.into(), offset),
        }
    }

    pub fn stw_label(src: Register, label: impl Into<Label>) -> Self {
        Self::Stw {
            src,
            dest: MemOperand::Label(label.into()),
        }
    }

    // Arithmetic
    pub fn add(src1: Register, src2: Register, dest: Register) -> Self {
        Self::Add { src1, src2, dest }
    }

    pub fn sub(src1: Register, src2: Register, dest: Register) -> Self {
        Self::Sub { src1, src2, dest }
    }

    pub fn and(src1: Register, src2: Register, dest: Register) -> Self {
        Self::And { src1, src2, dest }
    }

    pub fn or(src1: Register, src2: Register, dest: Register) -> Self {
        Self::Or { src1, src2, dest }
    }

    pub fn xor(src1: Register, src2: Register, dest: Register) -> Self {
        Self::Xor { src1, src2, dest }
    }

    pub fn not(src: Register, dest: Register) -> Self {
        Self::Not { src, dest }
    }

    pub fn shl(src1: Register, r_shamt: Register, i_shamt: u16, dest: Register) -> Self {
        Self::Shl {
            src1,
            r_shamt,
            i_shamt,
            dest,
        }
    }

    pub fn shr(src1: Register, r_shamt: Register, i_shamt: u16, dest: Register) -> Self {
        Self::Shr {
            src1,
            r_shamt,
            i_shamt,
            dest,
        }
    }

    pub fn iadd(src: Register, value: i64) -> Self {
        let imm = Imm(value.unsigned_abs() as u32);

        if value < 0 {
            Self::ISub {
                src,
                imm,
                dest: None,
            }
        } else {
            Self::IAdd {
                src,
                imm,
                dest: None,
            }
        }
    }

    pub fn iadd_dest(src: Register, value: i32, dest: Register) -> Self {
        let imm = Imm(value.unsigned_abs());

        if value < 0 {
            Self::ISub {
                src,
                imm,
                dest: Some(dest),
            }
        } else {
            Self::IAdd {
                src,
                imm,
                dest: Some(dest),
            }
        }
    }

    pub fn inc(reg: Register) -> Self {
        Self::Inc { reg }
    }

    pub fn dec(reg: Register) -> Self {
        Self::Dec { reg }
    }

    // Immediate loads
    pub fn lwi(value: u32, dest: Register) -> Self {
        if value > 0xFFFF {
            Self::Lwi {
                imm: Imm(value),
                dest,
            }
        } else {
            Self::Lli {
                imm: Imm(value),
                dest,
            }
        }
    }

    pub fn lwi_label(label: impl Into<String>, dest: Register) -> Self {
        Self::LwiLabel {
            label: label.into(),
            dest,
        }
    }

    // Control flow
    pub fn label(name: impl Into<String>) -> Self {
        Self::Label(Label(name.into()))
    }

    pub fn jmp(target: impl Into<Label>) -> Self {
        Self::Jmp {
            target: target.into(),
        }
    }

    pub fn jeq(target: impl Into<Label>) -> Self {
        Self::Jeq {
            target: target.into(),
        }
    }

    pub fn cmp(reg1: Register, reg2: Register) -> Self {
        Self::Cmp { reg1, reg2 }
    }

    // Stack
    pub fn push(reg: Register) -> Self {
        Self::Push { reg }
    }

    pub fn pop(reg: Register) -> Self {
        Self::Pop { reg }
    }

    // Functions
    pub fn call(target: impl Into<String>) -> Self {
        Self::Call {
            target: target.into(),
        }
    }

    pub fn int(code: u8) -> Self {
        Self::Int { code }
    }

    pub fn ret() -> Self {
        Self::Return
    }

    // Utilities
    pub fn comment(text: impl Into<String>) -> Self {
        Self::Comment {
            text: text.into(),
            top_level: false,
        }
    }

    pub fn global_comment(text: impl Into<String>) -> Self {
        Self::Comment {
            text: text.into(),
            top_level: true,
        }
    }

    pub fn include(name: impl Into<String>, path: impl Into<String>) -> Self {
        Self::Include {
            name: name.into(),
            path: path.into(),
        }
    }
}

// Convenience trait for Label conversion
impl From<String> for Label {
    fn from(s: String) -> Self {
        Label(s)
    }
}

impl From<&str> for Label {
    fn from(s: &str) -> Self {
        Label(s.to_string())
    }
}

fn reg_and_offset(op: &MemOperand) -> (String, i32) {
    match op {
        MemOperand::RegIndirect(reg) => (reg.to_string(), 0),
        MemOperand::RegOffset(reg, offset) => (reg.to_string(), *offset),
        MemOperand::Label(label) => (label.to_string(), 0),
        MemOperand::LabelOffset(label, offset) => (label.to_string(), *offset),
    }
}

/// Memory operand for loads/stores
#[derive(Debug, Clone)]
pub enum MemOperand {
    /// Register indirect: [reg]
    RegIndirect(Register),
    /// Register with offset: [reg + offset]
    RegOffset(Register, i32),
    /// Label: [label]
    Label(Label),
    /// Label with offset: [label + offset]
    LabelOffset(Label, i32),
}

/// Immediate value (16-bit or 32-bit)
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Imm(pub u32);

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

/// Label reference
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Label(pub String);

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