//! All types which a Mnemonic can have as some kind of operand.
//! This includes things such as immediates, ModRM byte targets, etc. etc.

// used in doc, but not code
#[allow(unused_imports)]
use crate::register::SegmentRegister;

use crate::{disasm::DisasmError, register::Register};
use core::fmt;
use std::ops::{Add, BitAnd, BitOr, BitXor, Div, Mul, Sub};

pub type Byte = u8; // b
pub type IByte = i8; // used for displacements of memory access
pub type Word = u16; // w or v
pub type IWord = i16; // used for displacement of memory access
pub type DWord = u32;

#[derive(Debug, Clone, Ord, Eq, PartialEq, PartialOrd, Copy)]
/// Encodes either Byte- or Word-sized immediate operands.
/// Also sometimes used to decide if an instruction is Byte- or Word-sized,
/// which is usually indicated by using a value of 0 and the disregarding
/// the value when read.
/// Can either be interpreted as signed or unsigned, depending on the context.
pub enum ImmediateOperand {
    Byte(Byte),
    Word(Word),
}

impl ImmediateOperand {
    // Sign-extend [`Self::Byte`] into [`Self::Word`].
    // Returns [`Self::Word`], if already a word.
    pub fn sign_extend(self) -> Self {
        match self {
            Self::Byte(_) => {
                return if self.msb() {
                    self.flip_sign().as_word().flip_sign()
                } else {
                    self.as_word()
                };
            }
            Self::Word(_) => self,
        }
    }

    // Interprets [`Self::Byte`] as [`Self::Word`].
    // Returns word, if already a [`Self::Word`].
    pub fn as_word(self) -> Self {
        match self {
            Self::Byte(b) => Self::Word(b as Word),
            Self::Word(_) => self,
        }
    }

    // Flip most significant bit.
    pub fn flip_sign(self) -> Self {
        match self {
            Self::Byte(b) => Self::Byte(b ^ (1 << 7)),
            Self::Word(w) => Self::Word(w ^ (1 << 15)),
        }
    }

    // Check if value is zero.
    pub fn zero(&self) -> bool {
        match self {
            Self::Byte(byte) => return *byte == 0,
            Self::Word(word) => return *word == 0,
        }
    }

    // Check if leasy significant byte has even number of 1's.
    pub fn parity(&self) -> bool {
        match self {
            Self::Byte(byte) => return byte.count_ones() % 2 != 0,
            Self::Word(word) => {
                let [low, _]: [u8; 2] = word.to_le_bytes();
                return low.count_ones() % 2 != 0;
            }
        }
    }

    /// Check if least significant bit is set.
    pub fn _lsb(&self) -> bool {
        match self {
            Self::Byte(byte) => return byte & 1 == 1,
            Self::Word(word) => return word & 1 == 1,
        }
    }

    /// Check if most significant bit is set.
    pub fn msb(&self) -> bool {
        match self {
            Self::Byte(byte) => return (byte >> 7) == 1,
            Self::Word(word) => return (word >> 15) == 1,
        }
    }
}

impl Into<usize> for ImmediateOperand {
    fn into(self) -> usize {
        match self {
            ImmediateOperand::Byte(b) => b as usize,
            ImmediateOperand::Word(w) => w as usize,
        }
    }
}

impl Add for ImmediateOperand {
    type Output = Self;

    fn add(self, other: Self) -> Self {
        match self {
            ImmediateOperand::Byte(lhsb) => match other {
                ImmediateOperand::Byte(rhsb) => ImmediateOperand::Byte(lhsb.wrapping_add(rhsb)),
                ImmediateOperand::Word(rhsw) => ImmediateOperand::Word(match other.sign_extend() {
                    ImmediateOperand::Word(lhsw) => lhsw.wrapping_add(rhsw),
                    _ => panic!("unreachable"),
                }),
            },
            ImmediateOperand::Word(lhsw) => match other {
                ImmediateOperand::Word(rhsw) => ImmediateOperand::Word(lhsw.wrapping_add(rhsw)),
                ImmediateOperand::Byte(_) => ImmediateOperand::Word(match other.sign_extend() {
                    ImmediateOperand::Word(rhsw) => lhsw.wrapping_add(rhsw),
                    _ => panic!("unreachable"),
                }),
            },
        }
    }
}

impl Add<Displacement> for ImmediateOperand {
    type Output = ImmediateOperand;

    fn add(self, disp: Displacement) -> Self::Output {
        // Warning: this gets rid of the sign, which is fine as long as it is
        // used for a memory index, which can never be negative.
        match disp {
            Displacement::IByte(byte) => {
                if byte < 0 {
                    return self - ImmediateOperand::Byte((byte * -1) as Byte);
                } else {
                    return self + ImmediateOperand::Byte(byte as Byte);
                }
            }
            Displacement::IWord(word) => {
                if word < 0 {
                    return self - ImmediateOperand::Word((word * -1) as Word);
                } else {
                    return self + ImmediateOperand::Word(word as Word);
                }
            }
        }
    }
}

impl Add<Byte> for ImmediateOperand {
    type Output = ImmediateOperand;

    fn add(self, imm: Byte) -> Self::Output {
        match self {
            Self::Byte(b) => Self::Byte(b.wrapping_add(imm)),
            Self::Word(w) => Self::Word(w.wrapping_add(imm as Word)),
        }
    }
}
impl Sub for ImmediateOperand {
    type Output = Self;

    fn sub(self, other: Self) -> Self {
        match self {
            ImmediateOperand::Byte(lhsb) => match other {
                ImmediateOperand::Byte(rhsb) => ImmediateOperand::Byte(lhsb.wrapping_sub(rhsb)),
                ImmediateOperand::Word(rhsw) => ImmediateOperand::Word(match other.sign_extend() {
                    ImmediateOperand::Word(lhsw) => lhsw.wrapping_sub(rhsw),
                    _ => panic!("unreachable"),
                }),
            },
            ImmediateOperand::Word(lhsw) => match other {
                ImmediateOperand::Word(rhsw) => ImmediateOperand::Word(lhsw.wrapping_sub(rhsw)),
                ImmediateOperand::Byte(_) => ImmediateOperand::Word(match other.sign_extend() {
                    ImmediateOperand::Word(rhsw) => lhsw.wrapping_sub(rhsw),
                    _ => panic!("unreachable"),
                }),
            },
        }
    }
}

impl Sub<Byte> for ImmediateOperand {
    type Output = ImmediateOperand;

    fn sub(self, imm: Byte) -> Self::Output {
        match self {
            Self::Byte(b) => Self::Byte(b.wrapping_sub(imm)),
            Self::Word(w) => Self::Word(w.wrapping_sub(imm as Word)),
        }
    }
}
impl Mul for ImmediateOperand {
    type Output = Self;

    fn mul(self, other: Self) -> Self {
        match self {
            ImmediateOperand::Byte(lhsb) => match other {
                ImmediateOperand::Byte(rhsb) => ImmediateOperand::Byte(lhsb.wrapping_mul(rhsb)),
                ImmediateOperand::Word(rhsw) => ImmediateOperand::Word(match other.sign_extend() {
                    ImmediateOperand::Word(lhsw) => lhsw.wrapping_mul(rhsw),
                    _ => panic!("unreachable"),
                }),
            },
            ImmediateOperand::Word(lhsw) => match other {
                ImmediateOperand::Word(rhsw) => ImmediateOperand::Word(lhsw.wrapping_mul(rhsw)),
                ImmediateOperand::Byte(_) => ImmediateOperand::Word(match other.sign_extend() {
                    ImmediateOperand::Word(rhsw) => lhsw.wrapping_mul(rhsw),
                    _ => panic!("unreachable"),
                }),
            },
        }
    }
}

impl Div for ImmediateOperand {
    type Output = Self;

    fn div(self, other: Self) -> Self {
        match self {
            ImmediateOperand::Byte(lhsb) => match other {
                ImmediateOperand::Byte(rhsb) => ImmediateOperand::Byte(lhsb.wrapping_div(rhsb)),
                ImmediateOperand::Word(rhsw) => ImmediateOperand::Word(match other.sign_extend() {
                    ImmediateOperand::Word(lhsw) => lhsw.wrapping_div(rhsw),
                    _ => panic!("unreachable"),
                }),
            },
            ImmediateOperand::Word(lhsw) => match other {
                ImmediateOperand::Word(rhsw) => ImmediateOperand::Word(lhsw.wrapping_div(rhsw)),
                ImmediateOperand::Byte(_) => ImmediateOperand::Word(match other.sign_extend() {
                    ImmediateOperand::Word(rhsw) => lhsw.wrapping_div(rhsw),
                    _ => panic!("unreachable"),
                }),
            },
        }
    }
}

impl BitOr for ImmediateOperand {
    type Output = Self;

    fn bitor(self, other: Self) -> Self {
        match self {
            ImmediateOperand::Byte(lhsb) => match other {
                ImmediateOperand::Byte(rhsb) => ImmediateOperand::Byte(lhsb | rhsb),
                ImmediateOperand::Word(rhsw) => match other.sign_extend() {
                    ImmediateOperand::Word(lhsw) => ImmediateOperand::Word(lhsw | rhsw),
                    _ => panic!("unreachable"),
                },
            },
            ImmediateOperand::Word(lhsw) => match other {
                ImmediateOperand::Word(rhsw) => ImmediateOperand::Word(lhsw | rhsw),
                ImmediateOperand::Byte(_) => match other.sign_extend() {
                    ImmediateOperand::Word(rhsw) => ImmediateOperand::Word(lhsw | rhsw),
                    _ => panic!("unreachable"),
                },
            },
        }
    }
}

impl BitAnd for ImmediateOperand {
    type Output = Self;

    fn bitand(self, other: Self) -> Self {
        match self {
            ImmediateOperand::Byte(lhsb) => match other {
                ImmediateOperand::Byte(rhsb) => ImmediateOperand::Byte(lhsb & rhsb),
                ImmediateOperand::Word(rhsw) => match other.sign_extend() {
                    ImmediateOperand::Word(lhsw) => ImmediateOperand::Word(lhsw & rhsw),
                    _ => panic!("unreachable"),
                },
            },
            ImmediateOperand::Word(lhsw) => match other {
                ImmediateOperand::Word(rhsw) => ImmediateOperand::Word(lhsw & rhsw),
                ImmediateOperand::Byte(_) => match other.sign_extend() {
                    ImmediateOperand::Word(rhsw) => ImmediateOperand::Word(lhsw & rhsw),
                    _ => panic!("unreachable"),
                },
            },
        }
    }
}

impl BitXor for ImmediateOperand {
    type Output = Self;

    fn bitxor(self, other: Self) -> Self {
        match self {
            ImmediateOperand::Byte(lhsb) => match other {
                ImmediateOperand::Byte(rhsb) => ImmediateOperand::Byte(lhsb ^ rhsb),
                ImmediateOperand::Word(rhsw) => match other.sign_extend() {
                    ImmediateOperand::Word(lhsw) => ImmediateOperand::Word(lhsw ^ rhsw),
                    _ => panic!("unreachable"),
                },
            },
            ImmediateOperand::Word(lhsw) => match other {
                ImmediateOperand::Word(rhsw) => ImmediateOperand::Word(lhsw ^ rhsw),
                ImmediateOperand::Byte(_) => match other.sign_extend() {
                    ImmediateOperand::Word(rhsw) => ImmediateOperand::Word(lhsw ^ rhsw),
                    _ => panic!("unreachable"),
                },
            },
        }
    }
}

impl fmt::Display for ImmediateOperand {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Self::Byte(byte) => write!(f, "{}", byte),
            Self::Word(word) => write!(f, "{}", word),
        }
    }
}

impl fmt::LowerHex for ImmediateOperand {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Byte(b) => fmt::LowerHex::fmt(b, f),
            Self::Word(v) => fmt::LowerHex::fmt(v, f),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Copy)]
/// ModRM byte can either target a [`MemoryIndex`] (location in memory) or some
/// [`Register`].
pub enum ModRmTarget {
    Memory(MemoryIndex),
    Register(Register),
}

impl std::fmt::Display for ModRmTarget {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            Self::Memory(idx) => write!(f, "{}", idx),
            Self::Register(reg) => write!(f, "{}", reg),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Copy)]
/// Memory displacements are concrete signed versions of Byte and Word operands.
/// Encodes either Byte- or Word-sized operands.
/// Generally, a [`Displacement`] is the result of a ModRM byte parse and
/// usually really is the signed displacement of a [`MemoryIndex`] for memory
/// operations.
/// Although, some instructions use this parsed displacement as an unsigned
/// pointer, hence the [`Displacement`] will be cast to [`Pointer16`], when this
/// is the case.
pub enum Displacement {
    IByte(IByte),
    IWord(IWord),
}

impl fmt::LowerHex for Displacement {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::IByte(b) => fmt::LowerHex::fmt(b, f),
            Self::IWord(v) => fmt::LowerHex::fmt(v, f),
        }
    }
}

impl std::fmt::Display for Displacement {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            Self::IByte(b) => {
                if *b > 0 {
                    write!(f, " + {:#x}", b)
                } else {
                    write!(f, " - {:#x}", b * -1)
                }
            }
            Self::IWord(w) => {
                if *w > 0 {
                    write!(f, " + {:#x}", w)
                } else {
                    write!(f, " - {:#x}", w * -1)
                }
            }
        }
    }
}

/// A memory index operand is usually created by ModRM bytes or words.
/// e.g. [bx+si]
#[derive(Debug, Clone, PartialEq, Eq, Copy)]
pub struct MemoryIndex {
    pub base: Option<Register>,
    pub index: Option<Register>,
    pub displacement: Option<Displacement>,
}

impl fmt::Display for MemoryIndex {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match &self.base {
            Some(base) => match &self.index {
                Some(index) => match &self.displacement {
                    Some(displacement) => {
                        write!(f, "[{} + {}{}]", base, index, displacement)
                    }
                    None => write!(f, "[{} + {}]", base, index),
                },
                None => match &self.displacement {
                    Some(displacement) => write!(f, "[{}{}]", base, displacement),
                    None => write!(f, "[{}]", base),
                },
            },
            None => match &self.index {
                Some(index) => match &self.displacement {
                    Some(displacement) => write!(f, "[{}{}]", index, displacement),
                    None => write!(f, "[{}]", index),
                },
                None => match &self.displacement {
                    Some(displacement) => write!(f, "[{:#x}]", displacement),
                    None => panic!("Memory Index without base, index and displacement"),
                },
            },
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Copy)]
/// 16-bit pointer for access, usually with a [`SegmentRegister`]  as segment
/// and [`Pointer16`] as offset.
/// Generally, this type only gets constructed in rare scenarios, when the
/// [`Displacement`] of ModRM byte is used as a raw pointer.
pub struct Pointer16 {
    pub word: Word,
}

impl std::fmt::Display for Pointer16 {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "ptr [{:#04x}]", self.word)
    }
}

impl TryFrom<ModRmTarget> for Pointer16 {
    type Error = DisasmError;

    fn try_from(target: ModRmTarget) -> Result<Self, Self::Error> {
        match target {
            ModRmTarget::Memory(mem) => match mem.displacement {
                Some(disp) => match disp {
                    Displacement::IWord(word) => Ok(Pointer16 { word: word as Word }),
                    _ => {
                        return Err(DisasmError::IllegalOperand(
                            "Tried to construct Pointer16 with Byte, when a Word is expected"
                                .into(),
                        ));
                    }
                },
                _ => {
                    return Err(DisasmError::IllegalOperand("Tried to construct Pointer16 with Register, when a Displacement was expected".into()));
                }
            },
            _ => {
                return Err(DisasmError::IllegalOperand(
                    "Tried to construct Pointer16 with Register, when a MemoryIndex expected"
                        .into(),
                ));
            }
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Copy)]
/// 32-bit segment:offset pointer for long jumps.
/// Both [`Word`]s are immediately encoded after the instruction
pub struct Pointer32 {
    pub raw: DWord,
    pub segment: Word,
    pub offset: Word,
}

impl std::fmt::Display for Pointer32 {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{:#04x}:{:#04x}", self.segment, self.offset)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn add_byte_byte() {
        let a = ImmediateOperand::Byte(5);
        let b = ImmediateOperand::Byte(7);
        assert_eq!(a + b, ImmediateOperand::Byte(5 + 7))
    }

    #[test]
    fn add_word_word() {
        let a = ImmediateOperand::Word(5);
        let b = ImmediateOperand::Word(7);
        assert_eq!(a + b, ImmediateOperand::Word(5 + 7))
    }

    #[test]
    fn add_byte_word() {
        let a = ImmediateOperand::Byte(5);
        let b = ImmediateOperand::Word(7);
        assert_eq!(b + a, ImmediateOperand::Word(5 + 7))
    }

    #[test]
    fn sub_byte_word() {
        let a = ImmediateOperand::Byte(5);
        let b = ImmediateOperand::Word(7);
        assert_eq!(b - a, ImmediateOperand::Word(7 - 5))
    }

    #[test]
    fn test_msb() {
        let pos = ImmediateOperand::Byte(1 << 4);
        let neg = ImmediateOperand::Byte(1 << 7);
        assert_eq!(pos.msb(), false);
        assert_eq!(neg.msb(), true);
    }

    #[test]
    fn test_as_word() {
        let b: u8 = 5;
        let byte = ImmediateOperand::Byte(b);
        let word = ImmediateOperand::Word(b as Word);
        assert_eq!(byte.as_word(), word);
        assert_eq!(word, word);
    }

    #[test]
    fn test_flip_sign_neg_to_pos() {
        let b = 0 << 2;
        let byte = ImmediateOperand::Byte((1 << 7) | b);
        let word = ImmediateOperand::Word((1 << 15) | b as Word);
        assert_eq!(byte.flip_sign(), ImmediateOperand::Byte(b));
        assert_eq!(word.flip_sign(), ImmediateOperand::Word(b as Word));
    }

    #[test]
    fn test_flip_sign_pos_to_neg() {
        let b = 1 << 2;
        let byte = ImmediateOperand::Byte(b);
        let word = ImmediateOperand::Word(b as Word);
        assert_eq!(byte.flip_sign(), ImmediateOperand::Byte((1 << 7) | b));
        assert_eq!(
            word.flip_sign(),
            ImmediateOperand::Word((1 << 15) | b as Word)
        );
    }

    #[test]
    fn test_sign_extend() {
        let byte = ImmediateOperand::Byte(0b01010101);
        let word = ImmediateOperand::Word(0b0000000001010101);
        assert_eq!(byte.sign_extend(), word);
    }

    #[test]
    fn test_sign_extend_neg() {
        let byte = ImmediateOperand::Byte(1 << 7);
        let word = ImmediateOperand::Word(1 << 15);
        assert_eq!(byte.sign_extend(), word);
    }
}