folio

use comemo::Tracked;
use ecow::{EcoString, EcoVec, eco_format};
use typst_syntax::Span;

use super::{History, Introspect};
use crate::diag::{HintedStrResult, SourceDiagnostic, StrResult, warning};
use crate::engine::Engine;
use crate::foundations::{
    Array, Content, Context, Label, LocatableSelector, Repr, Selector, Value, func,
};
use crate::introspection::Introspector;

/// Finds elements in the document.
///
/// The `query` function lets you search your document for elements of a
/// particular type or with a particular label. To use it, you first need to
/// ensure that [context] is available.
///

/// # Finding elements
/// In the example below, we manually create a table of contents instead of
/// using the [`outline`] function.
///
/// To do this, we first query for all headings in the document at level 1 and
/// where `outlined` is true. Querying only for headings at level 1 ensures
/// that, for the purpose of this example, sub-headings are not included in the
/// table of contents. The `outlined` field is used to exclude the "Table of
/// Contents" heading itself.
///
/// Note that we open a `context` to be able to use the `query` function.
///
/// ```example
/// >>> #set page(
/// >>>  width: 240pt,
/// >>>  height: 180pt,
/// >>>  margin: (top: 20pt, bottom: 35pt)
/// >>> )
/// #set page(numbering: "1")
///
/// #heading(outlined: false)[
///   Table of Contents
/// ]
/// #context {
///   let chapters = query(
///     heading.where(
///       level: 1,
///       outlined: true,
///     )
///   )
///   for chapter in chapters {
///     let loc = chapter.location()
///     let nr = counter(page).display(at: loc)
///     [#chapter.body #h(1fr) #nr \ ]
///   }
/// }
///
/// = Introduction
/// #lorem(10)
/// #pagebreak()
///
/// == Sub-Heading
/// #lorem(8)
///
/// = Discussion
/// #lorem(18)
/// ```
///
/// To get the page numbers, we first get the location of the elements returned
/// by `query` with [`location`]($content.location). We then also retrieve the
/// [page numbering]($location.page-numbering) and [page
/// counter]($counter/#page-counter) at that location and apply the numbering to
/// the counter.
///
/// # A word of caution { #caution }
/// To resolve all your queries, Typst evaluates and layouts parts of the
/// document multiple times. However, there is no guarantee that your queries
/// can actually be completely resolved. If you aren't careful a query can
/// affect itself—leading to a result that never stabilizes.
///
/// In the example below, we query for all headings in the document. We then
/// generate as many headings. In the beginning, there's just one heading,
/// titled `Real`. Thus, `count` is `1` and one `Fake` heading is generated.
/// Typst sees that the query's result has changed and processes it again. This
/// time, `count` is `2` and two `Fake` headings are generated. This goes on and
/// on. As we can see, the output has a finite amount of headings. This is
/// because Typst simply gives up after a few attempts.
///
/// In general, you should try not to write queries that affect themselves. The
/// same words of caution also apply to other introspection features like
/// [counters]($counter) and [state].
///
/// ```example
/// = Real
/// #context {
///   let elems = query(heading)
///   let count = elems.len()
///   count * [= Fake]
/// }
/// ```
///
/// # Command line queries
/// You can also perform queries from the command line, using the `typst eval`
/// command. This command evaluates Typst code, potentially in the context of a
/// document, and outputs the resulting value in serialized form. It takes the
/// code to evaluate as its first argument and (optionally) the path to a
/// document via `--in`.
///
/// Consider the following `example.typ` file which contains some invisible
/// [metadata]:
///
/// ```typ
/// #metadata("This is a note") <note>
/// ```
///
/// You can execute a query on it as follows using Typst's CLI.
/// ```sh
/// $ typst eval 'query(<note>)' --in example.typ
/// [
///   {
///     "func": "metadata",
///     "value": "This is a note",
///     "label": "<note>"
///   }
/// ]
/// ```
///
/// This command tells Typst to compile `example.typ` and then run the code
/// `{query(<note>)}` with access to the resulting document.
///
/// **Note:** The code is surrounded with quotes to avoid special characters
/// being interpreted by the shell. How to quote strings depends on your
/// platform/shell.
///
/// ## Retrieving a specific field
///
/// Frequently, you're interested in only one specific field of the resulting
/// elements. In the case of the `metadata` element, the `value` field is the
/// interesting one. You can extract just this field by adjusting the code.
///
/// ```sh
/// $ typst eval 'query(<note>).map(it => it.value)' --in example.typ
/// ["This is a note"]
/// ```
///
/// If you are interested in just a single element, you can also use the
/// [`first()`]($array.first) method to extract just it.
///
/// ```sh
/// $ typst eval 'query(<note>).first().value' --in example.typ
/// "This is a note"
/// ```
///
/// ## Querying for a specific export target
///
/// In case you need to query a document when exporting for a specific target,
/// you can use the `--target` argument. Valid values are `paged`, and `html`
/// (if the [`html`] feature is enabled).
#[func(contextual)]
pub fn query(
    engine: &mut Engine,
    context: Tracked<Context>,
    span: Span,
    /// Can be
    /// - an element function like a `heading` or `figure`,
    /// - a `{<label>}`,
    /// - a more complex selector like `{heading.where(level: 1)}`,
    /// - or `{selector(heading).before(here())}`.
    ///
    /// Only [locatable]($location/#locatable) element functions are supported.
    target: LocatableSelector,
) -> HintedStrResult<Array> {
    context.introspect()?;
    let vec = engine.introspect(QueryIntrospection(target.0, span));
    Ok(vec.into_iter().map(Value::Content).collect())
}

/// Retrieves all matches of a selector in the document.
#[derive(Debug, Clone, PartialEq, Hash)]
pub struct QueryIntrospection(pub Selector, pub Span);

impl Introspect for QueryIntrospection {
    type Output = EcoVec<Content>;

    fn introspect(
        &self,
        _: &mut Engine,
        introspector: Tracked<dyn Introspector + '_>,
    ) -> Self::Output {
        introspector.query(&self.0)
    }

    fn diagnose(&self, history: &History<Self::Output>) -> SourceDiagnostic {
        let lengths = history.as_ref().map(|vec| vec.len());
        let things = format_selector(&self.0, "elements");
        let what = if !lengths.converged() {
            eco_format!("number of {things}")
        } else {
            eco_format!("query for {things}")
        };
        format_convergence_warning(self.1, &lengths, &what)
    }
}

/// Retrieves the first match of a selector in the document.
#[derive(Debug, Clone, PartialEq, Hash)]
pub struct QueryFirstIntrospection(pub Selector, pub Span);

impl Introspect for QueryFirstIntrospection {
    type Output = Option<Content>;

    fn introspect(
        &self,
        _: &mut Engine,
        introspector: Tracked<dyn Introspector + '_>,
    ) -> Self::Output {
        introspector.query_first(&self.0)
    }

    fn diagnose(&self, history: &History<Self::Output>) -> SourceDiagnostic {
        let lengths = history.as_ref().map(|vec| vec.is_some() as usize);
        let thing = format_selector(&self.0, "element");
        let what = eco_format!("query for the first {thing}");
        format_convergence_warning(self.1, &lengths, &what)
    }
}

/// Retrieves the only match of a selector in the document.
///
/// Fails if there are multiple occurrences.
#[derive(Debug, Clone, PartialEq, Hash)]
pub struct QueryUniqueIntrospection(pub Selector, pub Span);

impl Introspect for QueryUniqueIntrospection {
    type Output = StrResult<Content>;

    fn introspect(
        &self,
        _: &mut Engine,
        introspector: Tracked<dyn Introspector + '_>,
    ) -> Self::Output {
        introspector.query_unique(&self.0)
    }

    fn diagnose(&self, history: &History<Self::Output>) -> SourceDiagnostic {
        let lengths = history.as_ref().map(|vec| vec.is_ok() as usize);
        let thing = format_selector(&self.0, "element");
        let what = eco_format!("query for a unique {thing}");
        format_convergence_warning(self.1, &lengths, &what)
    }
}

/// Retrieves the only occurrence of a label in the document.
///
/// Fails if there are multiple occurrences.
#[derive(Debug, Clone, PartialEq, Hash)]
pub struct QueryLabelIntrospection(pub Label, pub Span);

impl Introspect for QueryLabelIntrospection {
    type Output = StrResult<Content>;

    fn introspect(
        &self,
        _: &mut Engine,
        introspector: Tracked<dyn Introspector + '_>,
    ) -> Self::Output {
        introspector.query_label(self.0).cloned()
    }

    fn diagnose(&self, history: &History<Self::Output>) -> SourceDiagnostic {
        QueryUniqueIntrospection(Selector::Label(self.0), self.1).diagnose(history)
    }
}

/// The warning when an introspection on a [`Location`] did not converge.
fn format_convergence_warning(
    span: Span,
    lengths: &History<usize>,
    what: &str,
) -> SourceDiagnostic {
    let mut diag = warning!(span, "{what} did not stabilize");
    if !lengths.converged() {
        diag.hint(lengths.hint("numbers of elements", |c| eco_format!("{c}")));
    }
    diag
}

/// Formats a selector human-readably.
fn format_selector(selector: &Selector, kind: &str) -> EcoString {
    match selector {
        Selector::Elem(elem, None) => eco_format!("{} {kind}", elem.name()),
        Selector::Elem(elem, _) => eco_format!("matching {} {kind}", elem.name()),
        Selector::Label(label) => eco_format!("{kind} labelled `{}`", label.repr()),
        other => eco_format!("{kind} matching `{}`", other.repr()),
    }
}