//! Private binary encoding helpers for `.beamr_native` AOT cache bundles.

use super::aot::{AotError, MAGIC, VERSION};
use super::{RootLocation, StackMapEntry};
use crate::atom::{Atom, AtomTable};
use crate::loader::load_beam_chunks;
use std::collections::HashMap;

pub(crate) struct DecodedBundle {
    pub(crate) module_checksum: u64,
    pub(crate) module_bytes: Vec<u8>,
    pub(crate) entries: Vec<(Atom, u8, Vec<StackMapEntry>)>,
}

impl DecodedBundle {
    pub(crate) fn read(bytes: &[u8], target_hash: u64) -> Result<Self, AotError> {
        let mut reader = Reader::new(bytes);
        let magic = reader.read_exact(MAGIC.len())?;
        if magic != MAGIC {
            return Err(AotError::InvalidMagic);
        }
        let version = reader.read_u8()?;
        if version != VERSION {
            return Err(AotError::UnsupportedVersion(version));
        }
        let actual_target = reader.read_u64()?;
        if actual_target != target_hash {
            return Err(AotError::TargetMismatch {
                expected: target_hash,
                actual: actual_target,
            });
        }
        let module_checksum = reader.read_u64()?;
        let _module_index = reader.read_u32()?;
        let module_bytes = reader.read_bytes()?;
        let atom_table = AtomTable::with_common_atoms();
        let parsed = load_beam_chunks(&module_bytes, &atom_table).map_err(AotError::Load)?;
        let functions_by_index: HashMap<u32, Atom> = parsed
            .exports
            .iter()
            .map(|export| (export.function.index(), export.function))
            .collect();
        let entry_count = reader.read_u32()? as usize;
        let mut entries = Vec::with_capacity(entry_count);
        for _ in 0..entry_count {
            let function_index = reader.read_u32()?;
            let function = functions_by_index
                .get(&function_index)
                .copied()
                .ok_or_else(|| {
                    AotError::Malformed(format!("unknown function atom index {function_index}"))
                })?;
            let arity = reader.read_u8()?;
            let stack_maps = reader.read_stack_maps()?;
            entries.push((function, arity, stack_maps));
        }
        if !reader.is_empty() {
            return Err(AotError::Malformed(
                "trailing bytes after AOT bundle".into(),
            ));
        }
        Ok(Self {
            module_checksum,
            module_bytes,
            entries,
        })
    }
}

pub(crate) fn write_atom(output: &mut Vec<u8>, atom: Atom) {
    write_u32(output, atom.index());
}

pub(crate) fn write_stack_maps(output: &mut Vec<u8>, stack_maps: &[StackMapEntry]) {
    write_u32(output, stack_maps.len() as u32);
    for entry in stack_maps {
        write_u32(output, entry.offset_from_entry);
        write_u32(output, entry.live_roots.len() as u32);
        for root in &entry.live_roots {
            match root {
                RootLocation::Register(register) => {
                    output.push(0);
                    write_u16(output, *register);
                }
                RootLocation::StackSlot(slot) => {
                    output.push(1);
                    write_i32(output, *slot);
                }
            }
        }
    }
}

pub(crate) fn write_bytes(output: &mut Vec<u8>, bytes: &[u8]) {
    write_u64(output, bytes.len() as u64);
    output.extend_from_slice(bytes);
}

pub(crate) fn write_u64(output: &mut Vec<u8>, value: u64) {
    output.extend_from_slice(&value.to_le_bytes());
}

fn write_u16(output: &mut Vec<u8>, value: u16) {
    output.extend_from_slice(&value.to_le_bytes());
}

pub(crate) fn write_u32(output: &mut Vec<u8>, value: u32) {
    output.extend_from_slice(&value.to_le_bytes());
}

fn write_i32(output: &mut Vec<u8>, value: i32) {
    output.extend_from_slice(&value.to_le_bytes());
}

struct Reader<'a> {
    bytes: &'a [u8],
    offset: usize,
}

impl<'a> Reader<'a> {
    const fn new(bytes: &'a [u8]) -> Self {
        Self { bytes, offset: 0 }
    }

    fn is_empty(&self) -> bool {
        self.offset == self.bytes.len()
    }

    fn read_exact(&mut self, len: usize) -> Result<&'a [u8], AotError> {
        let end = self
            .offset
            .checked_add(len)
            .ok_or_else(|| AotError::Malformed("bundle offset overflow".into()))?;
        let slice = self
            .bytes
            .get(self.offset..end)
            .ok_or_else(|| AotError::Malformed("truncated AOT bundle".into()))?;
        self.offset = end;
        Ok(slice)
    }

    fn read_u8(&mut self) -> Result<u8, AotError> {
        self.read_exact(1).map(|bytes| bytes[0])
    }

    fn read_u32(&mut self) -> Result<u32, AotError> {
        let bytes = self.read_exact(4)?;
        Ok(u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
    }

    fn read_u64(&mut self) -> Result<u64, AotError> {
        let bytes = self.read_exact(8)?;
        Ok(u64::from_le_bytes([
            bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7],
        ]))
    }

    fn read_i32(&mut self) -> Result<i32, AotError> {
        let bytes = self.read_exact(4)?;
        Ok(i32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
    }

    fn read_u16(&mut self) -> Result<u16, AotError> {
        let bytes = self.read_exact(2)?;
        Ok(u16::from_le_bytes([bytes[0], bytes[1]]))
    }

    fn read_bytes(&mut self) -> Result<Vec<u8>, AotError> {
        let len = usize::try_from(self.read_u64()?)
            .map_err(|_| AotError::Malformed("byte section length overflows usize".into()))?;
        self.read_exact(len).map(<[u8]>::to_vec)
    }

    fn read_stack_maps(&mut self) -> Result<Vec<StackMapEntry>, AotError> {
        let count = self.read_u32()? as usize;
        let mut entries = Vec::with_capacity(count);
        for _ in 0..count {
            let offset_from_entry = self.read_u32()?;
            let root_count = self.read_u32()? as usize;
            let mut live_roots = Vec::with_capacity(root_count);
            for _ in 0..root_count {
                let tag = self.read_u8()?;
                let root = match tag {
                    0 => RootLocation::Register(self.read_u16()?),
                    1 => RootLocation::StackSlot(self.read_i32()?),
                    other => {
                        return Err(AotError::Malformed(format!(
                            "unknown stack-map root tag {other}"
                        )));
                    }
                };
                live_roots.push(root);
            }
            entries.push(StackMapEntry {
                offset_from_entry,
                live_roots,
            });
        }
        Ok(entries)
    }
}
