use std::collections::HashMap;

use cfg_if::cfg_if;
use yggdrasil_rt::mem::MappingFlags;

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

cfg_if! {
    if #[cfg(rust_analyzer)] {
        mod variant1;
        mod variant2;
        pub use variant2::TlsLayoutImpl;
    } else if #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] {
        mod variant2;
        pub use variant2::TlsLayoutImpl;
    } else if #[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))] {
        mod variant1;
        pub use variant1::TlsLayoutImpl;
    }
}

pub trait TlsLayoutBuilder {
    fn build(root: &Object, libraries: &HashMap<u32, Object>) -> Option<(TlsLayout, usize)>;
}

#[derive(Debug)]
pub struct TlsLayout {
    pub segments: Vec<TlsSegment>,
    pub total_size: usize,
    pub tp_offset: usize,
    pub prefix_len: usize,
}

#[derive(Debug)]
pub struct TlsSegment {
    pub offset: usize,
    pub size: usize,
    pub module_id: u32,
}

pub struct TlsImage {
    // TLS master image that will be passed to the program
    pub data: Mapping,
    pub size: usize,
    pub align: usize,
    pub tp_offset: usize,
}

impl TlsLayout {
    pub fn offset(&self, module_id: u32, offset: usize) -> Option<usize> {
        let segment = self
            .segments
            .iter()
            .find(|seg| seg.module_id == module_id)?;
        Some(segment.offset + offset)
    }
}

pub fn build_tls_image(
    root: &mut Object,
    libraries: &mut HashMap<u32, Object>,
) -> Result<Option<(TlsImage, TlsLayout)>, Error> {
    let Some((layout, align)) = TlsLayoutImpl::build(root, libraries) else {
        return Ok(None);
    };

    // Create a master image based on this layout
    let mut image_data = Mapping::new(layout.total_size, MappingFlags::WRITE)?;

    let tp_offset = layout.tp_offset;

    for segment in layout.segments.iter() {
        log::info!(
            "Load TLS segment: tlsoffset_i={:#x?} (-{:#x}), module_id={}",
            segment.offset..segment.offset + segment.size,
            layout.prefix_len,
            segment.module_id
        );

        let object = match segment.module_id {
            0 => unreachable!(),
            1 => &mut *root,
            _ => libraries.get_mut(&(segment.module_id - 1)).unwrap(),
        };
        object.tls_offset = Some(segment.offset);
        let load_offset = segment.offset - layout.prefix_len;
        object.place_tls_copy(&mut image_data[load_offset..load_offset + segment.size])?;
    }

    Ok(Some((
        TlsImage {
            data: image_data,
            size: layout.total_size,
            align,
            tp_offset,
        },
        layout,
    )))
}
