use abi::{
    error::Error,
    mem::{MappingFlags, MappingSource},
    process::ExitCode,
};

use super::{TlsImage, TlsInfo};

// Variant I TLS layout:
//
// [ TCB ] [ pad to p_align ] [ MODULE ] [ MODULE ] [ MODULE ] ...
// |                          |          |          |
// |                          |          |          |
// tp                        off1       off2       off3
//
// TCB size is fixed at 2 * sizeof(usize).

/// Creates a new TLS image in the process memory, copying data from the TLS master copy (if any).
/// Returns the resulting thread pointer.
pub fn clone_tls(image: &TlsImage) -> Result<TlsInfo, Error> {
    const TCB_SIZE: usize = size_of::<usize>() * 2;

    if !image.align.is_power_of_two() {
        crate::debug_trace!(
            Debug,
            "TLS layout not aligned to a power of two: {}",
            image.align
        );
        unsafe { crate::sys::exit_process(ExitCode::Exited(1)) };
    }
    if image.align > 0x1000 {
        crate::debug_trace!(
            Warn,
            "TODO: TLS alignment larger than a page size is not supported"
        );
        unsafe { crate::sys::exit_process(ExitCode::Exited(1)) };
    }

    // TCB size, padded to align. Also the start of the first module
    let tcb_aligned_size = (TCB_SIZE + image.align - 1) & !(image.align - 1);
    let full_size = tcb_aligned_size + image.full_size;
    let page_aligned_size = (full_size + 0xFFF) & !0xFFF;

    let base = unsafe {
        crate::sys::map_memory(
            None,
            page_aligned_size,
            MappingFlags::WRITE,
            &MappingSource::Anonymous,
        )
    }?;

    let ptr = core::ptr::with_exposed_provenance_mut::<u8>(base);

    let module_data =
        unsafe { core::slice::from_raw_parts_mut(ptr.add(tcb_aligned_size), image.full_size) };
    let tcb = unsafe { core::slice::from_raw_parts_mut(ptr, TCB_SIZE) };

    // Clone the Master Copy, if any
    let data_size = if let Some(master_copy) = image.master_copy {
        let source = unsafe { master_copy.as_ref() };
        module_data[..source.len()].copy_from_slice(source);
        source.len()
    } else {
        0
    };

    // Zero the rest
    module_data[data_size..image.full_size].fill(0);

    // Place the self-pointer to point at the TP itself
    let tp = base;
    assert_eq!(
        tp % size_of::<usize>(),
        0,
        "TLS self-pointer should at least be word-aligned"
    );
    tcb[..size_of::<usize>()].copy_from_slice(&tp.to_ne_bytes());

    // Zero the TCB after the self-pointer
    tcb[size_of::<usize>()..].fill(0);

    crate::debug_trace!(Debug, "TLS: base={:#x}, tp={:#x}", base, tp);

    Ok(TlsInfo {
        base,
        tp,
        module0_offset: Some(tcb_aligned_size),
    })
}

pub(super) fn get_tcb_raw(tp: usize) -> *mut u8 {
    core::ptr::with_exposed_provenance_mut(tp)
}
