use std::{fmt::Display, ops::Add};

use crate::{backend::register_allocator::{REG_FP, REG_LR, REG_PC, REG_R0, REG_R1, REG_R2, REG_R3, REG_R12, REG_SP, Register}, ir::types::CmpOp as IRCmpOp};
pub enum ARMInstr{
    Move(MoveInstr),
    Movw(MovwInstr),
    Movt(MovtInstr),
    Load(LoadInstr),
    Store(StoreInstr),
    Mul(MulInstr),
    SDiv(SDivInstr),
    Add(AddInstr),
    Sub(SubInstr),
    Rsb(RsbInstr),
    Cmp(CmpInstr),
    Push(PushInstr),
    Pop(PopInstr),
    FunctionHead(String, usize),
    Bl(BlInstr),
    B(BInstr),
    Label(String),
}

impl Display for ARMInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ARMInstr::Move(instr) => write!(f, "{}", instr),
            ARMInstr::Movw(instr) => write!(f, "{}", instr),
            ARMInstr::Movt(instr) => write!(f, "{}", instr),
            ARMInstr::Load(instr) => write!(f, "{}", instr),
            ARMInstr::Store(instr) => write!(f, "{}", instr),
            ARMInstr::Mul(instr) => write!(f, "{}", instr),
            ARMInstr::SDiv(instr) => write!(f, "{}", instr),
            ARMInstr::Add(instr) => write!(f, "{}", instr),
            ARMInstr::Sub(instr) => write!(f, "{}", instr),
            ARMInstr::Rsb(instr) => write!(f, "{}", instr),
            ARMInstr::Cmp(instr) => write!(f, "{}", instr),
            ARMInstr::Push(instr) => write!(f, "{}", instr),
            ARMInstr::Pop(instr) => write!(f, "{}", instr),
            ARMInstr::Bl(instr) => write!(f, "{}", instr),
            ARMInstr::B(instr) => write!(f, "{}", instr),
            ARMInstr::FunctionHead(name, align_size) => write!(f, ".align {}\n.global {}\n.type {}, %function\n{}:", align_size, name, name, name),
            ARMInstr::Label(name) => write!(f, "{}:", name),
        }
    }
}
pub enum RegisterOrImm {
    Reg(Register),
    Imm(i32),
}   
impl Display for RegisterOrImm {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            RegisterOrImm::Reg(reg) => write!(f, "{}", reg),
            RegisterOrImm::Imm(imm) => write!(f, "#{}", imm),
        }
    }
}
pub enum ConditionCode {
    Eq,
    Ne,
    Lt,
    Le,
    Gt,
    Ge,
}
impl From<IRCmpOp> for ConditionCode {
    fn from(cmp_op: IRCmpOp) -> Self {
        match cmp_op {
            IRCmpOp::Eq => ConditionCode::Eq,
            IRCmpOp::Ne => ConditionCode::Ne,
            IRCmpOp::Lt => ConditionCode::Lt,
            IRCmpOp::Le => ConditionCode::Le,
            IRCmpOp::Gt => ConditionCode::Gt,
            IRCmpOp::Ge => ConditionCode::Ge,
        }
    }
}
impl Display for ConditionCode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let code_str = match self {
            ConditionCode::Eq => "eq",
            ConditionCode::Ne => "ne",
            ConditionCode::Lt => "lt",
            ConditionCode::Le => "le",
            ConditionCode::Gt => "gt",
            ConditionCode::Ge => "ge",
        };
        write!(f, "{}", code_str)
    }
}
// pub enum RegisterOrMemory {
// 	Reg(Register),
// 	Mem(Register, Option<RegisterOrImm>), // Base register and optional offset
// 	MemPesudoName(String),
// }
// impl Display for RegisterOrMemory {
// 	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
// 		match self {
// 			RegisterOrMemory::Reg(reg) => write!(f, "{}", reg),
// 			RegisterOrMemory::Mem(base, Some(offset)) => write!(f, "[{}, {}]", base, offset),
// 			RegisterOrMemory::Mem(base, None) => write!(f, "[{}]", base),
// 			RegisterOrMemory::MemPesudoName(name) => write!(f, "{}", name),
// 		}
// 	}
// }
pub struct MoveInstr(Option<ConditionCode>, Register, RegisterOrImm);
impl MoveInstr {
    pub fn new_uncond(dest: Register, src: RegisterOrImm) -> ARMInstr {
        ARMInstr::Move(MoveInstr(None, dest, src))
    }
    pub fn new_cond(condition: ConditionCode, dest: Register, src: RegisterOrImm) -> ARMInstr {
        ARMInstr::Move(MoveInstr(Some(condition), dest, src))
    }
    pub fn new_sp_to_fp() -> ARMInstr {
        ARMInstr::Move(MoveInstr(None, REG_FP, RegisterOrImm::Reg(REG_SP)))
    }
    pub fn new_fp_to_sp() -> ARMInstr {
        ARMInstr::Move(MoveInstr(None, REG_SP, RegisterOrImm::Reg(REG_FP)))
    }
}
impl Display for MoveInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let MoveInstr(condition, dest, src) = self;
        if let Some(condition) = condition {
            write!(f, "mov{} {}, {}", condition, dest, src)
        } else {
            write!(f, "mov {}, {}", dest, src)
        }
    }
}

pub enum MoveWideOperand {
    Imm(u16),
    Symbol(String),
}

pub struct MovwInstr(Register, MoveWideOperand);
impl MovwInstr {
    pub fn new_imm(dest: Register, value: u16) -> ARMInstr {
        ARMInstr::Movw(MovwInstr(dest, MoveWideOperand::Imm(value)))
    }

    pub fn new_symbol(dest: Register, name: String) -> ARMInstr {
        ARMInstr::Movw(MovwInstr(dest, MoveWideOperand::Symbol(name)))
    }
}
impl Display for MovwInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let MovwInstr(dest, operand) = self;
        match operand {
            MoveWideOperand::Imm(value) => write!(f, "movw {}, #{}", dest, value),
            MoveWideOperand::Symbol(name) => write!(f, "movw {}, #:lower16:{}", dest, name),
        }
    }
}

pub struct MovtInstr(Register, MoveWideOperand);
impl MovtInstr {
    pub fn new_imm(dest: Register, value: u16) -> ARMInstr {
        ARMInstr::Movt(MovtInstr(dest, MoveWideOperand::Imm(value)))
    }

    pub fn new_symbol(dest: Register, name: String) -> ARMInstr {
        ARMInstr::Movt(MovtInstr(dest, MoveWideOperand::Symbol(name)))
    }
}
impl Display for MovtInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let MovtInstr(dest, operand) = self;
        match operand {
            MoveWideOperand::Imm(value) => write!(f, "movt {}, #{}", dest, value),
            MoveWideOperand::Symbol(name) => write!(f, "movt {}, #:upper16:{}", dest, name),
        }
    }
}

pub struct LoadInstr(Register, Register, Option<RegisterOrImm>);
impl LoadInstr {
    pub fn new(dest: Register, base: Register, offset: Option<RegisterOrImm>) -> ARMInstr {
        ARMInstr::Load(LoadInstr(dest, base, offset))
    }
    pub fn new_stack(dest: Register, offset: i32) -> ARMInstr {
        ARMInstr::Load(LoadInstr(dest, REG_FP, Some(RegisterOrImm::Imm(-(offset as i32)))))
    }
}
impl Display for LoadInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let LoadInstr(dest, base, offset) = self;
        if let Some(offset) = offset {
            write!(f, "ldr {}, [{}, {}]", dest, base, offset)
        } else {
            write!(f, "ldr {}, [{}]", dest, base)
        }
    }
}
pub struct StoreInstr(Register, Register, Option<RegisterOrImm>);
impl StoreInstr {
    pub fn new(src: Register, dest: Register, offset: Option<RegisterOrImm>) -> ARMInstr {
        ARMInstr::Store(StoreInstr(src, dest, offset))
    }
    pub fn new_stack(dest: Register, offset: i32) -> ARMInstr {
        ARMInstr::Store(StoreInstr(dest, REG_FP, Some(RegisterOrImm::Imm(-(offset as i32)))))
    }
}
impl Display for StoreInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let StoreInstr(src, dest, offset) = self;
        if let Some(offset) = offset {
            write!(f, "str {}, [{}, {}]", src, dest, offset)
        } else {
            write!(f, "str {}, [{}]", src, dest)
        }
    }
}

pub struct MulInstr(Register, Register, RegisterOrImm);
impl MulInstr {
    pub fn new(dest: Register, left: Register, right: RegisterOrImm) -> ARMInstr {
        ARMInstr::Mul(MulInstr(dest, left, right))
    }
}
impl Display for MulInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let MulInstr(dest, left, right) = self;
        write!(f, "mul {}, {}, {}", dest, left, right)
    }
}
pub struct SDivInstr(Register, Register, RegisterOrImm);
impl SDivInstr {
    pub fn new(dest: Register, left: Register, right: RegisterOrImm) -> ARMInstr {
        ARMInstr::SDiv(SDivInstr(dest, left, right))
    }
}
impl Display for SDivInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let SDivInstr(dest, left, right) = self;
        write!(f, "sdiv {}, {}, {}", dest, left, right)
    }
}
pub struct AddInstr(Register, Register, RegisterOrImm);
impl AddInstr {
    pub fn new(dest: Register, left: Register, right: RegisterOrImm) -> ARMInstr {
        ARMInstr::Add(AddInstr(dest, left, right))
    }
    pub fn new_sp(offset: i32) -> ARMInstr {
        ARMInstr::Add(AddInstr(REG_SP, REG_SP, RegisterOrImm::Imm(offset)))
    }
}
impl Display for AddInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let AddInstr(dest, left, right) = self;
        write!(f, "add {}, {}, {}", dest, left, right)
    }
}
pub struct SubInstr(Register, Register, RegisterOrImm);
impl SubInstr {
    pub fn new(dest: Register, left: Register, right: RegisterOrImm) -> ARMInstr {
        ARMInstr::Sub(SubInstr(dest, left, right))
    }
    pub fn new_sp(offset: i32) -> ARMInstr {
        ARMInstr::Sub(SubInstr(REG_SP, REG_SP, RegisterOrImm::Imm(offset)))
    }
}

impl Display for SubInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let SubInstr(dest, left, right) = self;
        write!(f, "sub {}, {}, {}", dest, left, right)
    }
}
pub struct RsbInstr(Register, Register, RegisterOrImm);
impl RsbInstr {
    pub fn new(dest: Register, left: Register, right: RegisterOrImm) -> ARMInstr {
        ARMInstr::Rsb(RsbInstr(dest, left, right))
    }
}
impl Display for RsbInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let RsbInstr(dest, left, right) = self;
        write!(f, "rsb {}, {}, {}", dest, left, right)
    }
}
pub struct CmpInstr(Register, RegisterOrImm);
impl Display for CmpInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let CmpInstr(left, right) = self;
        write!(f, "cmp {}, {}", left, right)
    }
}
impl CmpInstr {
    pub fn new(left: Register, right: RegisterOrImm) -> ARMInstr {
        ARMInstr::Cmp(CmpInstr(left, right))
    }
}
pub struct PushInstr(Vec<Register>);
impl PushInstr {
    pub fn new_push_fp_lr() -> ARMInstr {
         ARMInstr::Push(PushInstr(vec![
            REG_FP,
            REG_LR,
        ]))
    }
    pub fn new_push_caller_save() -> ARMInstr {
        ARMInstr::Push(PushInstr(vec![
            REG_R0,
            REG_R1,
            REG_R2,
            REG_R3,
            REG_R12,
        ]))
    }
}
impl Display for PushInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) ->
        std::fmt::Result {
        let PushInstr(registers) = self;
        let regs_str = registers.iter().map(|reg| format!("{}", reg)).collect::<Vec<_>>().join(", ");
        write!(f, "push {{{}}}", regs_str)
    }
}

pub struct PopInstr(Vec<Register>);
impl PopInstr {
    pub fn new_pop_fp_pc() -> ARMInstr {
         ARMInstr::Pop(PopInstr(vec![
            REG_FP,
            REG_PC,
        ]))
    }
    pub fn new_pop_caller_save() -> ARMInstr {
        ARMInstr::Pop(PopInstr(vec![
            REG_R0,
            REG_R1,
            REG_R2,
            REG_R3,
            REG_R12,
        ]))
    }
}
impl Display for PopInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let PopInstr(registers) = self;
        let regs_str = registers.iter().map(|reg| format!("{}", reg)).collect::<Vec<_>>().join(", ");
        write!(f, "pop {{{}}}", regs_str)
    }
}

pub struct BlInstr(Option<ConditionCode>, String);
impl Display for BlInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match &self.0 {
            Some(condition) => write!(f, "bl{} {}", condition, self.1),
            None => write!(f, "bl {}", self.1),
        }
    }
}

impl BlInstr {
    pub fn new(label_name: String) -> ARMInstr {
        ARMInstr::Bl(BlInstr(None, label_name))
    }

    pub fn new_cond(condition: ConditionCode, label_name: String) -> ARMInstr {
        ARMInstr::Bl(BlInstr(Some(condition), label_name))
    }
}
pub struct BInstr(Option<ConditionCode>, String);
impl Display for BInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match &self.0 {
            Some(condition) => write!(f, "b{} {}", condition, self.1),
            None => write!(f, "b {}", self.1),
        }
    }
}

impl BInstr {
    pub fn new(label_name: String) -> ARMInstr {
        ARMInstr::B(BInstr(None, label_name))
    }

    pub fn new_cond(condition: ConditionCode, label_name: String) -> ARMInstr {
        ARMInstr::B(BInstr(Some(condition), label_name))
    }
}