use crate::error::Result;
use crate::input_data::FileId;
use crate::input_data::InputRef;
use crate::input_data::LoadedStubLibrary;
use crate::input_data::MAX_FILES_PER_GROUP;
use crate::input_section_id::InputSectionId;
use crate::input_section_id::SectionIdRange;
use crate::macho_stub_library::DefinedStubLibrary;
use crate::parsing::ParsedInputObject;
use crate::parsing::Prelude;
use crate::parsing::ProcessedLinkerScript;
use crate::parsing::SyntheticSymbols;
use crate::platform;
use crate::platform::ObjectFile;
use crate::platform::Platform;
use crate::sharding::ShardKey as _;
use crate::symbol::UnversionedSymbolName;
use crate::symbol_db::SymbolDb;
use crate::symbol_db::SymbolId;
use crate::symbol_db::SymbolIdRange;
use crate::symbol_db::SymbolStrength;
use crate::timing_phase;
use crate::verbose_timing_phase;
use std::fmt::Display;

#[derive(Debug)]
pub(crate) enum Group<'data, P: Platform> {
    Prelude(Prelude<'data, P>),
    Objects(&'data [SequencedInputObject<'data, P>]),
    StubLibraries(Vec<SequencedStubLibrary<'data>>),
    LinkerScripts(Vec<SequencedLinkerScript<'data, P>>),
    SyntheticSymbols(SyntheticSymbols),
    #[cfg(all(feature = "plugins", unix))]
    LtoInputs(Vec<crate::linker_plugins::LtoInput<'data>>),
}

#[derive(Debug)]
pub(crate) struct SequencedInputObject<'data, P: Platform> {
    pub(crate) parsed: Box<ParsedInputObject<'data, P>>,
    pub(crate) symbol_id_range: SymbolIdRange,
    pub(crate) section_id_range: SectionIdRange,
    pub(crate) file_id: FileId,
}

#[derive(Debug)]
pub(crate) struct SequencedLinkerScript<'data, P: Platform> {
    pub(crate) parsed: ProcessedLinkerScript<'data, P>,
    pub(crate) symbol_id_range: SymbolIdRange,
    pub(crate) file_id: FileId,
}

#[derive(Debug)]
pub(crate) struct SequencedStubLibrary<'data> {
    pub(crate) input: InputRef<'data>,
    pub(crate) defined_symbols: DefinedStubLibrary<'data>,
    pub(crate) symbol_id_range: SymbolIdRange,
    pub(crate) file_id: FileId,
}

#[derive(Debug)]
pub(crate) enum SequencedInput<'db, 'data, P: Platform> {
    Prelude(&'db Prelude<'data, P>),
    Object(&'data SequencedInputObject<'data, P>),
    StubLibrary(&'db SequencedStubLibrary<'data>),
    LinkerScript(&'db SequencedLinkerScript<'data, P>),
    SyntheticSymbols(&'db SyntheticSymbols),
    #[cfg(all(feature = "plugins", unix))]
    LtoInput(&'db crate::linker_plugins::LtoInput<'data>),
}

impl<'data, P: Platform> Group<'data, P> {
    // This is used when the verbose-ttttiming feature is enabled.
    pub(crate) fn group_id(&self) -> usize {
        match self {
            Group::Prelude(_) => 0,
            Group::Objects(objects) => objects[0].file_id.group(),
            Group::StubLibraries(stubs) => stubs[0].file_id.group(),
            Group::LinkerScripts(scripts) => scripts[0].file_id.group(),
            Group::SyntheticSymbols(s) => s.file_id.group(),
            #[cfg(all(feature = "plugins", unix))]
            Group::LtoInputs(s) => s[0].file_id.group(),
        }
    }

    pub(crate) fn start_symbol_id(&self) -> SymbolId {
        self.symbol_id_range().start()
    }

    pub(crate) fn symbol_id_range(&self) -> SymbolIdRange {
        match self {
            Group::Prelude(o) => SymbolIdRange::prelude(o.symbol_definitions.len()),
            Group::Objects(objects) => SymbolIdRange::covering(
                objects[0].symbol_id_range,
                objects[objects.len() - 1].symbol_id_range,
            ),
            Group::StubLibraries(stubs) => SymbolIdRange::covering(
                stubs[0].symbol_id_range,
                stubs[stubs.len() - 1].symbol_id_range,
            ),
            Group::LinkerScripts(scripts) => {
                if scripts.is_empty() {
                    SymbolIdRange::empty()
                } else {
                    SymbolIdRange::covering(
                        scripts[0].symbol_id_range,
                        scripts[scripts.len() - 1].symbol_id_range,
                    )
                }
            }
            Group::SyntheticSymbols(o) => o.symbol_id_range,
            #[cfg(all(feature = "plugins", unix))]
            Group::LtoInputs(objects) => SymbolIdRange::covering(
                objects[0].symbol_id_range,
                objects[objects.len() - 1].symbol_id_range,
            ),
        }
    }

    pub(crate) fn num_symbols(&self) -> usize {
        self.symbol_id_range().len()
    }
}

pub(crate) fn create_groups<'data, P: Platform>(
    symbol_db: &mut SymbolDb<'data, P>,
    parsed_objects: Vec<Box<ParsedInputObject<'data, P>>>,
    stub_libraries: Vec<LoadedStubLibrary<'data>>,
    linker_scripts: Vec<ProcessedLinkerScript<'data, P>>,
) {
    timing_phase!("Group files");

    let max_files_per_group = determine_max_files_per_group(symbol_db.args);
    let num_symbols = count_symbols(&parsed_objects);
    let symbols_per_group = determine_symbols_per_group(num_symbols, symbol_db.args);

    symbol_db.groups_reserve(parsed_objects.len() / max_files_per_group + 3);

    let mut next_symbol_id = symbol_db.next_symbol_id();
    let mut next_input_section_id = symbol_db.next_input_section_id;

    let mut objects = parsed_objects.into_iter().peekable();

    let mut num_symbols_in_group = 0;
    let mut group_objects = Vec::new();

    let allocator = symbol_db.herd.get();

    while let Some(parsed) = objects.next() {
        let file_id = FileId::new(symbol_db.next_group_index(), group_objects.len() as u32);
        let num_symbols_in_file = parsed.object.num_symbols();

        let section_count = if parsed.object.is_dynamic() {
            // We don't copy sections from dynamic objects into the output, so for our purposes,
            // there are no sections.
            0
        } else {
            parsed.object.num_sections()
        };

        group_objects.push(SequencedInputObject {
            parsed,
            symbol_id_range: SymbolIdRange::input(next_symbol_id, num_symbols_in_file),
            section_id_range: SectionIdRange::input(next_input_section_id, section_count),
            file_id,
        });

        next_symbol_id = next_symbol_id.add_usize(num_symbols_in_file);
        next_input_section_id = next_input_section_id.add_usize(section_count);

        num_symbols_in_group += num_symbols_in_file;

        // Finish the current group if we've reached the maximum number of files for the group, if
        // this is the last file or if the next file would put us over the per-group symbol limit.
        let finish_group = group_objects.len() >= max_files_per_group
            || objects.peek().is_none_or(|next_obj| {
                num_symbols_in_group + next_obj.object.num_symbols() > symbols_per_group
            });

        if finish_group {
            num_symbols_in_group = 0;

            debug_assert!(
                group_objects.len() <= MAX_FILES_PER_GROUP as usize,
                "Group is too large: {}",
                group_objects.len()
            );

            let objects_slice =
                allocator.alloc_slice_fill_iter(core::mem::take(&mut group_objects));

            symbol_db.add_group(Group::Objects(objects_slice));
        }
    }

    let linker_scripts: Vec<SequencedLinkerScript<'_, P>> = linker_scripts
        .into_iter()
        .enumerate()
        .map(|(i, script)| {
            let symbol_id_range = SymbolIdRange::input(next_symbol_id, script.num_symbols());
            next_symbol_id = next_symbol_id.add_usize(symbol_id_range.len());

            SequencedLinkerScript {
                parsed: script,
                file_id: FileId::new(symbol_db.next_group_index(), i as u32),
                symbol_id_range,
            }
        })
        .collect();

    if !linker_scripts.is_empty() {
        symbol_db.add_group(Group::LinkerScripts(linker_scripts));
    }

    let stub_libraries: Vec<SequencedStubLibrary<'data>> = stub_libraries
        .into_iter()
        .enumerate()
        .map(|(i, stub)| {
            let symbol_id_range =
                SymbolIdRange::input(next_symbol_id, stub.defined_symbols.total_symbols());
            next_symbol_id = next_symbol_id.add_usize(symbol_id_range.len());

            SequencedStubLibrary {
                input: stub.input,
                defined_symbols: stub.defined_symbols,
                symbol_id_range,
                file_id: FileId::new(symbol_db.next_group_index(), i as u32),
            }
        })
        .collect();

    if !stub_libraries.is_empty() {
        symbol_db.add_group(Group::StubLibraries(stub_libraries));
    }

    symbol_db.next_input_section_id = next_input_section_id;

    tracing::trace!(
        "GROUPS:\n{}",
        std::fmt::from_fn(|f| {
            for (i, group) in symbol_db.groups.iter().enumerate() {
                writeln!(f, "{i}: {group}")?;
            }
            Ok(())
        })
    );
}

/// Decides after how many symbols, we should start a new group.
fn determine_symbols_per_group(num_symbols: usize, args: &impl platform::Args) -> usize {
    let num_threads = args.common().available_threads.get();

    // If we're running with a single thread, then we might as well put everything into a single
    // group.
    if num_threads == 1 {
        return usize::MAX;
    }

    // If we have lots of threads, then we might benefit from a few more groups in order to properly
    // take advantage of the available parallelism.
    let groups_per_thread =
        args.common()
            .numeric_experiment(crate::args::Experiment::GroupsPerThread, 5) as usize;

    // If we don't have lots of threads, then we still want a reasonable number of groups. The need
    // for this was based on experimentation.
    let min_groups = args
        .common()
        .numeric_experiment(crate::args::Experiment::MinGroups, 150) as usize;

    let target_num_groups = (num_threads * groups_per_thread).max(min_groups);

    1.max(num_symbols / target_num_groups)
}

/// Decides the maximum number of files that we'll put into one group.
fn determine_max_files_per_group(args: &impl platform::Args) -> usize {
    if let Some(v) = args.common().files_per_group {
        return v as usize;
    }

    // We may eventually find that a lower value based on the number of threads is better, but for
    // now, if files are small, we allow lots of them in a single group.
    crate::input_data::MAX_FILES_PER_GROUP as usize
}

/// Compute the total number of symbols in the supplied objects.
fn count_symbols<P: Platform>(objects: &[Box<ParsedInputObject<P>>]) -> usize {
    verbose_timing_phase!("Count symbols");

    objects.iter().map(|o| o.num_symbols()).sum::<usize>()
}

impl<'data, P: Platform> SequencedInputObject<'data, P> {
    pub(crate) fn symbol_name(
        &self,
        symbol_id: crate::symbol_db::SymbolId,
    ) -> Result<UnversionedSymbolName<'data>> {
        let index = symbol_id.to_input(self.symbol_id_range);
        let symbol = self.parsed.object.symbol(index)?;
        Ok(UnversionedSymbolName::new(
            self.parsed.object.symbol_name(symbol)?,
        ))
    }

    /// Get the version of a symbol. Only intended for diagnostic purposes since it's potentially
    /// quite slow.
    pub(crate) fn symbol_version_debug(
        &self,
        symbol_id: crate::symbol_db::SymbolId,
    ) -> Option<String> {
        self.parsed
            .object
            .symbol_version_debug(symbol_id.to_input(self.symbol_id_range))
    }

    /// Returns whether this input should be skipped if there are no non-weak references to symbols
    /// it defines. This is true for archive entries for which --whole-archive is false and shared
    /// objects for which --as-needed is true.
    pub(crate) fn is_optional(&self) -> bool {
        (self.parsed.input.has_archive_semantics() && !self.parsed.modifiers.whole_archive)
            || (self.is_dynamic() && self.parsed.modifiers.as_needed)
    }

    pub(crate) fn is_dynamic(&self) -> bool {
        self.parsed.object.is_dynamic()
    }

    pub(crate) fn symbol_strength(&self, symbol_id: crate::symbol_db::SymbolId) -> SymbolStrength {
        let local_index = symbol_id.to_input(self.symbol_id_range);
        let Ok(obj_symbol) = self.parsed.object.symbol(local_index) else {
            return SymbolStrength::Undefined;
        };
        SymbolStrength::of(obj_symbol)
    }

    fn input_section_id_for_symbol(&self, symbol_id: SymbolId) -> Option<InputSectionId> {
        let symbol_index = self.symbol_id_range.id_to_input(symbol_id);
        let symtab_entry = self.parsed.object.symbol(symbol_index).ok()?;

        let section_index = self
            .parsed
            .object
            .symbol_section(symtab_entry, symbol_index)
            .ok()??;

        Some(self.section_id_range.input_to_id(section_index))
    }
}

impl<'data, P: Platform> SequencedLinkerScript<'data, P> {
    pub(crate) fn symbol_name(&self, symbol_id: SymbolId) -> UnversionedSymbolName<'data> {
        let local_index = self.symbol_id_range.id_to_offset(symbol_id);
        UnversionedSymbolName::new(self.parsed.symbol_defs[local_index].name)
    }
}

impl<'data> SequencedStubLibrary<'data> {
    pub(crate) fn symbol_name(&self, symbol_id: SymbolId) -> UnversionedSymbolName<'data> {
        let local_index = self.symbol_id_range.id_to_offset(symbol_id);
        UnversionedSymbolName::new(
            self.defined_symbols
                .symbols
                .get(local_index)
                .unwrap_or_else(|| {
                    &self.defined_symbols.weak_symbols
                        [local_index - self.defined_symbols.symbols.len()]
                })
                .as_bytes(),
        )
    }

    pub(crate) fn symbol_strength(&self, symbol_id: SymbolId) -> SymbolStrength {
        let local_index = self.symbol_id_range.id_to_offset(symbol_id);
        if local_index < self.defined_symbols.symbols.len() {
            SymbolStrength::Strong
        } else {
            SymbolStrength::Weak
        }
    }
}

impl<'db, 'data, P: Platform> SequencedInput<'db, 'data, P> {
    pub(crate) fn symbol_id_range(&self) -> SymbolIdRange {
        match self {
            SequencedInput::Prelude(o) => SymbolIdRange::prelude(o.symbol_definitions.len()),
            SequencedInput::Object(o) => o.symbol_id_range,
            SequencedInput::StubLibrary(o) => o.symbol_id_range,
            SequencedInput::LinkerScript(o) => o.symbol_id_range,
            SequencedInput::SyntheticSymbols(o) => o.symbol_id_range,
            #[cfg(all(feature = "plugins", unix))]
            SequencedInput::LtoInput(o) => o.symbol_id_range,
        }
    }

    pub(crate) fn is_dynamic(&self) -> bool {
        match self {
            SequencedInput::Object(o) if o.is_dynamic() => true,
            SequencedInput::StubLibrary(_) => true,
            _ => false,
        }
    }

    pub(crate) fn symbol_strength(&self, symbol_id: SymbolId) -> SymbolStrength {
        match self {
            SequencedInput::Object(o) => o.symbol_strength(symbol_id),
            SequencedInput::StubLibrary(o) => o.symbol_strength(symbol_id),
            _ => SymbolStrength::Undefined,
        }
    }

    /// Returns the ID of the input section in which the specified symbol is defined.
    pub(crate) fn input_section_id_for_symbol(
        &self,
        symbol_id: SymbolId,
    ) -> Option<InputSectionId> {
        match self {
            SequencedInput::Object(sequenced_input_object) => {
                sequenced_input_object.input_section_id_for_symbol(symbol_id)
            }
            _ => None,
        }
    }
}

impl<'data, P: Platform> std::fmt::Display for SequencedInputObject<'data, P> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        std::fmt::Display::fmt(&self.parsed.input, f)
    }
}

impl std::fmt::Display for SequencedStubLibrary<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        std::fmt::Display::fmt(&self.input, f)
    }
}

impl<'data, P: Platform> Display for Group<'data, P> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Group::Prelude(_) => write!(f, "<prelude>"),
            Group::Objects(parsed_input_objects) => {
                let num_symbols: usize = parsed_input_objects
                    .iter()
                    .map(|o| o.symbol_id_range.len())
                    .sum();

                write!(
                    f,
                    "num_objects: {num_objects} num_symbols: {num_symbols}",
                    num_objects = parsed_input_objects.len(),
                )
            }
            Group::StubLibraries(stubs) => write!(f, "{} Mach-O stub library(s)", stubs.len()),
            Group::LinkerScripts(scripts) => write!(f, "{} linker script(s)", scripts.len()),
            Group::SyntheticSymbols(_) => write!(f, "<epilogue>"),
            #[cfg(all(feature = "plugins", unix))]
            Group::LtoInputs(lto_inputs) => write!(f, "<{} lto inputs>", lto_inputs.len()),
        }
    }
}

impl<'db, 'data, P: Platform> std::fmt::Display for SequencedInput<'db, 'data, P> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SequencedInput::Prelude(_) => std::fmt::Display::fmt("<prelude>", f),
            SequencedInput::Object(o) => std::fmt::Display::fmt(o, f),
            SequencedInput::StubLibrary(o) => std::fmt::Display::fmt(o, f),
            SequencedInput::LinkerScript(o) => std::fmt::Display::fmt(&o.parsed, f),
            SequencedInput::SyntheticSymbols(_) => std::fmt::Display::fmt("<epilogue>", f),
            #[cfg(all(feature = "plugins", unix))]
            SequencedInput::LtoInput(o) => std::fmt::Display::fmt(o, f),
        }
    }
}
