//! Owned, bounded queries over exact current source buffers.
//!
//! Includes are inspected from these bytes only. A previously validated bundle
//! is an equality witness, never a source of missing text or navigation spans.
//! The caller must advance the workspace key on every indexed-input change and
//! reject asynchronous results whose key, query serial or byte readback changed.

use super::symbols::{self, Index};
use oa_core::Id;
use oa_core::language::{self, OutlineLimits, Span, TokenClass, TokenSpan};
use oa_core::source_project::{self, IncludeLocation, ProjectLimits};
use std::collections::BTreeMap;

pub use super::symbols::{Evidence, Identity, Membership, Position, SnapshotKey};

#[derive(Debug, Clone, Copy)]
pub struct DraftFile<'a> {
    pub path: &'a str,
    pub text: &'a str,
}

/// The caller asserts that this complete source bundle passed explicit semantic
/// validation. Capture independently compares its entry, membership and bytes.
#[derive(Debug, Clone, Copy)]
pub struct ValidatedSources<'a> {
    pub entry: &'a str,
    pub files: &'a [DraftFile<'a>],
}

#[derive(Debug, Clone, Copy)]
pub struct Input<'a> {
    pub key: SnapshotKey,
    pub entry: &'a str,
    pub files: &'a [DraftFile<'a>],
    pub applied: bool,
    pub validation: Option<ValidatedSources<'a>>,
}

#[derive(Debug, Clone, Copy)]
pub struct Limits {
    pub index: symbols::Limits,
    pub include_depth: usize,
    pub include_edges: usize,
    pub result_rows: usize,
    /// Upper bound on copied result strings, including quoted name input before
    /// JSON decoding. Row/vector overhead is separately bounded by result_rows.
    pub result_bytes: usize,
}

impl Default for Limits {
    fn default() -> Self {
        Self {
            index: symbols::Limits {
                files: 256,
                source_bytes: 64 * 1024 * 1024,
                path_bytes: 64 * 1024,
                tokens: 1_000_000,
                records: 250_000,
            },
            include_depth: 32,
            include_edges: 1024,
            result_rows: 200,
            result_bytes: 256 * 1024,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AdmissionError {
    Index(symbols::BuildError),
    Outline {
        file: usize,
        error: language::OutlineLimit,
    },
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ScopeIssueKind {
    MissingEntry,
    InvalidIncludes,
    MissingInclude,
    IncludeCycle,
    IncludeDepth,
    IncludeBudget,
}

/// File indices refer to Snapshot::file. Include spans retain the original
/// quoted include location; diagnostics do not copy unbounded source strings.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ScopeIssue {
    pub kind: ScopeIssueKind,
    pub file: Option<usize>,
    pub span: Span,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Scope {
    pub closure_complete: bool,
    pub applied: bool,
    pub validated: bool,
    pub included_files: usize,
    pub detached_files: usize,
    pub unverified_files: usize,
    pub issues: Vec<ScopeIssue>,
}

impl Scope {
    pub fn partial(&self) -> bool {
        !self.closure_complete || !self.applied
    }

    fn index_scope(&self) -> symbols::Scope {
        if !self.applied {
            symbols::Scope::Partial(symbols::PartialReason::UnappliedWorkspace)
        } else if !self.closure_complete {
            let missing = self.issues.iter().any(|issue| {
                matches!(
                    issue.kind,
                    ScopeIssueKind::MissingEntry | ScopeIssueKind::MissingInclude
                )
            });
            symbols::Scope::Partial(if missing {
                symbols::PartialReason::MissingFiles
            } else {
                symbols::PartialReason::IncludesUnverified
            })
        } else {
            symbols::Scope::CompleteProject
        }
    }
}

#[derive(Debug)]
struct File {
    path: String,
    text: String,
    tokens: Vec<TokenSpan>,
    declarations: Vec<symbols::DeclarationInput>,
    membership: Membership,
    evidence: Evidence,
}

/// Immutable and Send: callers can retain it in an Arc for asynchronous work.
/// No Workspace, accepted Project, filesystem handle or mutation is retained.
#[derive(Debug)]
pub struct Snapshot {
    key: SnapshotKey,
    entry: String,
    files: Vec<File>,
    scope: Scope,
    limits: Limits,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct QueryStamp {
    pub snapshot: SnapshotKey,
    pub serial: u64,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Action {
    Definition,
    References,
}

#[derive(Debug, Clone, Copy)]
pub struct Query<'a> {
    pub stamp: QueryStamp,
    pub path: &'a str,
    pub byte_offset: usize,
    pub action: Action,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Disposition {
    NoSymbol,
    NoDefinition,
    /// All candidates require an explicit choice, including a sole draft or
    /// detached candidate. Truncation can never imply unique resolution.
    Choices,
    Direct,
    LexicalReferences,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Declaration {
    pub alias: String,
    pub kind: String,
    pub name: Option<String>,
    pub span: Span,
    pub authored_id: Option<Id>,
    pub identity: Identity,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DefinitionRow {
    pub path: String,
    /// The alias token to select, rather than an enclosing declaration body.
    pub span: Span,
    pub position: Position,
    pub end_position: Position,
    pub declaration: Declaration,
    pub membership: Membership,
    pub evidence: Evidence,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ReferenceRow {
    pub path: String,
    pub alias: String,
    pub span: Span,
    /// Separate live @ token: intervening comments are never selected as a hit.
    pub marker: Span,
    pub position: Position,
    pub end_position: Position,
    /// Innermost tolerant declaration containing this occurrence, if available.
    /// This describes its lexical location, not its semantic target or owner.
    pub context: Option<Declaration>,
    pub membership: Membership,
    pub evidence: Evidence,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Results {
    pub stamp: QueryStamp,
    pub action: Action,
    pub origin_path: String,
    pub origin_offset: usize,
    pub alias: Option<String>,
    pub scope: Scope,
    pub disposition: Disposition,
    pub definitions: Vec<DefinitionRow>,
    pub references: Vec<ReferenceRow>,
    pub total_definitions: usize,
    pub total_references: usize,
    pub omitted_definitions: usize,
    pub omitted_references: usize,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum QueryError {
    Index(symbols::BuildError),
    Query(symbols::QueryError),
    MissingFile,
    ResultEnvelopeLimit,
}

impl From<symbols::QueryError> for QueryError {
    fn from(error: symbols::QueryError) -> Self {
        Self::Query(error)
    }
}

fn charge(
    total: &mut usize,
    value: usize,
    limit: usize,
    resource: symbols::Resource,
) -> Result<(), AdmissionError> {
    let next =
        total
            .checked_add(value)
            .ok_or(AdmissionError::Index(symbols::BuildError::Limit {
                resource,
                limit,
                observed_at_least: usize::MAX,
            }))?;
    if next > limit {
        return Err(AdmissionError::Index(symbols::BuildError::Limit {
            resource,
            limit,
            observed_at_least: next,
        }));
    }
    *total = next;
    Ok(())
}

/// No path normalization, outline extraction, accepted-source copy or index
/// allocation takes place until aggregate raw input admission succeeds.
fn admit_raw(input: Input<'_>, limits: Limits) -> Result<(), AdmissionError> {
    charge(
        &mut 0,
        input.files.len(),
        limits.index.files,
        symbols::Resource::Files,
    )?;
    let mut bytes = 0;
    let mut paths = 0;
    charge(
        &mut paths,
        input.entry.len(),
        limits.index.path_bytes,
        symbols::Resource::PathBytes,
    )?;
    for file in input.files {
        charge(
            &mut bytes,
            file.text.len(),
            limits.index.source_bytes,
            symbols::Resource::SourceBytes,
        )?;
        charge(
            &mut paths,
            file.path.len(),
            limits.index.path_bytes,
            symbols::Resource::PathBytes,
        )?;
    }
    Ok(())
}

struct Extracted {
    tokens: Vec<TokenSpan>,
    declarations: Vec<symbols::DeclarationInput>,
}

#[derive(Debug)]
struct OwnedFile {
    path: String,
    text: String,
}

#[derive(Debug, Default)]
struct ValidationProof {
    admitted: bool,
    file_count: usize,
    matches: Vec<bool>,
}

impl ValidationProof {
    fn capture(input: Input<'_>, limits: Limits) -> Self {
        let Some(validation) = input.validation else {
            return Self::default();
        };
        if validation.entry != input.entry
            || admit_raw(
                Input {
                    files: validation.files,
                    entry: validation.entry,
                    validation: None,
                    ..input
                },
                limits,
            )
            .is_err()
        {
            return Self::default();
        }
        let files: BTreeMap<_, _> = validation
            .files
            .iter()
            .map(|file| (file.path, file.text))
            .collect();
        if files.len() != validation.files.len() {
            return Self::default();
        }
        Self {
            admitted: true,
            file_count: files.len(),
            matches: input
                .files
                .iter()
                .map(|file| files.get(file.path).is_some_and(|text| *text == file.text))
                .collect(),
        }
    }

    fn matches(&self, memberships: &[Membership], scope: &Scope) -> bool {
        self.admitted
            && scope.closure_complete
            && scope.applied
            && self.file_count == scope.included_files
            && self.matches.len() == memberships.len()
            && memberships
                .iter()
                .zip(&self.matches)
                .all(|(membership, equal)| *membership != Membership::Included || *equal)
    }
}

/// Request-time data for spawn_blocking: raw admission plus one current-text
/// clone, with no lexer/parser work. The private equality proof contains only
/// booleans/counts, never a second accepted SourceBundle. These private fields
/// prevent changing the bytes after their equality evidence was captured.
#[derive(Debug)]
pub struct OwnedInput {
    key: SnapshotKey,
    entry: String,
    files: Vec<OwnedFile>,
    applied: bool,
    validation: ValidationProof,
    limits: Limits,
}

pub fn admit_owned(input: Input<'_>, limits: Limits) -> Result<OwnedInput, AdmissionError> {
    admit_raw(input, limits)?;
    let validation = ValidationProof::capture(input, limits);
    Ok(OwnedInput {
        key: input.key,
        entry: input.entry.to_owned(),
        files: input
            .files
            .iter()
            .map(|file| OwnedFile {
                path: file.path.to_owned(),
                text: file.text.to_owned(),
            })
            .collect(),
        applied: input.applied,
        validation,
        limits,
    })
}

impl OwnedInput {
    /// Run on the blocking worker. Token/declaration/reference budgets are
    /// enforced before extraction/index allocation, then strings are moved into
    /// Snapshot without another copy. A limit error never returns partial data.
    pub fn capture(self) -> Result<Snapshot, AdmissionError> {
        let files: Vec<_> = self
            .files
            .iter()
            .map(|file| DraftFile {
                path: &file.path,
                text: &file.text,
            })
            .collect();
        let input = Input {
            key: self.key,
            entry: &self.entry,
            files: &files,
            applied: self.applied,
            validation: None,
        };
        let prepared = prepare(input, self.limits, Some(&self.validation))?;
        Ok(finish(
            self.key,
            self.entry,
            self.files,
            self.limits,
            prepared,
        ))
    }
}

struct Prepared {
    extracted: Vec<Extracted>,
    memberships: Vec<Membership>,
    scope: Scope,
}

#[cfg(test)]
pub fn capture(input: Input<'_>, limits: Limits) -> Result<Snapshot, AdmissionError> {
    let prepared = prepare(input, limits, None)?;
    let files = input
        .files
        .iter()
        .map(|file| OwnedFile {
            path: file.path.to_owned(),
            text: file.text.to_owned(),
        })
        .collect();
    Ok(finish(
        input.key,
        input.entry.to_owned(),
        files,
        limits,
        prepared,
    ))
}

fn prepare(
    input: Input<'_>,
    limits: Limits,
    proof: Option<&ValidationProof>,
) -> Result<Prepared, AdmissionError> {
    admit_raw(input, limits)?;
    let mut extracted = Vec::with_capacity(input.files.len());
    let mut tokens_left = limits.index.tokens;
    let mut records_left = limits.index.records;
    for (file, draft) in input.files.iter().enumerate() {
        // Remaining aggregate budgets are passed before this file is lexed.
        // A stopped outline returns an error, never a complete zero-result file.
        let outline = language::outline(
            draft.text,
            OutlineLimits {
                max_source_bytes: limits.index.source_bytes,
                max_tokens: tokens_left,
                max_declarations: records_left,
            },
        )
        .map_err(|error| AdmissionError::Outline { file, error })?;
        tokens_left -= outline.tokens.len();
        records_left -= outline.declarations.len();
        extracted.push(Extracted {
            tokens: outline.tokens,
            declarations: outline
                .declarations
                .into_iter()
                .map(|declaration| symbols::DeclarationInput {
                    span: declaration.declaration.span,
                    kind_span: declaration.kind_span,
                    alias_span: declaration.alias_span,
                    identity: declaration
                        .declaration
                        .id
                        .zip(declaration.id_literal_span)
                        .map(|(id, literal_span)| symbols::AuthoredIdentity { id, literal_span }),
                })
                .collect(),
        });
    }
    // Reuse the reviewed live-reference counter and structural admission, rather
    // than maintain a second lexer/reference recognizer. No text has been copied.
    let inputs: Vec<_> = input
        .files
        .iter()
        .zip(&extracted)
        .map(|(file, outline)| symbols::FileInput {
            path: file.path,
            text: file.text,
            tokens: &outline.tokens,
            declarations: &outline.declarations,
            syntax: symbols::SyntaxCoverage::FullFile,
            membership: Membership::Unverified,
            evidence: Evidence::Draft,
        })
        .collect();
    Index::build(
        input.key,
        &inputs,
        symbols::Scope::Partial(symbols::PartialReason::IncludesUnverified),
        limits.index,
    )
    .map_err(AdmissionError::Index)?;
    let (memberships, mut scope) = inspect_closure(input, limits);
    scope.validated = proof.map_or_else(
        || validation_matches(input, &memberships, &scope, limits),
        |proof| proof.matches(&memberships, &scope),
    );
    Ok(Prepared {
        extracted,
        memberships,
        scope,
    })
}

fn finish(
    key: SnapshotKey,
    entry: String,
    files: Vec<OwnedFile>,
    limits: Limits,
    prepared: Prepared,
) -> Snapshot {
    let Prepared {
        extracted,
        memberships,
        scope,
    } = prepared;
    let files = files
        .into_iter()
        .zip(extracted)
        .zip(memberships)
        .map(|((draft, outline), membership)| File {
            path: draft.path,
            text: draft.text,
            tokens: outline.tokens,
            declarations: outline.declarations,
            membership,
            evidence: if scope.validated && membership == Membership::Included {
                Evidence::Validated
            } else {
                Evidence::Draft
            },
        })
        .collect();
    Snapshot {
        key,
        entry,
        files,
        scope,
        limits,
    }
}

fn validation_matches(
    input: Input<'_>,
    memberships: &[Membership],
    scope: &Scope,
    limits: Limits,
) -> bool {
    if !scope.closure_complete || !input.applied {
        return false;
    }
    ValidationProof::capture(input, limits).matches(memberships, scope)
}

struct Frame {
    file: usize,
    depth: usize,
    includes: std::vec::IntoIter<IncludeLocation>,
}

fn inspect_closure(input: Input<'_>, limits: Limits) -> (Vec<Membership>, Scope) {
    let paths: BTreeMap<_, _> = input
        .files
        .iter()
        .enumerate()
        .map(|(index, file)| (file.path, index))
        .collect();
    let mut memberships = vec![Membership::Unverified; input.files.len()];
    let mut states = vec![0_u8; input.files.len()];
    let mut issues = Vec::new();
    let mut stack: Vec<Frame> = Vec::new();
    let mut edges_left = limits.include_edges;
    let mut pending = paths.get(input.entry).copied().map(|file| (file, 0, None));
    if pending.is_none() {
        issues.push(ScopeIssue {
            kind: ScopeIssueKind::MissingEntry,
            file: None,
            span: Span::default(),
        });
    }
    loop {
        if let Some((file, depth, origin)) = pending.take() {
            memberships[file] = Membership::Included;
            let (at_file, at_span) = origin.unwrap_or((file, Span::default()));
            let problem = if states[file] == 1 {
                Some(ScopeIssueKind::IncludeCycle)
            } else if depth > limits.include_depth {
                Some(ScopeIssueKind::IncludeDepth)
            } else {
                None
            };
            if let Some(kind) = problem {
                issues.push(ScopeIssue {
                    kind,
                    file: Some(at_file),
                    span: at_span,
                });
            } else if states[file] == 0 {
                states[file] = 1;
                let draft = input.files[file];
                let project_limits = ProjectLimits {
                    max_files: limits.index.files,
                    max_include_depth: limits.include_depth,
                    max_include_edges: edges_left,
                    model: language::ParseLimits {
                        max_source_bytes: limits.index.source_bytes,
                        max_graph_objects: limits.index.records,
                        max_records: limits.index.records,
                        max_step_references: limits.index.records,
                    },
                };
                match source_project::inspect_includes(
                    draft.path,
                    draft.text,
                    draft.path == input.entry,
                    project_limits,
                ) {
                    Ok(includes) => {
                        edges_left -= includes.len();
                        stack.push(Frame {
                            file,
                            depth,
                            includes: includes.into_iter(),
                        });
                    }
                    Err(diagnostics) => {
                        let first = diagnostics.first();
                        let kind = if first.is_some_and(|d| d.diagnostic.code == "project_limit") {
                            ScopeIssueKind::IncludeBudget
                        } else {
                            ScopeIssueKind::InvalidIncludes
                        };
                        issues.push(ScopeIssue {
                            kind,
                            file: Some(file),
                            span: first.map_or(Span::default(), |d| d.diagnostic.span),
                        });
                        states[file] = 2;
                    }
                }
            }
        }
        let Some(frame) = stack.last_mut() else {
            break;
        };
        if let Some(include) = frame.includes.next() {
            if let Some(&file) = paths.get(include.target.as_str()) {
                pending = Some((
                    file,
                    frame.depth.saturating_add(1),
                    Some((frame.file, include.location.span)),
                ));
            } else {
                issues.push(ScopeIssue {
                    kind: ScopeIssueKind::MissingInclude,
                    file: Some(frame.file),
                    span: include.location.span,
                });
            }
        } else {
            states[frame.file] = 2;
            stack.pop();
        }
    }
    let complete = issues.is_empty();
    if complete {
        for membership in &mut memberships {
            if *membership == Membership::Unverified {
                *membership = Membership::Detached;
            }
        }
    }
    let scope = Scope {
        closure_complete: complete,
        applied: input.applied,
        validated: false,
        included_files: memberships
            .iter()
            .filter(|m| **m == Membership::Included)
            .count(),
        detached_files: memberships
            .iter()
            .filter(|m| **m == Membership::Detached)
            .count(),
        unverified_files: memberships
            .iter()
            .filter(|m| **m == Membership::Unverified)
            .count(),
        issues,
    };
    (memberships, scope)
}

impl Snapshot {
    pub fn key(&self) -> SnapshotKey {
        self.key
    }

    #[cfg(test)]
    pub fn scope(&self) -> &Scope {
        &self.scope
    }

    #[cfg(test)]
    pub fn file(&self, index: usize) -> Option<DraftFile<'_>> {
        self.files.get(index).map(|file| DraftFile {
            path: &file.path,
            text: &file.text,
        })
    }

    /// Exact loaded file-set/byte comparison, independent of input order. The
    /// evidence outcome is rechecked too; irrelevant old witness text is not kept.
    /// The key check is mandatory even if bytes happen to match after Undo.
    pub fn matches(&self, input: Input<'_>) -> bool {
        if input.key != self.key
            || input.entry != self.entry
            || input.applied != self.scope.applied
            || input.files.len() != self.files.len()
            || admit_raw(input, self.limits).is_err()
        {
            return false;
        }
        let mut memberships = Vec::with_capacity(input.files.len());
        for draft in input.files {
            let Some(file) = self
                .files
                .iter()
                .find(|file| file.path == draft.path && file.text == draft.text)
            else {
                return false;
            };
            // Duplicate paths are not the same file set, even if their bytes match.
            if input
                .files
                .iter()
                .filter(|other| other.path == draft.path)
                .count()
                != 1
            {
                return false;
            }
            memberships.push(file.membership);
        }
        validation_matches(input, &memberships, &self.scope, self.limits) == self.scope.validated
    }

    fn index(&self) -> Result<Index<'_>, QueryError> {
        let inputs: Vec<_> = self
            .files
            .iter()
            .map(|file| symbols::FileInput {
                path: &file.path,
                text: &file.text,
                tokens: &file.tokens,
                declarations: &file.declarations,
                syntax: symbols::SyntaxCoverage::FullFile,
                membership: file.membership,
                evidence: file.evidence,
            })
            .collect();
        Index::build(
            self.key,
            &inputs,
            self.scope.index_scope(),
            self.limits.index,
        )
        .map_err(QueryError::Index)
    }

    pub fn query(&self, query: Query<'_>) -> Result<Results, QueryError> {
        if query.stamp.snapshot != self.key {
            return Err(symbols::QueryError::StaleSnapshot.into());
        }
        let file = self
            .files
            .iter()
            .position(|file| file.path == query.path)
            .ok_or(QueryError::MissingFile)?;
        let index = self.index()?;
        let hit = index.hit(file, query.byte_offset)?;
        let alias = match hit {
            Some(symbols::Hit::Definition(handle)) => Some(index.definition(handle)?.alias),
            Some(symbols::Hit::Reference(handle)) => Some(index.reference(handle)?.alias),
            None => None,
        };
        let mut budget = RowBudget {
            rows: self.limits.result_rows,
            bytes: self.limits.result_bytes,
        };
        if !budget.strings(query.path.len().saturating_add(alias.map_or(0, str::len))) {
            return Err(QueryError::ResultEnvelopeLimit);
        }
        let mut result = Results {
            stamp: query.stamp,
            action: query.action,
            origin_path: query.path.to_owned(),
            origin_offset: query.byte_offset,
            alias: alias.map(str::to_owned),
            scope: self.scope.clone(),
            disposition: Disposition::NoSymbol,
            definitions: Vec::new(),
            references: Vec::new(),
            total_definitions: 0,
            total_references: 0,
            omitted_definitions: 0,
            omitted_references: 0,
        };
        let Some(alias) = alias else {
            return Ok(result);
        };
        result.total_definitions = index
            .definitions_for_alias(alias, symbols::Domain::Loaded)
            .count();
        for handle in index.definitions_for_alias(alias, symbols::Domain::Loaded) {
            let definition = index.definition(handle)?;
            let source = &self.files[definition.declaration.file];
            let display = self.display(&index, handle)?;
            if !budget.row(source.path.len().saturating_add(display.bytes())) {
                break;
            }
            result.definitions.push(DefinitionRow {
                path: source.path.clone(),
                span: definition.alias_location.span,
                position: Position {
                    line: 0,
                    byte_column: 0,
                },
                end_position: Position {
                    line: 0,
                    byte_column: 0,
                },
                declaration: display.owned(),
                membership: source.membership,
                evidence: source.evidence,
            });
        }
        result.omitted_definitions = result.total_definitions - result.definitions.len();
        result.disposition = if result.total_definitions == 0 {
            Disposition::NoDefinition
        } else {
            Disposition::Choices
        };
        if query.action == Action::Definition
            && result.total_definitions == 1
            && result.definitions.len() == 1
            && self.scope.validated
            && self.files[file].membership == Membership::Included
        {
            let only = &result.definitions[0];
            if only.membership == Membership::Included
                && only.evidence == Evidence::Validated
                && !matches!(only.declaration.identity, Identity::Collision { .. })
            {
                result.disposition = Disposition::Direct;
            }
        }
        if query.action == Action::References {
            result.disposition = Disposition::LexicalReferences;
            result.total_references = index
                .references_for_alias(alias, symbols::Domain::Loaded)
                .count();
            for handle in index.references_for_alias(alias, symbols::Domain::Loaded) {
                let reference = index.reference(handle)?;
                let source = &self.files[reference.alias_location.file];
                let context = self.context(&index, reference.alias_location)?;
                let bytes = source
                    .path
                    .len()
                    .saturating_add(alias.len())
                    .saturating_add(context.as_ref().map_or(0, Display::bytes));
                if !budget.row(bytes) {
                    break;
                }
                result.references.push(ReferenceRow {
                    path: source.path.clone(),
                    alias: alias.to_owned(),
                    span: reference.alias_location.span,
                    marker: reference.at.span,
                    position: Position {
                        line: 0,
                        byte_column: 0,
                    },
                    end_position: Position {
                        line: 0,
                        byte_column: 0,
                    },
                    context: context.map(Display::owned),
                    membership: source.membership,
                    evidence: source.evidence,
                });
            }
            result.omitted_references = result.total_references - result.references.len();
        }
        self.locate_rows(&mut result);
        Ok(result)
    }

    fn locate_rows(&self, result: &mut Results) {
        // Index admission already proved every row span and UTF-8 boundary.
        // Scan each selected source prefix once instead of rescanning a large
        // buffer for every occurrence. Scratch has at most two entries per row.
        for file in &self.files {
            let mut offsets = Vec::new();
            for (row, definition) in result.definitions.iter().enumerate() {
                if definition.path == file.path {
                    offsets.push((definition.span.start, true, row, false));
                    offsets.push((definition.span.end, true, row, true));
                }
            }
            for (row, reference) in result.references.iter().enumerate() {
                if reference.path == file.path {
                    offsets.push((reference.span.start, false, row, false));
                    offsets.push((reference.span.end, false, row, true));
                }
            }
            offsets.sort_unstable_by_key(|&(offset, _, _, _)| offset);
            let (mut scanned, mut line, mut line_start) = (0, 0, 0);
            for (offset, definition, row, end) in offsets {
                for (relative, byte) in file.text.as_bytes()[scanned..offset].iter().enumerate() {
                    if *byte == b'\n' {
                        line += 1;
                        line_start = scanned + relative + 1;
                    }
                }
                scanned = offset;
                let position = Position {
                    line,
                    byte_column: offset - line_start,
                };
                match (definition, end) {
                    (true, false) => result.definitions[row].position = position,
                    (true, true) => result.definitions[row].end_position = position,
                    (false, false) => result.references[row].position = position,
                    (false, true) => result.references[row].end_position = position,
                }
            }
        }
    }

    fn display<'a>(
        &'a self,
        index: &Index<'a>,
        handle: symbols::DefinitionHandle,
    ) -> Result<Display<'a>, QueryError> {
        let definition = index.definition(handle)?;
        let file = &self.files[definition.declaration.file];
        let alias_end = definition.alias_location.span.end;
        let next = file
            .tokens
            .partition_point(|token| token.span.start < alias_end);
        // The coordinated outline has already proved the declaration and alias
        // roles. This optional decoration only decodes its immediate name token.
        let name = (definition.kind != "node")
            .then(|| {
                file.tokens[next..].iter().find(|token| {
                    !matches!(token.class, TokenClass::Whitespace | TokenClass::Comment)
                })
            })
            .flatten()
            .filter(|token| {
                token.class == TokenClass::String
                    && token.span.end <= definition.declaration.span.end
            })
            .map(|token| &file.text[token.span.start..token.span.end]);
        Ok(Display {
            alias: definition.alias,
            kind: definition.kind,
            quoted_name: name,
            span: definition.declaration.span,
            authored_id: definition.authored_id,
            identity: index.identity(handle)?,
        })
    }

    fn context<'a>(
        &'a self,
        index: &Index<'a>,
        location: symbols::Location,
    ) -> Result<Option<Display<'a>>, QueryError> {
        let file = &self.files[location.file];
        let inner = file
            .declarations
            .iter()
            .filter(|declaration| {
                declaration.span.start <= location.span.start
                    && location.span.end <= declaration.span.end
            })
            .min_by_key(|declaration| declaration.span.end - declaration.span.start);
        let Some(inner) = inner else {
            return Ok(None);
        };
        let Some(symbols::Hit::Definition(handle)) =
            index.hit(location.file, inner.alias_span.start)?
        else {
            return Ok(None);
        };
        self.display(index, handle).map(Some)
    }
}

struct Display<'a> {
    alias: &'a str,
    kind: &'a str,
    quoted_name: Option<&'a str>,
    span: Span,
    authored_id: Option<Id>,
    identity: Identity,
}

impl Display<'_> {
    fn bytes(&self) -> usize {
        self.alias
            .len()
            .saturating_add(self.kind.len())
            .saturating_add(self.quoted_name.map_or(0, str::len))
    }

    fn owned(self) -> Declaration {
        Declaration {
            alias: self.alias.to_owned(),
            kind: self.kind.to_owned(),
            name: self
                .quoted_name
                .and_then(|name| serde_json::from_str(name).ok()),
            span: self.span,
            authored_id: self.authored_id,
            identity: self.identity,
        }
    }
}

struct RowBudget {
    rows: usize,
    bytes: usize,
}

impl RowBudget {
    fn strings(&mut self, bytes: usize) -> bool {
        if bytes > self.bytes {
            return false;
        }
        self.bytes -= bytes;
        true
    }

    fn row(&mut self, bytes: usize) -> bool {
        if self.rows == 0 || !self.strings(bytes) {
            return false;
        }
        self.rows -= 1;
        true
    }
}

#[cfg(test)]
#[path = "query_tests.rs"]
mod tests;