//! Yggdrasil Runtime functions for handling Thread-Local Storage
use core::{
    ffi::c_void,
    ptr::{null_mut, NonNull},
};

use alloc::{boxed::Box, vec::Vec};

use abi::{
    error::Error,
    process::{auxv, AuxValue},
};

#[cfg(any(target_arch = "aarch64", rust_analyzer))]
mod aarch64;
#[cfg(any(target_arch = "aarch64", rust_analyzer))]
use aarch64 as imp;

#[cfg(any(target_arch = "x86_64", rust_analyzer))]
mod x86_64;
#[cfg(any(target_arch = "x86_64", rust_analyzer))]
use x86_64 as imp;

#[cfg(any(target_arch = "riscv64", rust_analyzer))]
mod riscv64;
#[cfg(any(target_arch = "riscv64", rust_analyzer))]
use riscv64 as imp;

#[cfg(any(target_arch = "aarch64", target_arch = "riscv64", rust_analyzer))]
mod variant1;
#[cfg(any(target_arch = "aarch64", target_arch = "riscv64", rust_analyzer))]
use variant1 as layout;

#[cfg(any(target_arch = "x86", target_arch = "x86_64", rust_analyzer))]
mod variant2;
#[cfg(any(target_arch = "x86", target_arch = "x86_64", rust_analyzer))]
use variant2 as layout;

pub use imp::{get_thread_pointer, set_thread_pointer};
pub use layout::clone_tls;

/// Describes a TLS master copy of this executable
#[derive(Debug)]
pub struct TlsImage {
    master_copy: Option<NonNull<[u8]>>,
    full_size: usize,
    align: usize,
    already_initialized: bool,
    // Pairs of dynamic module_id:offset from allocation base
    module_offsets: Vec<(usize, usize)>,
}

/// DTV table, containing dynamic TLS module mappings
pub struct Dtv {
    // DTV entries for TLS allocated by the dynamic linker/loader
    entries: Vec<*mut c_void>,
    // Entries for pthread_setspecific()-like behavior
    specific: Vec<*mut c_void>,
}

#[allow(missing_docs)]
pub struct TlsInfo {
    base: usize,
    tp: usize,
    module0_offset: Option<usize>,
}

struct TcbHeader {
    #[allow(unused)]
    self_pointer: usize,
    dtv_pointer: *mut Dtv,
}

impl TlsImage {
    /// Extracts information about the TLS master copy from aux vector
    pub fn from_auxv<'a, I: Iterator<Item = &'a AuxValue>>(auxv: I) -> Option<Self> {
        let mut master_copy = None;
        let mut data_size = None;
        let mut full_size = None;
        let mut align = None;
        let mut current_module = None;
        let mut module_offsets = Vec::new();
        let mut already_initialized = false;

        for item in auxv {
            match item.tag {
                auxv::TLS_MASTER_COPY => master_copy = Some(item.val as usize),
                auxv::TLS_DATA_SIZE => data_size = Some(item.val as usize),
                auxv::TLS_FULL_SIZE => full_size = Some(item.val as usize),
                auxv::TLS_ALIGN => align = Some(item.val as usize),
                auxv::TLS_MODULE_ID => current_module = Some(item.val as usize),
                auxv::TLS_MODULE_OFFSET => {
                    if let Some(module_id) = current_module.take() {
                        module_offsets.push((module_id, item.val as usize));
                    }
                }
                auxv::TLS_ALREADY_INITIALIZED => {
                    already_initialized = true;
                }
                _ => (),
            }
        }

        let full_size = full_size?;
        let align = align?;

        let master_copy = if let (Some(master_copy), Some(data_size)) = (master_copy, data_size) {
            assert_ne!(master_copy, 0);
            assert_ne!(data_size, 0);

            let base = unsafe {
                NonNull::new_unchecked(core::ptr::with_exposed_provenance_mut::<u8>(master_copy))
            };

            Some(NonNull::slice_from_raw_parts(base, data_size))
        } else {
            None
        };

        Some(Self {
            master_copy,
            full_size,
            align,
            module_offsets,
            already_initialized,
        })
    }
}

unsafe impl Sync for TlsImage {}

impl Dtv {
    fn new() -> Self {
        Self {
            entries: Vec::new(),
            specific: Vec::new(),
        }
    }

    #[inline]
    fn get_key(list: &[*mut c_void], key: usize) -> Option<*mut c_void> {
        if key == 0 || key > list.len() {
            None
        } else {
            Some(list[key - 1])
        }
    }

    #[inline]
    fn set_key(list: &mut [*mut c_void], key: usize, value: *mut c_void) -> bool {
        if key == 0 || key > list.len() {
            return false;
        }
        list[key - 1] = value;
        true
    }

    /// Returns a value assocaited with a given thread-specific key.
    ///
    /// # Panics
    ///
    /// Will panic if key == 0.
    /// Will panic if key is longer than the DTV itself.
    pub fn get_specific(&self, key: usize) -> Option<*mut c_void> {
        Self::get_key(&self.specific, key)
    }

    /// Sets a DTV entry for a thread-specific key.
    ///
    /// # Panics
    ///
    /// Will panic if key == 0.
    /// Will panic if key is longer than the DTV itself.
    pub fn set_specific(&mut self, key: usize, value: *mut c_void, grow: bool) {
        if self.try_set_specific(key, value, grow).is_err() {
            crate::debug_trace!(Debug, "Dtv::set_specific(): invalid key {key}, grow={grow}");
            panic!("Dtv::set_specific(): invalid key {key})")
        }
    }

    /// Sets a DTV entry for a thread-specific key.
    pub fn try_set_specific(
        &mut self,
        key: usize,
        value: *mut c_void,
        grow: bool,
    ) -> Result<(), Error> {
        if key > self.specific.len() && grow {
            self.specific.resize(key, null_mut());
        }
        if !Self::set_key(&mut self.specific, key, value) {
            return Err(Error::InvalidArgument);
        }
        Ok(())
    }

    /// Allocates a new thread-specific key in the DTV, filling it with a NULL value.
    pub fn new_specific(&mut self) -> usize {
        self.specific.push(null_mut());
        self.specific.len()
    }

    /// Returns a value associated with a given key.
    ///
    /// # Panics
    ///
    /// Will panic if key == 0.
    /// Will panic if key is larger than the DTV itself.
    pub fn get(&self, key: usize) -> *mut c_void {
        match Self::get_key(&self.entries, key) {
            Some(value) => value,
            None => {
                crate::debug_trace!(Debug, "Dtv::get(): out-of-bounds DTV key: {key}");
                panic!("Dtv::get(): out-of-bounds DTV key: {key}");
            }
        }
    }

    /// Sets a DTV entry, growing the DTV allocation if necessary
    ///
    /// # Panics
    ///
    /// Will panic if key == 0.
    pub fn set(&mut self, key: usize, value: *mut c_void) {
        crate::debug_trace!(Debug, "Dtv::set({key}, {value:p})");
        if key > self.entries.len() {
            self.entries.resize(key, null_mut());
        }
        if !Self::set_key(&mut self.entries, key, value) {
            crate::debug_trace!(Debug, "Dtv::set(): invalid key {key}");
            panic!("Dtv::set(): invalid key {key}");
        }
    }
}

/// Initializes thread-local storage for this thread from aux vector provided to the executable.
pub fn init_tls_from_auxv<'a, I: Iterator<Item = &'a AuxValue>>(
    auxv: I,
    force: bool,
) -> Result<Option<TlsImage>, Error> {
    let Some(tls_image) = TlsImage::from_auxv(auxv) else {
        return Ok(None);
    };

    if force || !tls_image.already_initialized {
        let tls = clone_tls(&tls_image)?;
        unsafe { set_thread_pointer(tls.tp) }?;
        setup_dtv(&tls_image, &tls)?;
    }

    Ok(Some(tls_image))
}

/// Initializes thread-local storage for this thread from a [TlsImage]. If no image is provided,
/// a dummy (platform-specific) TLS copy will be created.
pub fn init_tls(image: Option<&TlsImage>, force: bool) -> Result<(), Error> {
    let Some(image) = image else {
        crate::debug_trace!(Debug, "TODO: handle executables with no TLS");
        return Err(Error::NotImplemented);
    };

    if force || !image.already_initialized {
        let tls = clone_tls(image)?;
        unsafe { set_thread_pointer(tls.tp) }?;
        setup_dtv(image, &tls)?;
    }

    Ok(())
}

fn get_tcb_raw() -> *mut TcbHeader {
    let tp = imp::get_thread_pointer();
    layout::get_tcb_raw(tp).cast()
}

fn get_tcb_mut() -> &'static mut TcbHeader {
    unsafe { &mut *get_tcb_raw() }
}

fn setup_dtv(image: &TlsImage, tls_info: &TlsInfo) -> Result<(), Error> {
    #[cfg(any(target_arch = "riscv64", rust_analyzer))]
    const DTV_OFFSET: usize = 0x800;
    #[cfg(any(not(target_arch = "riscv64"), rust_analyzer))]
    const DTV_OFFSET: usize = 0;

    let dtv = get_dtv();
    // Executable itself
    // NOTE if module 1 is specified again by the dynamic loader, it will be overriden with
    //      what dynamic loader says
    if let Some(module0_offset) = tls_info.module0_offset {
        crate::debug_trace!(
            Debug,
            "DTV[1] = {:#x}",
            tls_info.base + module0_offset + DTV_OFFSET
        );
        dtv.set(
            1,
            core::ptr::without_provenance_mut(tls_info.base + module0_offset + DTV_OFFSET),
        );
    }

    if image.module_offsets.is_empty() {
        return Ok(());
    }
    for &(module_id, module_offset) in image.module_offsets.iter() {
        assert!(module_offset < image.full_size);
        crate::debug_trace!(
            Debug,
            "DTV[{module_id}] = {:#x}",
            tls_info.base + module_offset + DTV_OFFSET
        );
        dtv.set(
            module_id,
            core::ptr::with_exposed_provenance_mut(tls_info.base + module_offset + DTV_OFFSET),
        );
    }
    Ok(())
}

/// Gets or allocates a DTV for current thread.
pub fn get_dtv() -> &'static mut Dtv {
    let tcb = get_tcb_mut();
    if tcb.dtv_pointer.is_null() {
        let dtv = Box::new(Dtv::new());
        let dtv = Box::into_raw(dtv);
        tcb.dtv_pointer = dtv;
    }
    unsafe { &mut *tcb.dtv_pointer }
}

#[cfg(all(
    any(target_arch = "x86", target_arch = "x86_64", target_arch = "riscv64"),
    any(feature = "__tls_get_addr", rust_analyzer)
))]
#[no_mangle]
unsafe extern "C" fn __tls_get_addr(index: *mut usize) -> *mut c_void {
    let module_id = index.read();
    let offset = index.add(1).read();
    get_dtv().get(module_id).add(offset)
}
