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

pub 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<RegMemImm>) -> ERes {
        let dst = dst.into();
        let src = src.into();
        match dst {
            RegMem::Reg(mut dst) => match src {
                RegMemImm::Reg(src) => {
                    if dst.width() != src.width() {
                        return Err("src and dst are not same width".into());
                    }
                    self.prefix16(dst);
                    self.rex(dst, src, 0, dst)?;
                    self.bytes.push(0x88 | dst.not8());
                    self.modrm(src, dst);
                }
                RegMemImm::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, 0, 0, dst), 0xc7, 0xc0 | dst.base()]);
                        self.imm(src, Width::B32);
                    } else {
                        if src_width <= Width::B32 {
                            dst = dst.lower64();
                        }
                        self.rex(dst, 0, 0, dst)?;
                        self.bytes.push(0xb0 | (dst.not8() << 3) | dst.base());
                        self.imm(src, dst.width());
                    }
                }
                RegMemImm::Mem(src) => {
                    if src.width != dst.width() {
                        return Err("register & memory sizes don't match".into());
                    }
                    self.prefix32(src)?;
                    self.prefix16(dst);
                    self.rex(dst, dst, 0, src)?;
                    self.bytes.push(0x8a | dst.not8());
                    self.modrm(dst, src);
                }
            },
            RegMem::Mem(dst) => match src {
                RegMemImm::Reg(src) => {
                    if src.width() != dst.width {
                        return Err("register & memory sizes don't match".into());
                    }
                    self.prefix32(dst)?;
                    self.prefix16(src);
                    self.rex(dst, src, 0, dst)?;
                    self.bytes.push(0x88 | src.not8());
                    self.modrm(src, dst);
                }
                RegMemImm::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());
                    }
                    self.prefix32(dst)?;
                    self.prefix16(encode_width);
                    self.rex(dst, 0, 0, dst)?;
                    self.bytes.push(0xc6 | encode_width.not8());
                    self.modrm(0, dst);
                    self.imm(src, encode_width);
                }
                RegMemImm::Mem(_) => return Err("cannot move memory to memory".into()),
            },
        }
        Ok(())
    }

    pub fn push(&mut self, reg: impl Into<RegMemImm>) -> ERes {
        match reg.into() {
            RegMemImm::Reg(reg) => match reg.width() {
                Width::B64 => {
                    if reg.gt8() {
                        self.bytes.push(0x41);
                    }
                    self.bytes.push(0x50 | reg.base());
                }
                Width::B16 => todo!(),
                _ => return Err("register must be 64 or 16 bit".into()),
            },
            RegMemImm::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()),
            },
            RegMemImm::Mem(mem) => todo!(),
        }
        Ok(())
    }

    pub fn pop(&mut self, reg: RegW) -> 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: RegW, sym: Symbol) -> ERes {
        self.rex(1, dst, 0, 0)?;
        self.bytes.push(0x8d);
        self.modrm(dst, sym);
        Ok(())
    }

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

    fn add_sub(&mut self, dst: impl Into<RegMem>, src: impl Into<RegMemImm>, ext: u8) -> ERes {
        let dst = dst.into();
        match src.into() {
            RegMemImm::Reg(src) => {
                if src.width() != dst.width() {
                    return Err("incompatible widths".into());
                }
                self.prefix32(dst)?;
                self.prefix16(src);
                self.rex(dst, src, 0, dst)?;
                self.bytes.push(src.not8());
                self.modrm(src, dst);
            }
            RegMemImm::Imm(mut src) => {
                let mut imm_width = src.width_signed()?;
                let dst_width = dst.width().min(Width::B32);

                if imm_width > dst_width {
                    imm_width = src.width_unsigned()?;
                    if dst.width() == Width::B64 || imm_width > dst_width {
                        return Err("immediate overflow".into());
                    }
                    src = src.reinterpret(dst_width);
                    imm_width = src.width_signed()?;
                }
                let code = if dst.width() == Width::B8 {
                    0x80
                } else if imm_width == Width::B8 {
                    0x83
                } else {
                    imm_width = dst_width;
                    0x81
                };

                self.prefix32(dst)?;
                self.prefix16(dst_width);
                self.rex(dst, 0, 0, dst)?;
                self.bytes.push(code);
                self.modrm(ext, dst);
                self.imm(src, imm_width);
            }
            RegMemImm::Mem(src) => {
                let RegMem::Reg(dst) = dst else {
                    return Err("cannot add memory to memory".into());
                };
                if src.width != dst.width() {
                    return Err("incompatible widths".into());
                }
                self.prefix32(src)?;
                self.prefix16(dst);
                self.rex(dst, dst, 0, src)?;
                self.bytes.push(0x2 | dst.not8());
                self.modrm(dst, src);
            }
        }
        Ok(())
    }

    pub fn add(&mut self, dst: impl Into<RegMem>, src: impl Into<RegMemImm>) -> ERes {
        self.add_sub(dst, src, 0)
    }

    pub fn sub(&mut self, dst: impl Into<RegMem>, src: impl Into<RegMemImm>) -> ERes {
        self.add_sub(dst, src, 5)
    }

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

    fn prefix32(&mut self, mem: impl MaybeMem) -> Result<(), CompilerMsg> {
        let Some(mem) = mem.mem() else {
            return Ok(());
        };
        match mem.reg.width() {
            Width::B8 | Width::B16 => return Err("invalid register width".into()),
            Width::B32 => self.bytes.push(0x67),
            Width::B64 => (),
        }
        Ok(())
    }

    fn rex(&mut self, w: impl RexW, r: impl RexBit, x: u8, b: impl RexBit) -> ERes {
        if r.req() && b.req_no() || r.req_no() && b.req() {
            return Err("registers incompatible (REX)".into());
        }
        if w.rexw() || r.rex() || x.rex() || b.rex() || r.req() || b.req() {
            self.bytes.push(rex(w, r, x, b));
        }
        Ok(())
    }

    fn modrm(&mut self, reg: impl ModRMReg, rm: impl ModRMRM) {
        let addr = rm.addr();
        let mod_ = match addr {
            EffAddr::Mem0 | EffAddr::Sym(_) => 0b00,
            EffAddr::Mem8(_) => 0b01,
            EffAddr::Mem32(_) => 0b10,
            EffAddr::None => 0b11,
        };
        self.bytes
            .push(((mod_ as u8) << 6) | (reg.val() << 3) | rm.rm());
        if !matches!(addr, EffAddr::None) && rm.rm() == 0b100 {
            // SIB
            self.bytes.push(0x24);
        }
        match addr {
            EffAddr::Mem8(disp) => self.bytes.push(disp as u8),
            EffAddr::Mem32(disp) => self.bytes.extend(disp.to_le_bytes()),
            EffAddr::Sym(sym) => self.sym_offset4(sym),
            _ => (),
        }
    }

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