//! Code items

use crate as androscalpel_serializer;
use crate::{Error, ReadSeek, Result, Serializable, Sleb128, Uleb128};
use std::io::Write;

/// <https://source.android.com/docs/core/runtime/dex-format#code-item>
/// alignment: 4 bytes
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CodeItem {
    pub registers_size: u16,
    pub ins_size: u16,
    pub outs_size: u16,
    // pub tries_size: u16,
    /// 0 if no debug info, else offset to a [`crate::DebugInfoItem`].
    pub debug_info_off: u32,
    // pub insns_size: u32,
    pub insns: Vec<u16>,
    // pub padding: Vec<u8>,
    /// try items must refere to non overlapping range an order from low to hight addresses.
    pub tries: Vec<TryItem>,
    pub handlers: Option<EncodedCatchHandlerList>,
}

impl CodeItem {
    pub fn tries_size_field(&self) -> u16 {
        self.tries.len() as u16
    }

    pub fn insns_size_field(&self) -> u32 {
        self.insns.len() as u32
    }

    pub fn sanity_check(&self) -> Result<()> {
        if self.tries.is_empty() && self.handlers.is_some() {
            return Err(Error::InconsistantStruct(
                "CodeItem cannot have a `handlers` value if `tries_size` is 0 (CodeItem.tries is empty)".into()
            ));
        }
        if !self.tries.is_empty() && self.handlers.is_none() {
            return Err(Error::InconsistantStruct(
                "CodeItem must have a `handlers` value if `tries_size` is not 0 (CodeItem.tries not empty)".into()
            ));
        }
        let mut max_addr = 0;
        for item in &self.tries {
            let addr = item.start_addr;
            if addr < max_addr {
                return Err(Error::InconsistantStruct(
                    "try_item in a code_item must be non overlapping and sorted from low to high address".into()
                ));
            }
            max_addr += addr + (item.insn_count as u32);
        }
        // Necessary? no in spec
        if max_addr > self.insns.len() as u32 {
            return Err(Error::InconsistantStruct(
                "found try_item whose block span outside of the insns array".into(),
            ));
        }

        for try_ in &self.tries {
            let handler = self
                .handlers
                .as_ref()
                .unwrap()
                .get_handler_at_offset(try_.handler_off)?;
            handler.sanity_check()?;
            for handler in &handler.handlers {
                // Necessary? no in spec
                if handler.addr.0 > self.insns.len() as u32 {
                    return Err(Error::InconsistantStruct(
                        "Found an handler whose address is outside of the insns array".into(),
                    ));
                }
            }
            if let Some(Uleb128(addr)) = handler.catch_all_addr {
                // Necessary? no in spec
                if addr > self.insns.len() as u32 {
                    return Err(Error::InconsistantStruct(
                        "Found a catch all handler whose address is outside of the insns array"
                            .into(),
                    ));
                }
            }
        }
        Ok(())
    }
}

impl Serializable for CodeItem {
    fn serialize(&self, output: &mut dyn Write) -> Result<()> {
        self.sanity_check().map_err(|err| match err {
            Error::InconsistantStruct(msg) => {
                Error::SerializationError(format!("Inconsistant CodeItem: {msg}"))
            }
            err => err,
        })?;
        self.registers_size.serialize(output)?;
        self.ins_size.serialize(output)?;
        self.outs_size.serialize(output)?;
        self.tries_size_field().serialize(output)?;
        self.debug_info_off.serialize(output)?;
        for insn in &self.insns {
            insn.serialize(output)?;
        }
        if !self.tries.is_empty() && self.insns.len() % 2 == 1 {
            0u16.serialize(output)?;
        }
        for item in &self.tries {
            item.serialize(output)?;
        }
        if let Some(ref handlers) = self.handlers {
            handlers.serialize(output)?;
        }
        Ok(())
    }

    fn deserialize(input: &mut dyn ReadSeek) -> Result<Self> {
        let registers_size = u16::deserialize(input)?;
        let ins_size = u16::deserialize(input)?;
        let outs_size = u16::deserialize(input)?;
        let tries_size = u16::deserialize(input)?;
        let debug_info_off = u32::deserialize(input)?;
        let insns_size = u32::deserialize(input)?;
        let mut insns = vec![];
        for _ in 0..insns_size {
            insns.push(u16::deserialize(input)?);
        }
        if tries_size != 0 && insns_size % 2 == 1 {
            let _ = u16::deserialize(input)?;
        }
        let mut tries = vec![];
        for _ in 0..tries_size {
            tries.push(TryItem::deserialize(input)?);
        }
        let handlers = if tries_size != 0 {
            Some(EncodedCatchHandlerList::deserialize(input)?)
        } else {
            None
        };
        Ok(Self {
            registers_size,
            ins_size,
            outs_size,
            debug_info_off,
            insns,
            tries,
            handlers,
        })
    }
    fn size(&self) -> usize {
        self.registers_size.size()
            + self.ins_size.size()
            + self.outs_size.size()
            + self.tries_size_field().size()
            + self.debug_info_off.size()
            + self.insns.iter().map(|val| val.size()).sum::<usize>()
            + if !self.tries.is_empty() && self.insns.len() % 2 == 1 {
                2
            } else {
                0
            }
            + self.tries.iter().map(|val| val.size()).sum::<usize>()
            + self.handlers.as_ref().map(|val| val.size()).unwrap_or(0)
    }
}

/// <https://source.android.com/docs/core/runtime/dex-format#type-item>
#[derive(Serializable, Clone, Copy, Debug, PartialEq, Eq)]
pub struct TryItem {
    /// Start address of the block of code covered. It's a count of 16-bit code unit to the
    /// start of the first covered instruction of the block.
    pub start_addr: u32,
    /// Number of 16-bit code unit covered by the entry.
    pub insn_count: u16,
    /// **Offset in bytes** from the start of the [`crate::EncodedCatchHandlerList`] to the
    /// [`crate::EncodedCatchHandler`] associated.
    pub handler_off: u16,
}

/// <https://source.android.com/docs/core/runtime/dex-format#encoded-catch-handlerlist>
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct EncodedCatchHandlerList {
    // pub size: Uleb128,
    pub list: Vec<EncodedCatchHandler>,
}

impl EncodedCatchHandlerList {
    pub fn size_field(&self) -> Uleb128 {
        Uleb128(self.list.len() as u32)
    }

    /// Return a reference to the [`crate::EncodedCatchHandler`] located at `offset` bytes after
    /// the begining of the [`crate::EncodedCatchHandlerList`]. Expected to be used to lookup
    /// the value refered to by [`crate::TryItem`]`.handler_off`.
    pub fn get_handler_at_offset(&self, offset: u16) -> Result<&EncodedCatchHandler> {
        let offset = offset as usize;
        let mut current_offset = 0;
        for handler in &self.list {
            if current_offset == offset {
                return Ok(handler);
            }
            current_offset += handler.size();
            if current_offset > offset {
                break;
            }
        }
        Err(Error::InconsistantStruct(format!(
            "Offset 0x{offset:x} does not match with the begining of a EncodedCatchHandler in this EncodedCatchHandlerList"
        )))
    }
}

impl Serializable for EncodedCatchHandlerList {
    fn serialize(&self, output: &mut dyn Write) -> Result<()> {
        self.size_field().serialize(output)?;
        for item in &self.list {
            item.serialize(output)?;
        }
        Ok(())
    }

    fn deserialize(input: &mut dyn ReadSeek) -> Result<Self> {
        let Uleb128(size) = Uleb128::deserialize(input)?;
        let mut list = vec![];
        for _ in 0..size {
            list.push(EncodedCatchHandler::deserialize(input)?);
        }
        Ok(Self {
            // size,
            list,
        })
    }

    fn size(&self) -> usize {
        self.size_field().size() + self.list.iter().map(|val| val.size()).sum::<usize>()
    }
}

/// <https://source.android.com/docs/core/runtime/dex-format#encoded-catch-handler>
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct EncodedCatchHandler {
    // pub size: Sleb128,
    /// List of handler, one by type caught, in the order of the type tests
    pub handlers: Vec<EncodedTypeAddrPair>,
    /// Bytecode address of the catchall handler.
    pub catch_all_addr: Option<Uleb128>,
}

impl EncodedCatchHandler {
    pub fn size_field(&self) -> Sleb128 {
        let sign = if self.catch_all_addr.is_some() { -1 } else { 1 };
        let len = self.handlers.len() as i32;
        //if len == 0 && self.catch_all_addr.is_none() {
        // Not good. Panic? Error?
        //}
        Sleb128(len * sign)
    }

    pub fn sanity_check(&self) -> Result<()> {
        if self.catch_all_addr.is_none() && self.handlers.is_empty() {
            return Err(Error::InconsistantStruct(
                "EncodedCatchHandler must have at least one handler or catch_all_addr defined"
                    .into(),
            ));
        }
        Ok(())
    }
}

impl Serializable for EncodedCatchHandler {
    fn serialize(&self, output: &mut dyn Write) -> Result<()> {
        self.sanity_check().map_err(|err| match err {
            Error::InconsistantStruct(msg) => {
                Error::SerializationError(format!("Inconsistant EncodedCatchHandler: {msg}"))
            }
            err => err,
        })?;
        self.size_field().serialize(output)?;
        for handler in &self.handlers {
            handler.serialize(output)?;
        }
        if let Some(catch_all_addr) = self.catch_all_addr {
            catch_all_addr.serialize(output)?;
        }
        Ok(())
    }
    fn deserialize(input: &mut dyn ReadSeek) -> Result<Self> {
        let Sleb128(size) = Sleb128::deserialize(input)?;
        let mut handlers = vec![];
        for _ in 0..size.abs() {
            handlers.push(EncodedTypeAddrPair::deserialize(input)?);
        }
        let catch_all_addr = if size <= 0 {
            Some(Uleb128::deserialize(input)?)
        } else {
            None
        };
        Ok(Self {
            handlers,
            catch_all_addr,
        })
    }
    fn size(&self) -> usize {
        self.size_field().size()
            + self.handlers.iter().map(|val| val.size()).sum::<usize>()
            + self
                .catch_all_addr
                .as_ref()
                .map(|val| val.size())
                .unwrap_or(0)
    }
}

/// <https://source.android.com/docs/core/runtime/dex-format#encoded-type-addr-pair>
#[derive(Serializable, Clone, Copy, Debug, PartialEq, Eq)]
pub struct EncodedTypeAddrPair {
    /// Index of the [`crate::TypeIdItem`] in `type_ids`
    pub type_idx: Uleb128,
    /// Bytecode address of the exception handler
    pub addr: Uleb128,
}

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

    const CODE_ITEM_RAW_1: &[u8] = &[
        0x04, 0x00, 0x01, 0x00, 0x02, 0x00, 0x01, 0x00, 0x3d, 0x25, 0x4f, 0x00, 0x15, 0x00, 0x00,
        0x00, 0x54, 0x30, 0x6c, 0x01, 0x71, 0x10, 0xb1, 0x0c, 0x00, 0x00, 0x28, 0x0e, 0x0d, 0x00,
        0x6e, 0x10, 0x26, 0x85, 0x00, 0x00, 0x0c, 0x01, 0x1a, 0x02, 0xcb, 0x15, 0x71, 0x20, 0xe3,
        0x05, 0x21, 0x00, 0x0a, 0x01, 0x38, 0x01, 0x03, 0x00, 0x0e, 0x00, 0x27, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x05, 0x00, 0x01, 0x00, 0x01, 0x01, 0xe9, 0x46, 0x06,
    ];

    const ENCODED_CATCH_HANDLER_LIST_1: &[u8] = &[0x01, 0x01, 0xe9, 0x46, 0x06];

    #[test]
    fn test_deserialize_code_item() {
        assert_eq!(
            CodeItem::deserialize_from_slice(CODE_ITEM_RAW_1).unwrap(),
            CodeItem {
                registers_size: 4,
                ins_size: 1,
                outs_size: 2,
                debug_info_off: 5186877,
                insns: vec![
                    0x3054, 0x16c, 0x1071, 0xcb1, 0x0, 0xe28, 0xd, 0x106e, 0x8526, 0x0, 0x10c,
                    0x21a, 0x15cb, 0x2071, 0x5e3, 0x21, 0x10a, 0x138, 0x3, 0xe, 0x27
                ],
                tries: vec![TryItem {
                    start_addr: 0,
                    insn_count: 5,
                    handler_off: 1,
                },],
                handlers: Some(EncodedCatchHandlerList {
                    list: vec![EncodedCatchHandler {
                        handlers: vec![EncodedTypeAddrPair {
                            type_idx: Uleb128(9065),
                            addr: Uleb128(6)
                        }],
                        catch_all_addr: None,
                    }]
                })
            }
        );
    }

    #[test]
    fn test_deserialize_catch_handler_list() {
        assert_eq!(
            EncodedCatchHandlerList::deserialize_from_slice(ENCODED_CATCH_HANDLER_LIST_1).unwrap(),
            EncodedCatchHandlerList {
                list: vec![EncodedCatchHandler {
                    handlers: vec![EncodedTypeAddrPair {
                        type_idx: Uleb128(9065),
                        addr: Uleb128(6)
                    }],
                    catch_all_addr: None,
                }]
            }
        );
    }
}