citum_engine/processor/bibliography/

grouping.rs

/*
SPDX-License-Identifier: MIT OR Apache-2.0
SPDX-FileCopyrightText: © 2023-2026 Bruce D'Arcus
*/

//! Grouped bibliography rendering with configurable selectors and sorting.

use super::RenderedBibliographyGroup;
use crate::api::AnnotationStyle;
use crate::grouping::{GroupSorter, SelectorEvaluator};
use crate::processor::Processor;
use crate::processor::disambiguation::Disambiguator;
use crate::processor::rendering::{CompoundRenderData, Renderer, RendererResources};
use crate::reference::{Bibliography, Reference};
use crate::render::ProcEntry;
use crate::render::format::{OutputFormat, ProcEntryMetadata};
use crate::values::{ProcHints, RenderContext, RenderOptions, format_contributors_short};
use citum_schema::grouping::{BibliographyGroup, DisambiguationScope, GroupHeading};
use citum_schema::options::{BibliographyPartitionHeading, BibliographySortPartitioning};
use std::borrow::Cow;
use std::collections::{HashMap, HashSet};

impl Processor {
    pub(super) fn resolve_group_heading(&self, heading: &GroupHeading) -> Option<String> {
        match heading {
            GroupHeading::Literal { literal } => Some(literal.clone()),
            GroupHeading::Term { term, form } => self.locale.resolved_general_term(
                term,
                &form.clone().unwrap_or(citum_schema::locale::TermForm::Long),
                None,
            ),
            GroupHeading::Localized { localized } => self.resolve_localized_heading(localized),
        }
    }

    fn resolve_localized_heading(&self, localized: &HashMap<String, String>) -> Option<String> {
        fn language_tag(locale: &str) -> &str {
            locale.split('-').next().unwrap_or(locale)
        }

        let mut candidates = Vec::new();
        let mut push_candidate = |locale: &str| {
            let candidate = locale.to_string();
            if !candidates.contains(&candidate) {
                candidates.push(candidate);
            }
        };

        push_candidate(&self.locale.locale);
        push_candidate(language_tag(&self.locale.locale));

        if let Some(default_locale) = self.style.info.default_locale.as_deref() {
            push_candidate(default_locale);
            push_candidate(language_tag(default_locale));
        }

        push_candidate("en-US");
        push_candidate("en");

        for locale in candidates {
            if let Some(value) = localized.get(&locale) {
                return Some(value.clone());
            }
        }

        localized
            .iter()
            .min_by(|left, right| left.0.cmp(right.0))
            .map(|(_locale, value)| value.clone())
    }

    fn resolve_partition_heading(&self, heading: &BibliographyPartitionHeading) -> Option<String> {
        match heading {
            BibliographyPartitionHeading::Literal { literal } => Some(literal.clone()),
            BibliographyPartitionHeading::Term { term, form } => self.locale.resolved_general_term(
                term,
                &form.clone().unwrap_or(citum_schema::locale::TermForm::Long),
                None,
            ),
            BibliographyPartitionHeading::Localized { localized } => {
                self.resolve_localized_heading(localized)
            }
        }
    }

    fn collect_matching_group_refs<'a>(
        &'a self,
        bibliography: &'a [ProcEntry],
        assigned: &HashSet<String>,
        evaluator: &SelectorEvaluator<'_>,
        group: &BibliographyGroup,
    ) -> Vec<&'a Reference> {
        bibliography
            .iter()
            .filter(|entry| !assigned.contains(&entry.id))
            .filter_map(|entry| {
                self.bibliography
                    .get(&entry.id)
                    .filter(|reference| evaluator.matches(reference, &group.selector))
            })
            .collect()
    }

    /// Returns `ProcEntry` stubs with only `id` populated, in sort order.
    ///
    /// Used for grouping paths that only need IDs for selector matching — avoids
    /// the full PlainText render pass that `process_references` performs.
    fn sorted_id_stubs(&self) -> Vec<ProcEntry> {
        self.initialize_numeric_bibliography_numbers();
        self.sort_references(self.bibliography.values().collect())
            .into_iter()
            .filter_map(|r| {
                r.id().map(|id| ProcEntry {
                    id: id.to_string(),
                    template: vec![],
                    metadata: ProcEntryMetadata::default(),
                })
            })
            .collect()
    }

    fn mark_group_members_assigned(assigned: &mut HashSet<String>, references: &[&Reference]) {
        for reference in references {
            if let Some(id) = reference.id() {
                assigned.insert(id.to_string());
            }
        }
    }

    fn build_group_local_hints(
        &self,
        sorted_refs: &[&Reference],
        group: &BibliographyGroup,
    ) -> Option<HashMap<String, ProcHints>> {
        if !matches!(group.disambiguate, Some(DisambiguationScope::Locally)) {
            return None;
        }

        let mut group_bibliography = Bibliography::new();
        for reference in sorted_refs {
            group_bibliography.insert(
                reference.id().unwrap_or_default().to_string(),
                (*reference).clone(),
            );
        }

        let resolved_sort = group
            .sort
            .as_ref()
            .map(citum_schema::GroupSortEntry::resolve);
        let bibliography_config = self.get_bibliography_config();
        let disambiguator = if let Some(sort) = resolved_sort.as_ref() {
            Disambiguator::with_group_sort(
                &group_bibliography,
                &bibliography_config,
                &self.locale,
                sort,
            )
        } else {
            Disambiguator::new(&group_bibliography, &bibliography_config, &self.locale)
        };

        Some(disambiguator.calculate_hints())
    }

    fn effective_group_style<'a>(
        &'a self,
        group: &'a BibliographyGroup,
    ) -> Cow<'a, citum_schema::Style> {
        if let Some(group_template) = &group.template {
            let mut local_style = self.style.clone();
            if let Some(bibliography) = local_style.bibliography.as_mut() {
                bibliography.template = Some(group_template.clone());
            }
            Cow::Owned(local_style)
        } else {
            Cow::Borrowed(&self.style)
        }
    }

    fn render_group_entries<F>(
        &self,
        _bibliography: &[ProcEntry],
        sorted_refs: Vec<&Reference>,
        group: &BibliographyGroup,
        local_hints: Option<&HashMap<String, ProcHints>>,
    ) -> Vec<ProcEntry>
    where
        F: OutputFormat<Output = String>,
    {
        // Always process entries with format F so that group components (pre_formatted=true)
        // contain markup in the target format rather than PlainText (_..._).
        let hints = local_hints.unwrap_or(&self.hints);
        let effective_style = self.effective_group_style(group);
        let bibliography_config = self.get_bibliography_config();
        let bibliography_options = self.get_bibliography_options().into_owned();
        let substitute = bibliography_options.subsequent_author_substitute.clone();
        let renderer = Renderer::new(
            RendererResources {
                style: &effective_style,
                bibliography: &self.bibliography,
                locale: &self.locale,
                config: &bibliography_config,
                bibliography_config: Some(bibliography_options),
            },
            hints,
            &self.citation_numbers,
            CompoundRenderData {
                set_by_ref: &self.compound_set_by_ref,
                member_index: &self.compound_member_index,
                sets: &self.compound_sets,
            },
            self.show_semantics,
            self.inject_ast_indices,
            self.abbreviation_map.as_ref(),
        );

        let mut entries = Vec::new();
        let mut previous_reference: Option<&Reference> = None;

        for (index, reference) in sorted_refs.into_iter().enumerate() {
            let ref_id = reference.id().unwrap_or_default().to_string();
            let entry_number = self
                .citation_numbers
                .borrow()
                .get(&ref_id)
                .copied()
                .unwrap_or(index + 1);

            if let Some(mut processed) =
                renderer.process_bibliography_entry_with_format::<F>(reference, entry_number)
            {
                if let Some(substitute_string) = substitute.as_deref()
                    && let Some(previous) = previous_reference
                    && self.contributors_match(previous, reference)
                {
                    renderer.apply_author_substitution_with_format::<F>(
                        &mut processed,
                        substitute_string,
                    );
                }

                entries.push(ProcEntry {
                    id: ref_id,
                    template: processed,
                    metadata: self.extract_metadata(reference),
                });
                previous_reference = Some(reference);
            }
        }

        entries
    }

    fn append_rendered_group<F>(
        &self,
        result: &mut String,
        group: &BibliographyGroup,
        entries: Vec<ProcEntry>,
        annotations: Option<&HashMap<String, String>>,
        annotation_style: Option<&AnnotationStyle>,
        suppress_heading: bool,
    ) where
        F: OutputFormat<Output = String>,
    {
        if !result.is_empty() {
            result.push_str("\n\n");
        }

        if !suppress_heading
            && let Some(heading) = group
                .heading
                .as_ref()
                .and_then(|group_heading| self.resolve_group_heading(group_heading))
        {
            result.push_str(&self.render_group_heading::<F>(&heading));
        }

        result.push_str(&crate::render::refs_to_string_with_format::<F>(
            entries,
            annotations,
            annotation_style,
        ));
    }

    fn append_rendered_partition<F>(
        &self,
        result: &mut String,
        heading: Option<&BibliographyPartitionHeading>,
        entries: Vec<ProcEntry>,
        annotations: Option<&HashMap<String, String>>,
        annotation_style: Option<&AnnotationStyle>,
    ) where
        F: OutputFormat<Output = String>,
    {
        if !result.is_empty() {
            result.push_str("\n\n");
        }

        if let Some(heading) =
            heading.and_then(|group_heading| self.resolve_partition_heading(group_heading))
        {
            result.push_str(&self.render_group_heading::<F>(&heading));
        }

        result.push_str(&crate::render::refs_to_string_with_format::<F>(
            entries,
            annotations,
            annotation_style,
        ));
    }

    pub(super) fn render_with_partition_sections<F>(
        &self,
        sorted_refs: Vec<&Reference>,
        partitioning: &BibliographySortPartitioning,
        annotations: Option<&HashMap<String, String>>,
        annotation_style: Option<&AnnotationStyle>,
    ) -> String
    where
        F: OutputFormat<Output = String>,
    {
        let fmt = F::default();
        let mut result = String::new();

        for (partition_key, references) in
            crate::sort_partitioning::partition_references(sorted_refs, &self.locale, partitioning)
        {
            let heading = partition_key
                .as_ref()
                .and_then(|key| partitioning.headings.get(key));
            let entries = self.merge_compound_entries::<F>(self.process_sorted_refs::<_, F>(
                references.into_iter(),
                |reference, entry_number| {
                    self.process_bibliography_entry_with_format::<F>(reference, entry_number)
                },
            ));
            self.append_rendered_partition::<F>(
                &mut result,
                heading,
                entries,
                annotations,
                annotation_style,
            );
        }

        fmt.finish(result)
    }

    pub(super) fn render_with_custom_groups<F>(
        &self,
        all_entries: &[ProcEntry],
        groups: &[BibliographyGroup],
    ) -> String
    where
        F: OutputFormat<Output = String>,
    {
        let selected: HashSet<String> = all_entries.iter().map(|e| e.id.clone()).collect();
        self.render_with_custom_groups_filtered::<F>(all_entries, groups, &selected, None, None)
    }

    pub(super) fn render_with_custom_groups_filtered<F>(
        &self,
        all_entries: &[ProcEntry],
        groups: &[BibliographyGroup],
        selected: &HashSet<String>,
        annotations: Option<&HashMap<String, String>>,
        annotation_style: Option<&AnnotationStyle>,
    ) -> String
    where
        F: OutputFormat<Output = String>,
    {
        let fmt = F::default();
        let cited_ids = self.cited_ids.borrow();
        let evaluator = SelectorEvaluator::new(&cited_ids);
        let sorter = GroupSorter::new(&self.locale);

        let mut assigned = HashSet::new();
        let mut result = String::new();

        // First pass: collect all populated groups with their rendered entries
        let mut populated_groups: Vec<(&BibliographyGroup, Vec<ProcEntry>)> = Vec::new();

        for group in groups {
            let matching_refs =
                self.collect_matching_group_refs(all_entries, &assigned, &evaluator, group);

            let matching_refs: Vec<&Reference> = matching_refs
                .into_iter()
                .filter(|r| r.id().as_deref().is_some_and(|id| selected.contains(id)))
                .collect();

            if matching_refs.is_empty() {
                continue;
            }

            Self::mark_group_members_assigned(&mut assigned, &matching_refs);

            let sorted_refs = if let Some(sort_spec) = &group.sort {
                sorter.sort_references(matching_refs, &sort_spec.resolve())
            } else {
                matching_refs
            };
            let local_hints = self.build_group_local_hints(&sorted_refs, group);
            let entries = self.merge_compound_entries::<F>(self.render_group_entries::<F>(
                all_entries,
                sorted_refs,
                group,
                local_hints.as_ref(),
            ));

            populated_groups.push((group, entries));
        }

        // Compute unassigned entries to determine if heading suppression applies
        let unassigned_refs: Vec<&Reference> = all_entries
            .iter()
            .filter(|entry| !assigned.contains(&entry.id) && selected.contains(&entry.id))
            .filter_map(|entry| self.bibliography.get(&entry.id))
            .collect();

        let suppress_heading = populated_groups.len() == 1 && unassigned_refs.is_empty();

        // Second pass: render populated groups with optional heading suppression
        for (group, entries) in populated_groups {
            self.append_rendered_group::<F>(
                &mut result,
                group,
                entries,
                annotations,
                annotation_style,
                suppress_heading,
            );
        }

        self.append_unassigned_entries_filtered::<F>(
            &mut result,
            all_entries,
            &assigned,
            selected,
            annotations,
            annotation_style,
        );
        fmt.finish(result)
    }

    fn append_unassigned_entries_filtered<F>(
        &self,
        result: &mut String,
        bibliography: &[ProcEntry],
        assigned: &HashSet<String>,
        selected: &HashSet<String>,
        annotations: Option<&HashMap<String, String>>,
        annotation_style: Option<&AnnotationStyle>,
    ) where
        F: OutputFormat<Output = String>,
    {
        let unassigned_refs: Vec<&Reference> = bibliography
            .iter()
            .filter(|entry| !assigned.contains(&entry.id) && selected.contains(&entry.id))
            .filter_map(|entry| self.bibliography.get(&entry.id))
            .collect();

        if unassigned_refs.is_empty() {
            return;
        }

        // Re-process references to ensure correct author substitution and disambiguation
        // within the unassigned subset.
        let unassigned = self.merge_compound_entries::<F>(self.process_sorted_refs::<_, F>(
            unassigned_refs.into_iter(),
            |reference, entry_number| {
                self.process_bibliography_entry_with_format::<F>(reference, entry_number)
            },
        ));

        if !result.is_empty() {
            result.push_str("\n\n");
        }

        result.push_str(&crate::render::refs_to_string_with_format::<F>(
            unassigned,
            annotations,
            annotation_style,
        ));
    }

    fn render_with_legacy_grouping<F>(
        &self,
        bibliography: &[ProcEntry],
        annotations: Option<&HashMap<String, String>>,
        annotation_style: Option<&AnnotationStyle>,
    ) -> String
    where
        F: OutputFormat<Output = String>,
    {
        let fmt = F::default();
        let cited_ids = self.cited_ids.borrow();
        let cited_entries: Vec<ProcEntry> = bibliography
            .iter()
            .filter(|entry| cited_ids.contains(&entry.id))
            .cloned()
            .collect();

        let mut result = String::new();
        if !cited_entries.is_empty() {
            result.push_str(&crate::render::refs_to_string_with_format::<F>(
                cited_entries,
                annotations,
                annotation_style,
            ));
        }

        fmt.finish(result)
    }

    fn render_bibliography_for_group<F>(
        &self,
        group: &BibliographyGroup,
        annotations: Option<&HashMap<String, String>>,
        annotation_style: Option<&AnnotationStyle>,
    ) -> String
    where
        F: OutputFormat<Output = String>,
    {
        let bibliography = self.sorted_id_stubs();
        let fmt = F::default();
        let cited_ids = self.cited_ids.borrow();
        let evaluator = SelectorEvaluator::new(&cited_ids);
        let sorter = GroupSorter::new(&self.locale);

        let matching_refs =
            self.collect_matching_group_refs(&bibliography, &HashSet::new(), &evaluator, group);

        if matching_refs.is_empty() {
            return fmt.finish(String::new());
        }

        let sorted_refs = if let Some(sort_spec) = &group.sort {
            sorter.sort_references(matching_refs, &sort_spec.resolve())
        } else {
            matching_refs
        };

        let local_hints = self.build_group_local_hints(&sorted_refs, group);
        let entries = self.merge_compound_entries::<F>(self.render_group_entries::<F>(
            &bibliography,
            sorted_refs,
            group,
            local_hints.as_ref(),
        ));

        fmt.finish(crate::render::refs_to_string_with_format::<F>(
            entries,
            annotations,
            annotation_style,
        ))
    }

    /// Render the bibliography with grouping for uncited (nocite) items.
    ///
    /// If `style.bibliography.groups` is defined, uses configurable grouping
    /// with per-group sorting. Group selectors apply to individual references
    /// before compound numeric rows are merged, so each rendered group only
    /// includes the members that matched its selector. Otherwise, falls back to
    /// hardcoded cited/uncited grouping for backward compatibility.
    pub fn render_grouped_bibliography_with_format<F>(&self) -> String
    where
        F: OutputFormat<Output = String>,
    {
        self.render_grouped_bibliography_with_format_and_annotations::<F>(None, None)
    }

    /// Render the bibliography with grouping and annotations.
    pub fn render_grouped_bibliography_with_format_and_annotations<F>(
        &self,
        annotations: Option<&HashMap<String, String>>,
        annotation_style: Option<&AnnotationStyle>,
    ) -> String
    where
        F: OutputFormat<Output = String>,
    {
        if let Some(groups) = self
            .style
            .bibliography
            .as_ref()
            .and_then(|bibliography| bibliography.groups.as_ref())
        {
            let id_stubs = self.sorted_id_stubs();
            let selected = id_stubs
                .iter()
                .map(|e| e.id.clone())
                .collect::<HashSet<_>>();
            return self.render_with_custom_groups_filtered::<F>(
                &id_stubs,
                groups,
                &selected,
                annotations,
                annotation_style,
            );
        }

        let bibliography_options = self.get_bibliography_options();
        if let Some(partitioning) = bibliography_options.sort_partitioning.as_ref()
            && crate::sort_partitioning::should_render_sections(partitioning)
        {
            self.initialize_numeric_bibliography_numbers();
            let sorted_refs = self.sort_references(self.bibliography.values().collect());
            return self.render_with_partition_sections::<F>(
                sorted_refs,
                partitioning,
                annotations,
                annotation_style,
            );
        }

        let all_entries = self.process_references().bibliography;
        self.render_with_legacy_grouping::<F>(
            &self.merge_compound_entries::<F>(all_entries),
            annotations,
            annotation_style,
        )
    }

    /// Render frontmatter-defined bibliography groups for document output.
    ///
    /// This uses the same pre-merge selector semantics as
    /// [`Self::render_grouped_bibliography_with_format`].
    pub(crate) fn render_document_bibliography_groups<F>(
        &self,
        groups: &[BibliographyGroup],
    ) -> String
    where
        F: OutputFormat<Output = String>,
    {
        let all_entries = self.sorted_id_stubs();
        self.render_with_custom_groups::<F>(&all_entries, groups)
    }

    /// Render one bibliography block for document output.
    ///
    /// Returns heading and body separately so callers can insert headings
    /// in their own output format.
    pub(crate) fn render_document_bibliography_block<F>(
        &self,
        group: &BibliographyGroup,
    ) -> RenderedBibliographyGroup
    where
        F: OutputFormat<Output = String>,
    {
        let mut headingless = group.clone();
        let heading = headingless
            .heading
            .take()
            .and_then(|group_heading| self.resolve_group_heading(&group_heading));
        let body = self.render_bibliography_for_group::<F>(&headingless, None, None);

        RenderedBibliographyGroup { heading, body }
    }

    pub(super) fn extract_metadata(&self, reference: &Reference) -> ProcEntryMetadata {
        let bibliography_config = self.get_bibliography_config();
        let options = RenderOptions {
            config: &bibliography_config,
            bibliography_config: Some(self.get_bibliography_options().into_owned()),
            locale: &self.locale,
            context: RenderContext::Bibliography,
            mode: citum_schema::citation::CitationMode::NonIntegral,
            suppress_author: false,
            locator_raw: None,
            ref_type: None,
            show_semantics: self.show_semantics,
            current_template_index: None,
            abbreviation_map: self.abbreviation_map.as_ref(),
        };

        ProcEntryMetadata {
            author: reference
                .author()
                .map(|authors| format_contributors_short(&authors.to_names_vec(), &options)),
            year: reference
                .csl_issued_date()
                .map(|issued| issued.year().clone()),
            title: reference.title().map(|title| title.to_string()),
        }
    }

    fn render_group_heading<F>(&self, heading: &str) -> String
    where
        F: OutputFormat<Output = String>,
    {
        if std::any::type_name::<F>() == std::any::type_name::<crate::render::html::Html>() {
            return format!("<h2>{heading}</h2>\n\n");
        }

        format!("# {heading}\n\n")
    }
}