use crate::io::BinaryStream;
use crate::search::{SectionHelper, SearchSection};
use crate::error::{Error, Result};

pub const WASM_MAGIC: u32 = 0x6D736100;
pub const WASM_VERSION: u32 = 1;

#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WasmSectionId {
    Custom = 0,
    Type = 1,
    Import = 2,
    Function = 3,
    Table = 4,
    Memory = 5,
    Global = 6,
    Export = 7,
    Start = 8,
    Element = 9,
    Code = 10,
    Data = 11,
    DataCount = 12,
    Unknown = 255,
}

impl From<u8> for WasmSectionId {
    fn from(v: u8) -> Self {
        match v {
            0 => Self::Custom,
            1 => Self::Type,
            2 => Self::Import,
            3 => Self::Function,
            4 => Self::Table,
            5 => Self::Memory,
            6 => Self::Global,
            7 => Self::Export,
            8 => Self::Start,
            9 => Self::Element,
            10 => Self::Code,
            11 => Self::Data,
            12 => Self::DataCount,
            _ => Self::Unknown,
        }
    }
}

#[derive(Debug, Clone)]
struct WasmSection {
    _id: WasmSectionId,
    size: usize,
    offset: u64,
}

#[derive(Debug, Clone, Default)]
struct WasmDataSegment {
    offset: i64,
    size: usize,
    data_offset: u64,
}

pub struct Wasm {
    pub stream: BinaryStream,
    pub is_32bit: bool,
    code_section: Option<WasmSection>,
    data_segments: Vec<WasmDataSegment>,
}

impl Wasm {
    pub fn new(data: Vec<u8>) -> Result<Self> {
        let mut wasm = Self {
            stream: BinaryStream::new(data),
            is_32bit: true,
            code_section: None,
            data_segments: Vec::new(),
        };
        wasm.stream.is_32bit = true;
        wasm.load()?;
        Ok(wasm)
    }

    fn load(&mut self) -> Result<()> {
        self.stream.set_position(0);
        let magic = self.stream.read_u32()?;
        if magic != WASM_MAGIC {
            return Err(Error::InvalidFormat("Invalid WebAssembly magic".into()));
        }
        let version = self.stream.read_u32()?;
        if version != WASM_VERSION {
            return Err(Error::InvalidFormat(format!("Unsupported WASM version: {}", version)));
        }

        let data_len = self.stream.len();
        while self.stream.position() < data_len {
            let section_id_byte = self.stream.read_u8()?;
            let section_id = WasmSectionId::from(section_id_byte);
            let section_size = self.read_leb128_unsigned()?;
            let section_offset = self.stream.position();

            match section_id {
                WasmSectionId::Code => {
                    self.code_section = Some(WasmSection {
                        _id: section_id,
                        size: section_size,
                        offset: section_offset,
                    });
                }
                WasmSectionId::Data => {
                    self.parse_data_section(section_offset, section_size)?;
                }
                WasmSectionId::Custom => {
                    let name_len = self.read_leb128_unsigned()?;
                    let _name = self.stream.read_bytes(name_len)?;
                }
                _ => {}
            }

            self.stream.set_position(section_offset + section_size as u64);
        }

        Ok(())
    }

    fn parse_data_section(&mut self, offset: u64, _size: usize) -> Result<()> {
        self.stream.set_position(offset);
        let num_segments = self.read_leb128_unsigned()?;

        for _ in 0..num_segments {
            let flags = self.read_leb128_unsigned()?;
            let mut segment = WasmDataSegment::default();

            match flags {
                0 => {
                    let opcode = self.stream.read_u8()?;
                    if opcode == 0x41 {
                        segment.offset = self.read_leb128_signed()?;
                    }
                    let _end = self.stream.read_u8()?;
                }
                1 => {
                    segment.offset = 0;
                }
                2 => {
                    let _memory_index = self.read_leb128_unsigned()?;
                    let opcode = self.stream.read_u8()?;
                    if opcode == 0x41 {
                        segment.offset = self.read_leb128_signed()?;
                    }
                    let _end = self.stream.read_u8()?;
                }
                _ => {}
            }

            segment.size = self.read_leb128_unsigned()?;
            segment.data_offset = self.stream.position();
            self.stream.set_position(self.stream.position() + segment.size as u64);
            self.data_segments.push(segment);
        }

        Ok(())
    }

    fn read_leb128_unsigned(&mut self) -> Result<usize> {
        let mut result = 0usize;
        let mut shift = 0;
        loop {
            let byte = self.stream.read_u8()?;
            result |= ((byte & 0x7F) as usize) << shift;
            if (byte & 0x80) == 0 {
                break;
            }
            shift += 7;
        }
        Ok(result)
    }

    fn read_leb128_signed(&mut self) -> Result<i64> {
        let mut result = 0i64;
        let mut shift = 0;
        let mut byte;
        loop {
            byte = self.stream.read_u8()?;
            result |= ((byte & 0x7F) as i64) << shift;
            shift += 7;
            if (byte & 0x80) == 0 {
                break;
            }
        }
        if shift < 64 && (byte & 0x40) != 0 {
            result |= !0i64 << shift;
        }
        Ok(result)
    }

    pub fn map_vatr(&self, addr: u64) -> Result<u64> {
        let addr_i = addr as i64;
        for segment in &self.data_segments {
            if addr_i >= segment.offset && addr_i < segment.offset + segment.size as i64 {
                return Ok(segment.data_offset + (addr_i - segment.offset) as u64);
            }
        }
        Ok(addr)
    }

    pub fn map_rtva(&self, offset: u64) -> u64 {
        for segment in &self.data_segments {
            if offset >= segment.data_offset && offset < segment.data_offset + segment.size as u64 {
                return (segment.offset + (offset as i64 - segment.data_offset as i64)) as u64;
            }
        }
        offset
    }

    pub fn get_section_helper(&self, method_count: usize, type_definitions_count: usize, metadata_usages_count: usize, image_count: usize, version: f64) -> SectionHelper<'_> {
        let mut exec_list = Vec::new();
        let mut data_list = Vec::new();
        let mut all = Vec::new();

        if let Some(code) = &self.code_section {
            let s = SearchSection::new(code.offset, code.offset + code.size as u64, code.offset, code.offset + code.size as u64);
            all.push(s.clone());
            exec_list.push(s);
        }

        for segment in &self.data_segments {
            let s = SearchSection::new(
                segment.data_offset,
                segment.data_offset + segment.size as u64,
                segment.offset as u64,
                segment.offset as u64 + segment.size as u64,
            );
            all.push(s.clone());
            data_list.push(s);
        }

        let bss = data_list.clone();

        SectionHelper::new(
            self.stream.data(),
            self.is_32bit,
            version,
            false,
            all,
            data_list,
            exec_list,
            bss,
            method_count,
            type_definitions_count,
            metadata_usages_count,
            image_count,
        )
    }

    pub fn check_dump(&self) -> bool {
        false
    }

    pub fn get_rva(&self, pointer: u64) -> u64 {
        pointer
    }
}
