use super::*;
use crate::backend::Symbol;

type ERes = Result<(), CompilerMsg>;

/// machine code
#[derive(Default)]
pub struct Code {
    pub(super) bytes: Vec<u8>,
    pub(super) missing: Vec<(usize, Symbol)>,
}

impl Code {
    pub fn mov(&mut self, dst: impl Into<RegMem>, src: impl Into<RegImmMem>) -> ERes {
        let dst = dst.into();
        let src = src.into();
        match dst {
            RegMem::Reg(mut dst) => match src {
                RegImmMem::Reg(src) => {
                    if dst.width() != src.width() {
                        return Err("src and dst are not same width".into());
                    }
                    if dst.incompatible(&src) {
                        return Err("incompatible registers due to rex".into());
                    }
                    let width = dst.width();
                    self.prefix16(width);
                    if src.requires_rex() || dst.requires_rex() {
                        self.bytes.push(rex(width, src, 0, dst));
                    }
                    self.bytes.push(0x88 | width.gt8() as u8);
                    self.bytes.push(modrm_regs(src, dst));
                }
                RegImmMem::Imm(src) => {
                    let src_width = src.width_unsigned()?;
                    if src_width > dst.width() {
                        return Err("immediate cannot fit in register".into());
                    }
                    self.prefix16(dst);
                    if dst.width() == Width::B64 && src_width <= Width::B32 && src.0 < 0 {
                        // use different op that sign extends for less bytes
                        self.bytes
                            .extend([rex(dst.width(), 0, 0, dst), 0xc7, 0xc0 | dst.base()]);
                        self.imm(src, Width::B32);
                    } else {
                        if src_width <= Width::B32 {
                            dst = dst.lower64();
                        }
                        if dst.requires_rex() {
                            self.bytes.push(rex(dst.width(), 0, 0, dst));
                        }
                        let opcode = 0xb0 | ((dst.width().gt8() as u8) << 3);
                        self.bytes.push(opcode | dst.base());
                        self.imm(src, dst.width());
                    }
                }
                RegImmMem::Mem(src) => todo!(),
            },
            RegMem::Mem(dst) => match src {
                RegImmMem::Reg(src) => todo!(),
                RegImmMem::Imm(src) => {
                    let encode_width = dst.width.min(Width::B32);
                    let src_width = if dst.width == Width::B64 {
                        src.width_signed()
                    } else {
                        src.width_unsigned()
                    }?;
                    if src_width == Width::B64 {
                        return Err("cannot move 64 bit immediate into memory".into());
                    }
                    if src_width > dst.width {
                        return Err("source cannot fit in destination".into());
                    }
                    match dst.reg.width() {
                        Width::B8 | Width::B16 => return Err("invalid register width".into()),
                        Width::B32 => self.bytes.push(0x67),
                        Width::B64 => (),
                    }
                    self.prefix16(encode_width);
                    if dst.reg.requires_mem_rex() || dst.width == Width::B64 {
                        self.bytes.push(rex(dst.width, 0, 0, dst.reg));
                    }
                    self.bytes.push(0xc6 | (encode_width != Width::B8) as u8);
                    self.modrm_regdisp(dst.reg, dst.disp);
                    self.imm(src, encode_width);
                }
                RegImmMem::Mem(_) => return Err("cannot move memory to memory".into()),
            },
        }
        Ok(())
    }

    pub fn push(&mut self, reg: impl Into<RegImmMem>) -> ERes {
        match reg.into() {
            RegImmMem::Reg(reg) => match reg.width() {
                Width::B64 => {
                    if reg.gt8() {
                        self.bytes.push(0x41);
                    }
                    self.bytes.push(0x50 | reg.base());
                }
                Width::B16 => {}
                _ => return Err("register must be 64 or 16 bit".into()),
            },
            RegImmMem::Imm(imm) => match imm.width_unsigned()? {
                Width::B8 => {
                    self.bytes.push(0x6a);
                    self.bytes.push(imm.0 as u8);
                }
                Width::B16 | Width::B32 => {
                    self.bytes.push(0x68);
                    self.bytes.extend((imm.0 as u32).to_le_bytes());
                }
                Width::B64 => return Err("immediate must be 32 bit or less".into()),
            },
            RegImmMem::Mem(mem) => todo!(),
        }
        Ok(())
    }

    pub fn pop(&mut self, reg: Reg) -> ERes {
        match reg.width() {
            Width::B64 | Width::B16 => (),
            _ => return Err("register must be 64 or 16 bit".into()),
        }
        self.prefix16(reg);
        if reg.gt8() {
            self.bytes.push(0x41);
        }
        self.bytes.push(0x58 | reg.base());
        Ok(())
    }

    pub fn lea(&mut self, dst: Reg, sym: Symbol) {
        self.bytes
            .extend([rex(1, dst, 0, 0), 0x8d, modrm_disp32(dst)]);
        self.sym_offset4(sym);
    }

    pub fn int(&mut self, code: u8) {
        self.bytes.extend([0xcd, code])
    }

    pub fn syscall(&mut self) {
        self.bytes.extend([0x0f, 0x05])
    }

    pub fn call(&mut self, sym: Symbol) {
        self.bytes.push(0xe8);
        self.sym_offset4(sym);
    }

    pub fn call_mem(&mut self, sym: Symbol) {
        self.bytes.extend([0xff, 0x15]);
        self.sym_offset4(sym);
    }

    pub fn ret(&mut self) {
        self.bytes.push(0xc3);
    }

    pub fn sub(&mut self) {
        // sub esp 40 iirc
        self.bytes.extend([0x48, 0x83, 0xec, 0x28]);
    }

    fn prefix16(&mut self, width: impl Into<Width>) {
        if width.into() == Width::B16 {
            self.bytes.push(0x66);
        }
    }

    fn modrm_regdisp(&mut self, reg: Reg, disp: i32) {
        const I8_MIN: i32 = i8::MIN as i32;
        const I8_MAX: i32 = i8::MAX as i32;
        let mod_ = match disp {
            0 => 0b00,
            I8_MIN..=I8_MAX => 0b01,
            _ => 0b10,
        };
        self.bytes.push(modrm(mod_, 0, reg.base()));
        if reg.base() == rsp.base() {
            // SIB
            self.bytes.push(0x24);
        }
        match mod_ {
            0b00 => (),
            0b01 => self.bytes.push(disp as u8),
            0b10 => self.bytes.extend(disp.to_le_bytes()),
            _ => unreachable!(),
        }
    }

    /// inserts a 32 bit offset from a symbol
    fn sym_offset4(&mut self, sym: Symbol) {
        let pos = self.bytes.len();
        self.bytes.extend([0; 4]);
        self.missing.push((pos, sym));
    }

    pub fn extend(&mut self, other: &Code) {
        let pos = self.bytes.len();
        self.bytes.extend(&other.bytes);
        self.missing
            .extend(other.missing.iter().map(|&(p, s)| (pos + p, s)));
    }

    fn imm(&mut self, imm: Imm, width: Width) {
        self.bytes.extend(&imm.0.to_le_bytes()[..width.bytes()]);
    }
}

pub fn encode(f: impl FnOnce(&mut Code) -> Result<(), CompilerMsg>) -> Result<Code, CompilerMsg> {
    let mut code = Code::default();
    f(&mut code)?;
    Ok(code)
}