use std::collections::HashMap;

use crate::object::Object;

use super::{TlsLayout, TlsLayoutBuilder, TlsSegment};

pub struct TlsLayoutImpl;

impl TlsLayoutBuilder for TlsLayoutImpl {
    fn build(root: &Object, libraries: &HashMap<u32, Object>) -> Option<(TlsLayout, usize)> {
        // Find max align
        let mut align = usize::MIN;
        if let Some(tls) = root.tls.as_ref() {
            align = align.max(tls.p_align as usize);
        }
        for tls in libraries.values().filter_map(|lib| lib.tls.as_ref()) {
            align = align.max(tls.p_align as usize);
        }

        // No TLS segment in any of the objects
        if align == usize::MIN {
            return None;
        }
        let align = align.max(size_of::<usize>());

        let mut total_size = 0;
        // Layout:
        //
        // [ Module N | padding ] ... [ Module 1 | padding ] [ padding | Module 0 ] [ TP ]

        let mut modules = vec![];

        if let Some(tls) = root.tls.as_ref() {
            let size = tls.p_memsz as usize;
            let aligned_size = (size + align - 1) & !(align - 1);

            // module_id, offset from tp, full size
            modules.push((0, size, aligned_size));

            total_size += aligned_size;
        }

        for (&id, lib) in libraries.iter() {
            let Some(tls) = lib.tls.as_ref() else {
                continue;
            };

            let size = tls.p_memsz as usize;
            let aligned_size = (size + align - 1) & !(align - 1);

            modules.push((id, total_size + aligned_size, aligned_size));
            total_size += aligned_size;
        }

        let tp_offset = total_size;

        // All the offsets calculated before are (tp_offset - x) offsets, recalculate those
        // before returning the module list
        for (_, off, _) in modules.iter_mut() {
            *off = tp_offset - *off;
        }

        let segments = modules
            .into_iter()
            .map(|(object_id, offset, size)| TlsSegment {
                offset,
                size,
                module_id: object_id + 1,
            })
            .collect();

        Some((
            TlsLayout {
                segments,
                total_size,
                tp_offset,
                prefix_len: 0,
            },
            align,
        ))
    }
}
