use std::{
    collections::{HashMap, HashSet},
    ffi::{CStr, CString},
    iter,
    path::{Path, PathBuf},
    rc::Rc,
};

use elf::abi::{STT_FUNC, STT_NOTYPE, STT_OBJECT, STT_TLS};

use crate::{
    config::Config,
    error::Error,
    object::{DynamicSymbol, Object},
    search,
    thread_local::{self, TlsImage, TlsLayout},
};

pub struct SymbolByAddress<'a> {
    pub object_name: &'a CStr,
    pub object_base: usize,
    pub symbol_name: &'a CStr,
    pub symbol_base: usize,
}

pub struct ExportedNormalSymbol {
    pub source: PathBuf,
    pub value: usize,
    pub size: usize,
    pub weak: bool,
    pub object_id: u32,
    pub c_name: CString,
}

pub struct ExportedTlsSymbol {
    pub source: PathBuf,
    // Offset from TLS block start
    pub offset: usize,
    // TLS module ID
    pub module_id: u32,
    pub weak: bool,
}

pub enum ExportedSymbol<'a> {
    Normal(&'a ExportedNormalSymbol),
    Tls(&'a ExportedTlsSymbol),
}

pub struct State {
    symbol_table: HashMap<Rc<str>, ExportedNormalSymbol>,
    tls_symbol_table: HashMap<Rc<str>, ExportedTlsSymbol>,
    undefined_references: HashMap<PathBuf, HashSet<Rc<str>>>,
    pub tls_layout: Option<TlsLayout>,
}

pub struct ObjectSet {
    pub root: Object,
    libraries: HashMap<u32, Object>,
    library_path_map: HashMap<PathBuf, u32>,

    // Linker state
    pub state: State,

    last_object_id: u32,
}

impl State {
    pub fn new() -> Self {
        Self {
            symbol_table: HashMap::new(),
            tls_symbol_table: HashMap::new(),
            undefined_references: HashMap::new(),
            tls_layout: None,
        }
    }

    pub fn export(
        &mut self,
        source: impl AsRef<Path>,
        sym: &DynamicSymbol,
        offset: isize,
        weak: bool,
        object_id: u32,
    ) {
        let source = source.as_ref().to_owned();
        match sym.raw.st_symtype() {
            STT_FUNC | STT_OBJECT | STT_NOTYPE => {
                let value: usize = (isize::try_from(sym.raw.st_value).unwrap() + offset)
                    .try_into()
                    .unwrap();
                let size: usize = sym.raw.st_size as usize;
                let c_name = CString::new(sym.name.as_ref()).unwrap();

                match self.symbol_table.get_mut(&sym.name) {
                    // Stronger binding exported
                    Some(export) if export.weak && !weak => {
                        *export = ExportedNormalSymbol {
                            source,
                            value,
                            size,
                            weak,
                            c_name,
                            object_id,
                        };
                    }
                    // Do nothing, already strong or already weak
                    Some(_) => (),
                    None => {
                        self.symbol_table.insert(
                            sym.name.clone(),
                            ExportedNormalSymbol {
                                source,
                                value,
                                size,
                                weak,
                                c_name,
                                object_id,
                            },
                        );
                    }
                }
            }
            STT_TLS => {
                let offset = sym.raw.st_value as usize;

                match self.tls_symbol_table.get_mut(&sym.name) {
                    Some(export) if export.weak && !weak => {
                        *export = ExportedTlsSymbol {
                            source,
                            offset,
                            module_id: object_id + 1,
                            weak,
                        };
                    }
                    Some(_) => (),
                    None => {
                        log::debug!(
                            "{:?}: TLS {:?} -> {}:{:#x}",
                            source,
                            sym.name,
                            object_id + 1,
                            offset
                        );
                        self.tls_symbol_table.insert(
                            sym.name.clone(),
                            ExportedTlsSymbol {
                                source,
                                offset,
                                module_id: object_id + 1,
                                weak,
                            },
                        );
                    }
                }
            }
            _ => todo!(),
        }
    }

    pub fn ensure_resolved(&mut self, source: impl AsRef<Path>, sym: &DynamicSymbol) {
        match sym.raw.st_symtype() {
            STT_FUNC | STT_OBJECT | STT_NOTYPE => {
                if !self.symbol_table.contains_key(&sym.name) {
                    self.undefined(source, &sym.name);
                }
            }
            STT_TLS => {
                if !self.tls_symbol_table.contains_key(&sym.name) {
                    self.undefined(source, &sym.name);
                }
            }
            _ => todo!(),
        }
    }

    pub fn lookup(&self, sym: &DynamicSymbol) -> Option<ExportedSymbol> {
        match sym.raw.st_symtype() {
            STT_FUNC | STT_OBJECT | STT_NOTYPE => {
                self.symbol_table.get(&sym.name).map(ExportedSymbol::Normal)
            }
            STT_TLS => self
                .tls_symbol_table
                .get(&sym.name)
                .map(ExportedSymbol::Tls),
            _ => todo!(),
        }
    }

    pub fn undefined(&mut self, source: impl AsRef<Path>, sym: &Rc<str>) {
        let source = source.as_ref();
        let list = if let Some(list) = self.undefined_references.get_mut(source) {
            list
        } else {
            self.undefined_references
                .try_insert(source.to_owned(), HashSet::new())
                .unwrap()
        };
        list.insert(sym.clone());
    }

    pub fn no_undefined_symbols(&self) -> Result<(), &HashMap<PathBuf, HashSet<Rc<str>>>> {
        if self.undefined_references.is_empty() {
            Ok(())
        } else {
            Err(&self.undefined_references)
        }
    }

    fn export_magic_symbol(&mut self, name: &str, value: usize) {
        let c_name = CString::new(name).unwrap();
        self.symbol_table.insert(
            name.into(),
            ExportedNormalSymbol {
                source: "".into(),
                weak: false,
                value,
                size: 0,
                c_name,
                object_id: u32::MAX,
            },
        );
    }

    pub fn add_magic_symbols(&mut self) {
        self.export_magic_symbol("dlopen", crate::builtins::dlopen as usize);
        self.export_magic_symbol("dlclose", crate::builtins::dlclose as usize);
        self.export_magic_symbol("dlsym", crate::builtins::dlsym as usize);
        self.export_magic_symbol("dlerror", crate::builtins::dlerror as usize);
        self.export_magic_symbol("dladdr", crate::builtins::dladdr as usize);
        self.export_magic_symbol("dladdr1", crate::builtins::dladdr1 as usize);
    }
}

impl ObjectSet {
    pub fn new(mut root: Object) -> Self {
        root.set_id(0);
        Self {
            root,
            libraries: HashMap::new(),
            library_path_map: HashMap::new(),
            state: State::new(),
            last_object_id: 0,
        }
    }

    pub fn objects(&self) -> impl Iterator<Item = &Object> {
        iter::chain(iter::once(&self.root), self.libraries.values())
    }

    fn open_dependencies_of(&mut self, config: &Config, object_id: u32) -> Result<(), Error> {
        let object = match object_id {
            0 => &self.root,
            _ => self.libraries.get(&object_id).ok_or(Error::NotLoaded)?,
        };
        let needed = object.needed().cloned().collect::<Vec<_>>();
        for needed in needed {
            let path = search::find_library(config, needed.as_str())?;
            if self.library_path_map.contains_key(&path) {
                continue;
            }
            let mut object = Object::open(&path, false)?;
            self.last_object_id += 1;
            let id = self.last_object_id;
            object.id = id;
            self.libraries.insert(id, object);
            self.library_path_map.insert(path, id);

            self.open_dependencies_of(config, id)?;
        }
        Ok(())
    }

    pub fn open_root_dependencies(&mut self, config: &Config) -> Result<(), Error> {
        self.open_dependencies_of(config, 0)
    }

    pub fn load_all(&mut self) -> Result<Option<TlsImage>, Error> {
        // Load all objects somewhere, order doesn't matter
        self.root.load()?;
        for (_, lib) in self.libraries.iter_mut() {
            lib.load()?;
        }

        // Build a TLS master copy somewhere
        match thread_local::build_tls_image(&mut self.root, &mut self.libraries)? {
            Some((tls_image, tls_layout)) => {
                self.state.tls_layout = Some(tls_layout);
                Ok(Some(tls_image))
            }
            None => Ok(None),
        }
    }

    pub fn resolve_symbols(&mut self) -> Result<(), Error> {
        // Export all symbols
        self.root.export_symbols(&mut self.state)?;
        for (_, lib) in self.libraries.iter_mut() {
            lib.export_symbols(&mut self.state)?;
        }

        // Resolve all symbols
        self.root.resolve_symbols(&mut self.state);
        for (_, lib) in self.libraries.iter_mut() {
            lib.resolve_symbols(&mut self.state);
        }

        Ok(())
    }

    pub fn relocate(&mut self) -> Result<(), Error> {
        self.root.relocate(&mut self.state)?;
        for (_, lib) in self.libraries.iter_mut() {
            lib.relocate(&mut self.state)?;
        }
        Ok(())
    }

    pub fn initialize(&mut self) {
        // Call global constructors before handing off control to the program
        // .preinit_array, .preinit, ...
        for (_, lib) in self.libraries.iter_mut() {
            unsafe { lib.call_early_constructors() };
        }
        unsafe { self.root.call_early_constructors() };
        // .init_array, .init, ...
        for (_, lib) in self.libraries.iter_mut() {
            unsafe { lib.call_constructors() };
        }
        unsafe { self.root.call_constructors() };
    }

    pub fn is_library_loaded(&self, p: impl AsRef<Path>) -> bool {
        self.library_path_map.contains_key(p.as_ref())
    }

    pub fn get(&self, id: u32) -> Option<&Object> {
        match id {
            0 => Some(&self.root),
            _ => self.libraries.get(&id),
        }
    }

    pub fn get_mut(&mut self, id: u32) -> Option<&mut Object> {
        match id {
            0 => Some(&mut self.root),
            _ => self.libraries.get_mut(&id),
        }
    }

    pub fn lookup_address(&self, address: usize) -> Option<SymbolByAddress> {
        // Lookup in stuff exported from .dynsym
        for (_, symbol) in self.state.symbol_table.iter() {
            if address >= symbol.value && address - symbol.value < symbol.size {
                let object = self.get(symbol.object_id).unwrap();
                let object_base = object.load_range.as_ref().map(|r| r.start).unwrap_or(0);
                return Some(SymbolByAddress {
                    object_name: object.c_name.as_c_str(),
                    object_base,
                    symbol_name: symbol.c_name.as_c_str(),
                    symbol_base: symbol.value,
                });
            }
        }
        // Fall back to main object's .strtab
        for object in self.objects() {
            if let Some(range) = object.load_range.clone() {
                if range.contains(&address) {
                    let address = address - range.start;
                    let (name, base) = object.lookup_address(address)?;
                    return Some(SymbolByAddress {
                        object_name: object.c_name.as_c_str(),
                        object_base: range.start,
                        symbol_name: name,
                        symbol_base: base,
                    });
                }
            }
        }

        None
    }
}
