use bytemuck::Pod;
use yggdrasil_rt::{
    mem::MappingFlags,
    process::{auxv, AuxValue},
};

use crate::{error::Error, mapping::Mapping};

struct ArgPlacer<'a> {
    buffer: &'a mut [u8],
    position: usize,
}

impl<'a> ArgPlacer<'a> {
    fn new(buffer: &'a mut [u8]) -> Self {
        Self {
            buffer,
            position: 0,
        }
    }

    fn align(&mut self, align: usize) -> Result<(), Error> {
        debug_assert!(align.is_power_of_two());
        let aligned = (self.position + align - 1) & !(align - 1);
        if aligned > self.buffer.len() {
            todo!()
        }
        self.position = aligned;
        Ok(())
    }

    fn put_str(&mut self, s: &str) -> Result<usize, Error> {
        if self.position + s.len() >= self.buffer.len() {
            todo!()
        }
        let off = self.position;
        self.buffer[off..off + s.len()].copy_from_slice(s.as_bytes());
        self.buffer[off + s.len()] = 0;
        self.position += s.len() + 1;
        Ok(off)
    }

    fn put<T: Pod>(&mut self, v: &T) -> Result<usize, Error> {
        if self.position + size_of::<T>() > self.buffer.len() {
            todo!()
        }
        let off = self.position;
        self.buffer[off..off + size_of::<T>()].copy_from_slice(bytemuck::bytes_of(v));
        self.position += size_of::<T>();
        Ok(off)
    }

    fn put_aux_array(&mut self, s: &[AuxValue]) -> Result<usize, Error> {
        self.align(size_of::<u64>())?;
        let off = self.position;
        for item in s {
            self.put(&item.tag)?;
            self.put(&item.val)?;
        }
        self.put(&auxv::NULL)?;
        self.put(&0u64)?;
        Ok(off)
    }

    fn put_ptr_array(&mut self, s: &[usize]) -> Result<usize, Error> {
        self.align(size_of::<usize>())?;
        let off = self.position;
        for item in s {
            self.put(item)?;
        }
        self.put(&0usize)?;
        Ok(off)
    }
}

pub fn build_argument(args: &[String], binary: &str, auxv: &[AuxValue]) -> Result<usize, Error> {
    let mut buffer = Mapping::new(0x1000, MappingFlags::WRITE)?;
    let arg_base = buffer.as_ptr().addr();
    let mut placer = ArgPlacer::new(&mut buffer[..]);

    let mut argv = vec![];
    let mut envp = vec![];

    for arg in args {
        argv.push(placer.put_str(arg)? + arg_base);
    }
    for (key, value) in std::env::vars() {
        let env_string = format!("{key}={value}");
        envp.push(placer.put_str(&env_string)? + arg_base);
    }

    // TODO env
    let argv = placer.put_ptr_array(&argv)? + arg_base;
    let envp = placer.put_ptr_array(&envp)? + arg_base;
    let auxv = placer.put_aux_array(auxv)? + arg_base;
    let real_program = placer.put_str(binary)? + arg_base;

    placer.align(size_of::<usize>())?;

    let argument = placer.position + arg_base;

    placer.put(&argv)?;
    placer.put(&envp)?;
    placer.put(&auxv)?;
    placer.put(&real_program)?;

    buffer.leak();

    Ok(argument)
}
