use std::{
    cmp,
    ffi::{CStr, CString},
    fs::File,
    io::{BufReader, Read, Seek, SeekFrom},
    mem,
    ops::Range,
    path::{Path, PathBuf},
    ptr,
    rc::Rc,
};

use elf::{
    abi::{
        DF_1_PIE, DT_FINI, DT_FINI_ARRAY, DT_FINI_ARRAYSZ, DT_FLAGS_1, DT_INIT, DT_INIT_ARRAY,
        DT_INIT_ARRAYSZ, DT_NEEDED, DT_PREINIT_ARRAY, DT_PREINIT_ARRAYSZ, ET_DYN, ET_EXEC,
        PT_DYNAMIC, PT_GNU_EH_FRAME, PT_GNU_RELRO, PT_GNU_STACK, PT_INTERP, PT_LOAD, PT_NOTE,
        PT_NULL, PT_PHDR, PT_RISCV_ATTRIBUTES, PT_TLS, SHN_UNDEF, SHT_REL, SHT_RELA, STB_GLOBAL,
        STB_LOCAL, STB_WEAK,
    },
    endian::AnyEndian,
    parse::ParsingIterator,
    segment::ProgramHeader,
    symbol::Symbol,
    ElfStream,
};
use yggdrasil_rt::mem::MappingFlags;

use crate::{
    error::Error,
    mapping::Mapping,
    relocation::{RelocationValue, SymbolRelocation},
    state::{ExportedNormalSymbol, ExportedSymbol, ExportedTlsSymbol, State},
};

pub enum ElfType {
    // ET_EXEC
    Static,
    // ET_DYN
    Relocatable(bool),
}

pub enum ResolvedSymbol<'s> {
    Global(&'s ExportedNormalSymbol),
    Tls(&'s ExportedTlsSymbol),
    Local,
    Null(u32),
}

pub struct DynamicSymbol {
    pub raw: Symbol,
    pub name: Rc<str>,
}

pub struct Object {
    pub path: PathBuf,
    pub c_name: CString,
    pub load_range: Option<Range<usize>>,
    pub file: BufReader<File>,
    pub elf: ElfStream<AnyEndian, BufReader<File>>,

    pub vma_start: usize,
    pub vma_end: usize,

    pub ty: ElfType,
    needed: Vec<String>,

    pub tls: Option<ProgramHeader>,
    pub tls_offset: Option<usize>,
    pub id: u32,

    mapping: Option<Mapping>,
    dynamic_symbol_array: Vec<DynamicSymbol>,

    init_array: Option<Range<usize>>,
    pre_init_array: Option<Range<usize>>,

    symtab: Vec<(Symbol, Option<CString>)>,
}

impl ResolvedSymbol<'_> {
    pub fn value(&self) -> i64 {
        match *self {
            Self::Local => todo!(),
            Self::Global(sym) => sym.value as i64,
            Self::Tls(_) => 0,
            Self::Null(_) => 0,
        }
    }
}

impl Object {
    pub fn open<P: AsRef<Path>>(path: P, need_symtab: bool) -> Result<Self, Error> {
        let path = path.as_ref().to_owned();
        let c_name = CString::new(path.to_str().unwrap()).unwrap();
        let file = BufReader::new(File::open(&path)?);
        let mut elf = ElfStream::open_stream(file)?;
        let file = BufReader::new(File::open(&path)?);

        // TODO check ELF validity
        let ty = match elf.ehdr.e_type {
            ET_EXEC => ElfType::Static,
            ET_DYN => {
                // Look for PIE flag in .dynamic
                let is_executable = if let Some(dynamic) = elf.dynamic()? {
                    dynamic
                        .iter()
                        .any(|e| e.d_tag == DT_FLAGS_1 && e.d_val() == DF_1_PIE as _)
                } else {
                    // Weird, ET_DYN without .dynamic?
                    false
                };

                ElfType::Relocatable(is_executable)
            }
            t => return Err(Error::UnhandledElfType(t)),
        };

        let (vma_start, vma_end) = Self::bounds(&elf);

        // Extract DT_NEEDED entries
        let needed = if let Some(dynamic) = elf.dynamic()? {
            dynamic
                .iter()
                .filter_map(|entry| (entry.d_tag == DT_NEEDED).then(|| entry.d_val() as usize))
                .collect()
        } else {
            vec![]
        };

        let needed = if !needed.is_empty() {
            // TODO error?
            let (_, dynstr) = elf.dynamic_symbol_table()?.unwrap();

            needed
                .into_iter()
                .map(|off| dynstr.get(off).map(ToOwned::to_owned))
                .collect::<Result<Vec<_>, _>>()?
        } else {
            vec![]
        };

        // Extract dynamic symbols
        let mut dynamic_symbol_array = vec![];
        if let Some((dynsyms, dynstr)) = elf.dynamic_symbol_table()? {
            for symbol in dynsyms {
                let name = dynstr.get(symbol.st_name as usize)?.into();

                dynamic_symbol_array.push(DynamicSymbol {
                    raw: symbol.clone(),
                    name,
                });
            }
        }

        let tls_segments = elf
            .segments()
            .iter()
            .filter(|p| p.p_type == PT_TLS)
            .collect::<Vec<_>>();
        if tls_segments.len() > 1 {
            todo!("Support objects with more than one TLS segment");
        }
        let tls = tls_segments.first().map(|s| **s);

        let mut symtab = vec![];
        if need_symtab {
            if let Some((symtab_items, strtab)) = elf.symbol_table()? {
                for symbol in symtab_items {
                    let name = strtab.get(symbol.st_name as _)?;
                    let c_name = CString::new(name).ok();
                    symtab.push((symbol, c_name));
                }
            }
            symtab.sort_by(|(s0, _), (s1, _)| cmp::Ord::cmp(&s0.st_value, &s1.st_value));
        }

        Ok(Self {
            path,
            c_name,
            load_range: None,
            file,
            elf,

            ty,
            vma_start,
            vma_end,
            needed,

            tls,
            tls_offset: None,
            id: 0,

            mapping: None,
            dynamic_symbol_array,

            init_array: None,
            pre_init_array: None,
            symtab,
        })
    }

    pub fn set_id(&mut self, id: u32) {
        self.id = id;
    }

    pub fn load(&mut self) -> Result<(), Error> {
        // Already loaded
        if self.mapping.is_some() {
            return Ok(());
        }

        log::info!("Load {:?}", self.path);
        log::info!("Image range: {:#x?}", self.vma_start..self.vma_end);

        // TODO segment granularity mappings
        let mapping_size = self.vma_end - self.vma_start;
        let mut mapping = match self.ty {
            ElfType::Relocatable(_) => {
                Mapping::new(mapping_size, MappingFlags::WRITE | MappingFlags::EXEC)?
            }
            // TODO fixed mapping for this one
            ElfType::Static => todo!(),
        };

        let load_base = (mapping.base() as isize - self.vma_start as isize) as usize;
        self.load_range = Some(load_base..load_base + mapping_size);

        log::info!(
            "Actual range: {:#x?} ({:+#x})",
            mapping.range(),
            mapping.base() as i64 - self.vma_start as i64
        );

        // Parse .dynamic section again to extract extra info
        if let Some(dynamic) = self.elf.dynamic()? {
            let mut dt_init_array = None;
            let mut dt_init_array_sz = None;
            let mut dt_preinit_array = None;
            let mut dt_preinit_array_sz = None;

            for dynamic in dynamic {
                match dynamic.d_tag {
                    DT_INIT_ARRAY => dt_init_array = Some(dynamic.d_ptr()),
                    DT_INIT_ARRAYSZ => dt_init_array_sz = Some(dynamic.d_val()),
                    DT_FINI_ARRAY => todo!(),
                    DT_FINI_ARRAYSZ => todo!(),
                    DT_INIT => todo!(),
                    DT_FINI => todo!(),
                    DT_PREINIT_ARRAY => dt_preinit_array = Some(dynamic.d_ptr()),
                    DT_PREINIT_ARRAYSZ => dt_preinit_array_sz = Some(dynamic.d_val()),
                    _ => (),
                }
            }

            if let (Some(ptr), Some(size)) = (dt_preinit_array, dt_preinit_array_sz) {
                log::debug!("{:?}: DT_PREINIT_ARRAY: {:#x?}", self.path, ptr..ptr + size);
                let ptr = ptr as i64 + mapping.base() as i64 - self.vma_start as i64;
                let ptr = ptr as usize;
                self.pre_init_array = Some(ptr..ptr + size as usize);
            }
            if let (Some(ptr), Some(size)) = (dt_init_array, dt_init_array_sz) {
                // This address is subject to relocation
                log::debug!("{:?}: DT_INIT_ARRAY: {:#x?}", self.path, ptr..ptr + size);
                let ptr = ptr as i64 + mapping.base() as i64 - self.vma_start as i64;
                let ptr = ptr as usize;
                self.init_array = Some(ptr..ptr + size as usize);
            }
        }

        // Load segments
        for segment in self.elf.segments() {
            let mem_size = segment.p_memsz as usize;
            let file_size = segment.p_filesz as usize;

            match segment.p_type {
                PT_NULL => continue,
                PT_LOAD => {
                    // Map offset from image VMA to "real" VMA
                    let load_offset = segment.p_vaddr as usize - self.vma_start;
                    let segment_data = &mut mapping[load_offset..load_offset + mem_size];

                    // Load segment data
                    if file_size != 0 {
                        self.file.seek(SeekFrom::Start(segment.p_offset))?;
                        self.file.read_exact(&mut segment_data[..file_size])?;
                    }

                    segment_data[file_size..].fill(0);
                }
                // Handled separately
                PT_TLS => (),
                // TODO handle .riscv.attributes and extract float ABI
                PT_RISCV_ATTRIBUTES => (),
                // TODO handle GNU_STACK
                PT_DYNAMIC | PT_GNU_RELRO | PT_GNU_STACK | PT_INTERP | PT_PHDR
                | PT_GNU_EH_FRAME | PT_NOTE => (),
                _ => {
                    log::error!("Unhandled segment type: {:#x}", segment.p_type);
                    return Err(Error::UnhandledObject);
                }
            }
        }

        self.mapping = Some(mapping);
        Ok(())
    }

    pub fn lookup_address(&self, address: usize) -> Option<(&CStr, usize)> {
        const NULL_STRING: &CStr = c"";
        let address = address as u64;

        let base = self.load_range.as_ref().map(|r| r.start).unwrap_or(0);

        for (symbol, name) in self.symtab.iter() {
            if address >= symbol.st_value && address - symbol.st_value < symbol.st_size {
                let name = name.as_ref().map(|c| c.as_c_str()).unwrap_or(NULL_STRING);
                return Some((name, symbol.st_value as usize + base));
            }
        }

        None
    }

    pub fn place_tls_copy(&mut self, dst: &mut [u8]) -> Result<(), Error> {
        let tls = self.tls.as_ref().unwrap();
        if tls.p_filesz > 0 {
            self.file.seek(SeekFrom::Start(tls.p_offset))?;
            self.file.read_exact(&mut dst[..tls.p_filesz as usize])?;
        }
        // Zero the rest
        dst[tls.p_filesz as usize..].fill(0);
        Ok(())
    }

    pub fn export_symbols(&self, state: &mut State) -> Result<(), Error> {
        let mapping = self.mapping.as_ref().ok_or(Error::NotLoaded)?;

        // Export dynamic symbols into the linking table
        for dynsym in self.dynamic_symbol_array.iter() {
            // Don't export undefined symbols
            if dynsym.raw.st_shndx == SHN_UNDEF {
                continue;
            }

            // TODO: don't export hidden symbols?
            // Symbol actual address offset: V + B - S
            // Where V: symbol value as defined in the table
            //       B: "actual" load base
            //       S: "image" load base
            let offset = mapping.base() as isize - self.vma_start as isize;
            match dynsym.raw.st_bind() {
                STB_GLOBAL => state.export(&self.path, dynsym, offset, false, self.id),
                STB_WEAK => state.export(&self.path, dynsym, offset, true, self.id),
                STB_LOCAL => todo!(),
                _ => todo!(),
            }
        }

        Ok(())
    }

    unsafe fn call_constructor_list(range: Range<usize>) {
        type Constructor = extern "C" fn();

        for slot in range.step_by(size_of::<usize>()) {
            let func = ptr::with_exposed_provenance::<usize>(slot).read_unaligned();

            if func == 0 || func == usize::MAX {
                continue;
            }

            let func = mem::transmute::<_, Constructor>(func);
            log::debug!("ld: call constructor {func:p}");

            func();
        }
    }

    /// # Safety
    ///
    /// Only allowed to be called once. The object must be loaded and relocated before calling
    /// this function.
    pub unsafe fn call_early_constructors(&mut self) {
        if let Some(pre_init_array) = self.pre_init_array.as_ref() {
            Self::call_constructor_list(pre_init_array.clone());
        }
    }

    /// # Safety
    ///
    /// Only allowed to be called once. The object must be loaded and relocated before calling
    /// this function.
    pub unsafe fn call_constructors(&mut self) {
        if let Some(init_array) = self.init_array.as_ref() {
            Self::call_constructor_list(init_array.clone());
        }
    }

    pub fn entry(&self) -> Option<extern "C" fn(usize)> {
        let entry = self.elf.ehdr.e_entry as i64;
        let offset = match &self.ty {
            // Fixed entry
            ElfType::Static => 0,
            // PIE executable
            ElfType::Relocatable(true) => {
                let mapping = self.mapping.as_ref()?;
                mapping.base() as i64 - self.vma_start as i64
            }
            // Shared library
            ElfType::Relocatable(false) => return None,
        };

        let entry: usize = (entry + offset).try_into().unwrap();
        Some(unsafe { core::mem::transmute(entry) })
    }

    pub fn needed(&self) -> impl Iterator<Item = &String> {
        self.needed.iter()
    }

    pub fn resolve_symbols(&mut self, state: &mut State) {
        for dynsym in self.dynamic_symbol_array.iter() {
            // Don't resolve local symbols
            if dynsym.raw.st_bind() == STB_LOCAL {
                continue;
            }

            if dynsym.raw.st_shndx == SHN_UNDEF {
                state.ensure_resolved(&self.path, dynsym);
            }
        }
    }

    fn apply_relocations<'s, R: SymbolRelocation, F: Fn(usize) -> &'s DynamicSymbol>(
        state: &mut State,
        mapping: &mut Mapping,
        object_id: u32,
        rel_section: ParsingIterator<AnyEndian, R>,
        get_dynsym: F,
    ) -> Result<(), Error> {
        for rel in rel_section {
            let dynsym = get_dynsym(rel.symbol());
            let sym = match dynsym.raw.st_bind() {
                STB_GLOBAL | STB_WEAK => match state.lookup(dynsym).unwrap() {
                    ExportedSymbol::Normal(export) => ResolvedSymbol::Global(export),
                    ExportedSymbol::Tls(export) => ResolvedSymbol::Tls(export),
                },
                STB_LOCAL => {
                    if dynsym.name.is_empty() {
                        ResolvedSymbol::Null(object_id)
                    } else {
                        todo!("Relocation against local symbol: {:?}", dynsym.name)
                    }
                }
                _ => todo!("Unhandled relocation symbol binding"),
            };

            if let Some(value) = rel.resolve(state, &dynsym.name, &sym, mapping.base())? {
                value.write(mapping, rel.offset());
            }
        }
        Ok(())
    }

    pub fn relocate(&mut self, state: &mut State) -> Result<(), Error> {
        let mapping = self.mapping.as_mut().ok_or(Error::NotLoaded)?;

        let rela_sections = self
            .elf
            .section_headers()
            .iter()
            .filter(|s| s.sh_type == SHT_RELA)
            .cloned()
            .collect::<Vec<_>>();

        for rela_section in rela_sections {
            let rela_section = self.elf.section_data_as_relas(&rela_section)?;
            Self::apply_relocations(state, mapping, self.id, rela_section, |idx| {
                &self.dynamic_symbol_array[idx]
            })?;
        }

        let rel_sections = self
            .elf
            .section_headers()
            .iter()
            .filter(|s| s.sh_type == SHT_REL)
            .cloned()
            .collect::<Vec<_>>();

        for rel_section in rel_sections {
            let rel_section = self.elf.section_data_as_rels(&rel_section)?;
            Self::apply_relocations(state, mapping, self.id, rel_section, |idx| {
                &self.dynamic_symbol_array[idx]
            })?;
        }

        Ok(())
    }

    fn bounds<F: Read + Seek>(elf: &ElfStream<AnyEndian, F>) -> (usize, usize) {
        let mut start = u64::MAX;
        let mut end = u64::MIN;

        for segment in elf.segments() {
            if segment.p_type != PT_LOAD {
                continue;
            }

            let start_vma = segment.p_vaddr & !(segment.p_align - 1);
            let end_vma =
                (segment.p_vaddr + segment.p_memsz + segment.p_align - 1) & !(segment.p_align - 1);

            if start_vma < start {
                start = start_vma;
            }
            if end_vma > end {
                end = end_vma;
            }
        }

        (start.try_into().unwrap(), end.try_into().unwrap())
    }
}
