//! Various types of arguments that instructions can take, alongside encoding and decoding logic.

use crate::{
    instructions::{RegisterParseError, encode::Encode},
    prelude::Register,
};

/// A list of errors that can be returned when decoding instruction arguments.
#[derive(Debug)]
pub enum ArgsDecodeError {
    /// The register was not valid.
    InvalidRegister(u8),
}

impl From<RegisterParseError> for ArgsDecodeError {
    fn from(value: RegisterParseError) -> Self {
        match value {
            RegisterParseError::InvalidIndex(idx) => Self::InvalidRegister(idx),
            RegisterParseError::InvalidName(_) => Self::InvalidRegister(0xFF),
        }
    }
}

impl std::fmt::Display for ArgsDecodeError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::InvalidRegister(idx) => {
                write!(f, "invalid register index, got {idx:x}")?;
            }
        }

        Ok(())
    }
}

impl std::error::Error for ArgsDecodeError {}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
/// Used by instructions with 2 registers and an immediate argument.
pub struct ITypeArgs {
    pub immediate: u16,
    pub r1: Register,
    /// May not actually be used by some instructions taking an immediate e.g. LUI. This is solved by making the constructor take Options.
    pub r2: Register,
}

impl ITypeArgs {
    #[must_use]
    /// Creates a new [`ITypeArgs`]. If r1 or r2 is unset, they will be replaced with [`Register::NoReg`].
    pub fn new(immediate: u16, r1: Option<Register>, r2: Option<Register>) -> Self {
        let r1 = r1.unwrap_or_default();
        let r2 = r2.unwrap_or_default();

        Self { immediate, r1, r2 }
    }
}

impl Encode for ITypeArgs {
    /// Encodes an I-type instruction from its fields. These must have some unused high-order
    /// bits set to 0 else the bit shifting logic gets fucked.
    fn encode(self, opcode: u8) -> u32 {
        let opcode = u32::from(opcode);
        let r1 = self.r1 as u32;
        let dr = self.r2 as u32;
        let immediate = u32::from(self.immediate);

        (opcode << 26) | (r1 << 21) | (dr << 16) | immediate
    }
}

impl TryFrom<u32> for ITypeArgs {
    type Error = ArgsDecodeError;

    fn try_from(data: u32) -> Result<Self, Self::Error> {
        let r1 = ((data >> 21) & 0x1F) as u8;
        let r2 = ((data >> 16) & 0x1F) as u8;
        let immediate = data as u16;

        let r1 = r1.try_into()?;
        let r2 = r2.try_into()?;

        Ok(Self { immediate, r1, r2 })
    }
}

/// Used by instructions not using immediates (besides 5 bit shift values).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RTypeArgs {
    pub sr1: Register,
    pub sr2: Register,
    pub dr: Register,
    /// 5 bit shift amount.
    pub shamt: u8,
}

impl RTypeArgs {
    #[must_use]
    /// Creates a new [`RTypeArgs`]. If any registers are unset, they will be replaced with [`Register::NoReg`]. If `shamt` is unset, it will be set to 0.
    pub fn new(
        sr1: Option<Register>,
        sr2: Option<Register>,
        dr: Option<Register>,
        shamt: Option<u8>,
    ) -> Self {
        let sr1 = sr1.unwrap_or_default();
        let shamt = shamt.unwrap_or_default();
        let sr2 = sr2.unwrap_or_default();
        let dr = dr.unwrap_or_default();

        Self {
            sr1,
            sr2,
            dr,
            shamt,
        }
    }
}

impl Encode for RTypeArgs {
    /// Encodes an R-type instruction from its fields. These must have unused high-order
    /// bits set to 0 else the bit shifting logic is fucked.
    ///
    /// # Arguments
    ///
    /// - `shamt`: The amount to shift value (used only in shift instructions, otherwise 0).
    fn encode(self, opcode: u8) -> u32 {
        let opcode = u32::from(opcode);
        let sr1 = self.sr1 as u32;
        let sr2 = self.sr2 as u32;
        let dr = self.dr as u32;
        let shamt = u32::from(self.shamt);

        (opcode << 26) | (sr1 << 21) | (sr2 << 16) | (dr << 11) | (shamt << 6)
    }
}

impl TryFrom<u32> for RTypeArgs {
    type Error = ArgsDecodeError;

    fn try_from(data: u32) -> Result<Self, Self::Error> {
        let sr1 = ((data >> 21) & 0x1F) as u8;
        let sr2 = ((data >> 16) & 0x1F) as u8;
        let dr = ((data >> 11) & 0x1F) as u8;
        let shamt = ((data >> 6) & 0x1F) as u8;

        let sr1_reg = sr1.try_into()?;
        let sr2_reg = sr2.try_into()?;
        let dr_reg = dr.try_into()?;

        Ok(Self {
            sr1: sr1_reg,
            sr2: sr2_reg,
            dr: dr_reg,
            shamt,
        })
    }
}

#[macro_export]
macro_rules! args {
    // R-type arguments - allows omitting any field
    (R $(, $field:ident: $value:expr)* $(,)?) => {{
        let mut sr1: Option<Register> = None;
        let mut sr2: Option<Register> = None;
        let mut dr: Option<Register> = None;
        let mut shamt: Option<u8> = None;

        $(
            args!(@assign_r_option sr1, sr2, dr, shamt, $field, $value);
        )*

        RTypeArgs::new(sr1, sr2, dr, shamt)
    }};

    // I-type arguments - requires immediate, allows omitting registers
    (I, immediate: $immediate:expr $(, $field:ident: $value:expr)* $(,)?) => {{
        let mut r1: Option<Register> = None;
        let mut r2: Option<Register> = None;

        $(
            args!(@assign_i_option r1, r2, $field, $value);
        )*

        ITypeArgs::new($immediate, r1, r2)
    }};

    // Internal helpers (same as above for R-type)
    (@assign_r_option $sr1:ident, $sr2:ident, $dr:ident, $shamt:ident, sr1, $value:expr) => {
        $sr1 = Some($value);
    };
    (@assign_r_option $sr1:ident, $sr2:ident, $dr:ident, $shamt:ident, sr2, $value:expr) => {
        $sr2 = Some($value);
    };
    (@assign_r_option $sr1:ident, $sr2:ident, $dr:ident, $shamt:ident, dr, $value:expr) => {
        $dr = Some($value);
    };
    (@assign_r_option $sr1:ident, $sr2:ident, $dr:ident, $shamt:ident, shamt, $value:expr) => {
        $shamt = Some($value);
    };

    // Internal helpers for I-type (without immediate handling)
    (@assign_i_option $r1:ident, $r2:ident, r1, $value:expr) => {
        $r1 = Some($value);
    };
    (@assign_i_option $r1:ident, $r2:ident, r2, $value:expr) => {
        $r2 = Some($value);
    };
}
