use std::collections::HashMap;

use crate::{object::Object, thread_local::TlsSegment};

use super::{TlsLayout, TlsLayoutBuilder};

pub struct TlsLayoutImpl;

impl TlsLayoutBuilder for TlsLayoutImpl {
    fn build(root: &Object, libraries: &HashMap<u32, Object>) -> Option<(TlsLayout, usize)> {
        const TCB_SIZE: usize = size_of::<usize>() * 2;

        // Variant I doesn't use global align for modules. Instead, it aligns each of those
        // individually

        let mut segments = vec![];
        let mut total_size = TCB_SIZE;
        let mut first_module_align = 0;

        if let Some(tls) = root.tls.as_ref() {
            let align = tls.p_align as usize;
            assert!(align.is_power_of_two());

            if first_module_align == 0 {
                first_module_align = align;
            }

            total_size = (total_size + align - 1) & !(align - 1);

            segments.push(TlsSegment {
                offset: total_size,
                size: tls.p_memsz as _,
                module_id: 1,
            });

            total_size += tls.p_memsz as usize;
        }

        for (&object_id, lib) in libraries {
            let Some(tls) = lib.tls.as_ref() else {
                continue;
            };
            let align = tls.p_align as usize;
            assert!(align.is_power_of_two());

            if first_module_align == 0 {
                first_module_align = align;
            }

            total_size = (total_size + align - 1) & !(align - 1);

            segments.push(TlsSegment {
                offset: total_size,
                size: tls.p_memsz as _,
                module_id: object_id + 1,
            });

            total_size += tls.p_memsz as usize;
        }

        if first_module_align == 0 {
            assert!(segments.is_empty());
            None
        } else {
            let prefix_len = (TCB_SIZE + first_module_align - 1) & !(first_module_align - 1);
            Some((
                TlsLayout {
                    segments,
                    total_size,
                    tp_offset: 0,
                    prefix_len,
                },
                first_module_align,
            ))
        }
    }
}
