nuwiki/crates/nuwiki-lsp/src/index.rs

//! Workspace index — the in-memory model of every `.wiki` page nuwiki
//! has seen, used to power nav features.
//!
//! The initial scan runs as a background tokio task so
//! the server stays responsive on startup. Per-document updates land here
//! synchronously via `upsert` whenever the LSP backend re-parses on
//! `didOpen` / `didChange`.
//!
//! The index keeps only the projections nav needs (headings + outgoing
//! links) — the full AST already lives in the backend's `documents`
//! `DashMap`.

use std::collections::HashMap;
use std::path::{Path, PathBuf};

use nuwiki_core::ast::{
    BlockNode, BlockquoteNode, DocumentNode, InlineNode, LinkKind, LinkTarget, ListNode, Span,
    TableNode, TagScope,
};
use nuwiki_core::date::DiaryDate;
use tower_lsp::lsp_types::Url;

#[derive(Debug, Clone)]
pub struct IndexedPage {
    pub uri: Url,
    /// Page name relative to the workspace root, without the `.wiki`
    /// extension. Falls back to the URL's last path segment if no root.
    pub name: String,
    pub title: Option<String>,
    pub headings: Vec<HeadingInfo>,
    pub outgoing: Vec<OutgoingLink>,
    /// Every `TagNode` on the page. `tags_by_name` on the
    /// containing `WorkspaceIndex` is the reverse map.
    pub tags: Vec<TagInfo>,
    /// `Some(date)` when this page is a *daily* diary entry —
    /// stem parses as `YYYY-MM-DD`. Equivalent to
    /// `diary_period.and_then(|p| if let DiaryPeriod::Day(d) = p { Some(d) } else { None })`.
    /// Kept as its own field for back-compat with prev/next/list callers
    /// that pre-date the multi-frequency upgrade.
    pub diary_date: Option<DiaryDate>,
    /// Cluster 4 (vimwiki parity): `Some(period)` for any diary entry
    /// (daily, weekly, monthly, yearly). The stem must match the
    /// canonical format for one of the four frequencies.
    pub diary_period: Option<nuwiki_core::date::DiaryPeriod>,
}

/// One tag occurrence on a page. `name` is the bare tag string (no
/// surrounding colons). Multiple tags on the same source line each get
/// their own `TagInfo` so anchor lookup can return a precise location.
#[derive(Debug, Clone)]
pub struct TagInfo {
    pub name: String,
    pub scope: TagScope,
    pub span: Span,
}

#[derive(Debug, Clone)]
pub struct HeadingInfo {
    pub title: String,
    pub level: u8,
    pub anchor: String,
    pub span: Span,
}

impl IndexedPage {
    /// Find a heading whose anchor matches `request`. The on-disk wikilink
    /// syntax lets users write either the slug (`[[Page#some-heading]]`) or
    /// the raw heading text (`[[Page#Some Heading]]`). `slugify` is
    /// idempotent on valid slug input, so applying it to the request
    /// before comparison handles both forms in a single pass.
    pub fn find_heading_by_anchor(&self, request: &str) -> Option<&HeadingInfo> {
        let needle = slugify(request);
        self.headings.iter().find(|h| h.anchor == needle)
    }
}

#[derive(Debug, Clone)]
pub struct OutgoingLink {
    /// Resolved target page name (for `LinkKind::Wiki`) or `None` for kinds
    /// that don't address another page in the workspace (`AnchorOnly`,
    /// `Raw`, `File`, `Local`, `Diary`, …).
    pub target_page: Option<String>,
    pub anchor: Option<String>,
    pub kind: LinkKind,
    pub span: Span,
    /// Mirrors `LinkTarget::is_absolute` so `classify_outgoing` and the
    /// `links` query can distinguish `[[Page]]` (source-relative) from
    /// `[[/Page]]` (root-relative) without re-parsing.
    pub is_absolute: bool,
}

#[derive(Debug, Clone)]
pub struct Backlink {
    pub source_uri: Url,
    pub source_span: Span,
    pub target_anchor: Option<String>,
}

/// A tag's location, used by `WorkspaceIndex::tags_for` for cross-page
/// `[[Page#tag]]` resolution and `nuwiki.tags.search`.
#[derive(Debug, Clone)]
pub struct TagOccurrence {
    pub uri: Url,
    pub span: Span,
    pub scope: TagScope,
}

#[derive(Debug, Default)]
pub struct WorkspaceIndex {
    pub root: Option<PathBuf>,
    /// Subdir relative to `root` that holds diary entries.
    /// Mirrors `WikiConfig::diary_rel_path`; stored here so `upsert` can
    /// classify each indexed URI without a back-reference to the config.
    pub diary_rel_path: Option<String>,
    /// File extension for this wiki's pages (e.g. `".wiki"` or `"wiki"`).
    /// Used to normalise link targets that include the extension so that
    /// `[[page.wiki]]` resolves the same way as `[[page]]`.
    pub file_extension: Option<String>,
    pub pages_by_uri: HashMap<Url, IndexedPage>,
    pub pages_by_name: HashMap<String, Url>,
    pub backlinks: HashMap<String, Vec<Backlink>>,
    /// Tag name → every occurrence across the workspace. Used by
    /// the nav handlers (tag-as-anchor `definition`, `workspace/symbol`
    /// inclusion) and the `nuwiki.tags.search` command.
    pub tags_by_name: HashMap<String, Vec<TagOccurrence>>,
    /// For each source URI, the backlink-bucket keys and tag-bucket names it
    /// contributed. Lets `remove` touch only the affected buckets instead of
    /// scanning every bucket in the workspace (O(touched), not O(total)).
    contributions: HashMap<Url, SourceContributions>,
}

/// The bucket keys a single source page wrote into, so they can be undone
/// in `remove` without a full scan. See [`WorkspaceIndex::contributions`].
#[derive(Debug, Default)]
struct SourceContributions {
    backlink_keys: Vec<String>,
    tag_names: Vec<String>,
}

impl WorkspaceIndex {
    pub fn new(root: Option<PathBuf>) -> Self {
        Self {
            root,
            ..Default::default()
        }
    }

    pub fn with_diary_rel_path(mut self, rel: Option<String>) -> Self {
        self.diary_rel_path = rel;
        self
    }

    pub fn with_file_extension(mut self, ext: Option<String>) -> Self {
        self.file_extension = ext;
        self
    }

    /// Strip the wiki's configured file extension from a link target path.
    /// Handles both `"page.wiki"` → `"page"` and `"subdir/page.wiki"` →
    /// `"subdir/page"`. Returns `path` unchanged when the extension doesn't
    /// match or no extension is configured.
    pub fn strip_link_extension<'a>(&self, path: &'a str) -> &'a str {
        strip_wiki_extension(path, self.file_extension.as_deref())
    }

    /// Insert or update an indexed page. Removes any prior indexing for
    /// the same URI first (handles renames and re-parses without leaking
    /// stale backlinks).
    pub fn upsert(&mut self, uri: Url, ast: &DocumentNode) {
        self.remove(&uri);
        let name = page_name_from_uri(&uri, self.root.as_deref());
        let diary_period =
            diary_period_for_uri(&uri, self.root.as_deref(), self.diary_rel_path.as_deref());
        let diary_date = diary_period.and_then(|p| match p {
            nuwiki_core::date::DiaryPeriod::Day(d) => Some(d),
            _ => None,
        });
        let mut page = IndexedPage {
            uri: uri.clone(),
            name: name.clone(),
            title: ast.metadata.title.clone(),
            headings: Vec::new(),
            outgoing: Vec::new(),
            tags: Vec::new(),
            diary_date,
            diary_period,
        };
        index_blocks(&ast.children, &mut page);

        let mut contrib = SourceContributions::default();
        let ext = self.file_extension.as_deref();
        for link in &page.outgoing {
            if let Some(tgt) = &link.target_page {
                let key = canonical_link_name(tgt, &name, link.is_absolute, ext);
                self.backlinks
                    .entry(key.clone())
                    .or_default()
                    .push(Backlink {
                        source_uri: uri.clone(),
                        source_span: link.span,
                        target_anchor: link.anchor.clone(),
                    });
                contrib.backlink_keys.push(key);
            }
        }

        // Maintain the reverse tag map alongside the per-page list.
        for tag in &page.tags {
            self.tags_by_name
                .entry(tag.name.clone())
                .or_default()
                .push(TagOccurrence {
                    uri: uri.clone(),
                    span: tag.span,
                    scope: tag.scope,
                });
            contrib.tag_names.push(tag.name.clone());
        }
        self.contributions.insert(uri.clone(), contrib);

        self.pages_by_name.insert(name, uri.clone());
        self.pages_by_uri.insert(uri, page);
    }

    pub fn remove(&mut self, uri: &Url) {
        if let Some(page) = self.pages_by_uri.remove(uri) {
            self.pages_by_name.remove(&page.name);
        }
        // Only revisit the buckets this source actually wrote into, instead
        // of scanning every backlink/tag bucket in the workspace.
        if let Some(contrib) = self.contributions.remove(uri) {
            for key in contrib.backlink_keys {
                if let Some(entries) = self.backlinks.get_mut(&key) {
                    entries.retain(|b| &b.source_uri != uri);
                    if entries.is_empty() {
                        self.backlinks.remove(&key);
                    }
                }
            }
            for name in contrib.tag_names {
                if let Some(entries) = self.tags_by_name.get_mut(&name) {
                    entries.retain(|t| &t.uri != uri);
                    if entries.is_empty() {
                        self.tags_by_name.remove(&name);
                    }
                }
            }
        }
    }

    /// All occurrences of a tag across the workspace. Empty slice when the
    /// tag isn't indexed.
    pub fn tags_for(&self, name: &str) -> &[TagOccurrence] {
        self.tags_by_name
            .get(name)
            .map(Vec::as_slice)
            .unwrap_or(&[])
    }

    /// Resolve a `LinkTarget` to a page URI in this workspace.
    /// `source_page` is the name of the page the link originates from —
    /// needed for `AnchorOnly` resolution and source-relative `Wiki` links.
    pub fn resolve(&self, target: &LinkTarget, source_page: &str) -> Option<&Url> {
        match target.kind {
            LinkKind::Wiki => target
                .path
                .as_deref()
                .and_then(|p| self.resolve_wiki_path(p, source_page, target.is_absolute)),
            LinkKind::AnchorOnly => self.pages_by_name.get(source_page),
            // Interwiki / Diary / File / Local / Raw aren't resolved
            // against the workspace index here. The handlers fall back to
            // best-effort resolution (e.g. building file:// URLs).
            _ => None,
        }
    }

    /// Look up a wiki link target by name. Tries source-relative first
    /// (vimwiki's default — `[[Page]]` in `posts/index.wiki` opens
    /// `posts/Page.wiki`) then falls back to root-relative. Skips the
    /// source-relative attempt for explicitly absolute (`[[/Page]]`) and
    /// for source pages already at the wiki root (`source_page` with no
    /// `/`). Strips the configured `.wiki` extension and collapses any
    /// `.`/`..` segments before lookup.
    pub fn resolve_wiki_path(
        &self,
        path: &str,
        source_page: &str,
        is_absolute: bool,
    ) -> Option<&Url> {
        let normalized = self.strip_link_extension(path);
        if !is_absolute {
            if let Some(slash) = source_page.rfind('/') {
                let candidate = collapse_dots(&format!("{}/{}", &source_page[..slash], normalized));
                if let Some(uri) = self.pages_by_name.get(&candidate) {
                    return Some(uri);
                }
            }
        }
        let root_candidate = collapse_dots(normalized);
        self.pages_by_name.get(&root_candidate)
    }

    pub fn page(&self, uri: &Url) -> Option<&IndexedPage> {
        self.pages_by_uri.get(uri)
    }

    pub fn page_by_name(&self, name: &str) -> Option<&IndexedPage> {
        self.pages_by_name
            .get(name)
            .and_then(|u| self.pages_by_uri.get(u))
    }

    /// Resolve a wiki link target to the page it points at. Source-relative
    /// first, then root-relative — mirrors [`Self::resolve_wiki_path`] but
    /// returns the `IndexedPage` so anchor checks have everything they need.
    pub fn page_for_wiki_target(
        &self,
        path: &str,
        source_page: &str,
        is_absolute: bool,
    ) -> Option<&IndexedPage> {
        self.resolve_wiki_path(path, source_page, is_absolute)
            .and_then(|u| self.pages_by_uri.get(u))
    }

    pub fn page_names(&self) -> Vec<String> {
        let mut v: Vec<String> = self.pages_by_name.keys().cloned().collect();
        v.sort();
        v
    }

    pub fn backlinks_for(&self, page_name: &str) -> &[Backlink] {
        self.backlinks
            .get(page_name)
            .map(Vec::as_slice)
            .unwrap_or(&[])
    }
}

/// Return a parsed [`DiaryPeriod`] when `uri` is a diary entry at any of
/// the four supported frequencies. Recognises:
///
///   `YYYY-MM-DD`    → Day
///   `YYYY-Www`      → Week (ISO)
///   `YYYY-MM`       → Month
///   `YYYY`          → Year
///
/// The path must sit under `<root>/<diary_rel_path>/` to be accepted;
/// without a `root` we accept any stem (so ad-hoc test setups still work).
pub fn diary_period_for_uri(
    uri: &Url,
    root: Option<&Path>,
    diary_rel_path: Option<&str>,
) -> Option<nuwiki_core::date::DiaryPeriod> {
    let path = uri.to_file_path().ok()?;
    let stem = path.file_stem()?.to_str()?;
    let parsed = nuwiki_core::date::DiaryPeriod::parse(stem)?;
    let Some(root) = root else {
        return Some(parsed);
    };
    let parent = path.parent()?;
    let expected = root.join(diary_rel_path.unwrap_or("diary"));
    if parent.starts_with(&expected) {
        Some(parsed)
    } else {
        None
    }
}

/// Derive a page name from a `file://` URL. With a workspace root the name
/// is the URL's path relative to root, sans `.wiki`. Without a root it's
/// just the filename stem.
pub fn page_name_from_uri(uri: &Url, root: Option<&Path>) -> String {
    if let (Some(root), Ok(p)) = (root, uri.to_file_path()) {
        if let Ok(rel) = p.strip_prefix(root) {
            let stripped = rel.with_extension("");
            return stripped
                .to_string_lossy()
                .replace(std::path::MAIN_SEPARATOR, "/");
        }
    }
    // Fallback: last path segment, strip .wiki if present.
    if let Ok(p) = uri.to_file_path() {
        return p
            .file_stem()
            .map(|s| s.to_string_lossy().into_owned())
            .unwrap_or_default();
    }
    uri.path_segments()
        .and_then(|mut s| s.next_back())
        .map(|s| {
            s.strip_suffix(".wiki")
                .unwrap_or(s)
                .trim_start_matches('/')
                .to_string()
        })
        .unwrap_or_default()
}

/// Normalise heading text into a URL-safe anchor slug.
pub fn slugify(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    let mut prev_dash = false;
    for ch in s.chars() {
        if ch.is_alphanumeric() {
            for c in ch.to_lowercase() {
                out.push(c);
            }
            prev_dash = false;
        } else if (ch.is_whitespace() || ch == '-' || ch == '_') && !prev_dash && !out.is_empty() {
            out.push('-');
            prev_dash = true;
        }
        // other punctuation: drop
    }
    while out.ends_with('-') {
        out.pop();
    }
    out
}

// ===== Indexing helpers =====

fn index_blocks(blocks: &[BlockNode], page: &mut IndexedPage) {
    for block in blocks {
        index_block(block, page);
    }
}

fn index_block(block: &BlockNode, page: &mut IndexedPage) {
    match block {
        BlockNode::Heading(h) => {
            let title = inline_to_text(&h.children);
            page.headings.push(HeadingInfo {
                title: title.clone(),
                level: h.level,
                anchor: slugify(&title),
                span: h.span,
            });
            index_inlines(&h.children, page);
        }
        BlockNode::Paragraph(p) => index_inlines(&p.children, page),
        BlockNode::Blockquote(BlockquoteNode { children, .. }) => {
            for c in children {
                index_block(c, page);
            }
        }
        BlockNode::List(ListNode { items, .. }) => {
            for item in items {
                index_inlines(&item.children, page);
                if let Some(sub) = &item.sublist {
                    index_block(&BlockNode::List(sub.clone()), page);
                }
            }
        }
        BlockNode::DefinitionList(dl) => {
            for item in &dl.items {
                if let Some(term) = &item.term {
                    index_inlines(term, page);
                }
                for def in &item.definitions {
                    index_inlines(def, page);
                }
            }
        }
        BlockNode::Table(TableNode { rows, .. }) => {
            for row in rows {
                for cell in &row.cells {
                    index_inlines(&cell.children, page);
                }
            }
        }
        BlockNode::Tag(t) => {
            for name in &t.tags {
                page.tags.push(TagInfo {
                    name: name.clone(),
                    scope: t.scope,
                    span: t.span,
                });
            }
        }
        _ => {}
    }
}

fn index_inlines(inlines: &[InlineNode], page: &mut IndexedPage) {
    for n in inlines {
        match n {
            InlineNode::WikiLink(w) => {
                page.outgoing.push(OutgoingLink {
                    target_page: w.target.path.clone(),
                    anchor: w.target.anchor.clone(),
                    kind: w.target.kind,
                    span: w.span,
                    is_absolute: w.target.is_absolute,
                });
                if let Some(desc) = &w.description {
                    index_inlines(desc, page);
                }
            }
            InlineNode::ExternalLink(e) => {
                if let Some(desc) = &e.description {
                    index_inlines(desc, page);
                }
            }
            InlineNode::Bold(b) => index_inlines(&b.children, page),
            InlineNode::Italic(i) => index_inlines(&i.children, page),
            InlineNode::BoldItalic(bi) => index_inlines(&bi.children, page),
            InlineNode::Strikethrough(s) => index_inlines(&s.children, page),
            InlineNode::Superscript(s) => index_inlines(&s.children, page),
            InlineNode::Subscript(s) => index_inlines(&s.children, page),
            InlineNode::Color(c) => index_inlines(&c.children, page),
            _ => {}
        }
    }
}

fn inline_to_text(inlines: &[InlineNode]) -> String {
    let mut out = String::new();
    inline_to_text_into(inlines, &mut out);
    out.trim().to_string()
}

fn inline_to_text_into(inlines: &[InlineNode], out: &mut String) {
    for n in inlines {
        match n {
            InlineNode::Text(t) => out.push_str(&t.content),
            InlineNode::Bold(b) => inline_to_text_into(&b.children, out),
            InlineNode::Italic(i) => inline_to_text_into(&i.children, out),
            InlineNode::BoldItalic(bi) => inline_to_text_into(&bi.children, out),
            InlineNode::Strikethrough(s) => inline_to_text_into(&s.children, out),
            InlineNode::Code(c) => out.push_str(&c.content),
            InlineNode::Superscript(s) => inline_to_text_into(&s.children, out),
            InlineNode::Subscript(s) => inline_to_text_into(&s.children, out),
            InlineNode::MathInline(m) => out.push_str(&m.content),
            InlineNode::Color(c) => inline_to_text_into(&c.children, out),
            InlineNode::WikiLink(w) => {
                if let Some(d) = &w.description {
                    inline_to_text_into(d, out);
                } else if let Some(p) = &w.target.path {
                    out.push_str(p);
                }
            }
            InlineNode::ExternalLink(e) => match &e.description {
                Some(d) => inline_to_text_into(d, out),
                None => out.push_str(&e.url),
            },
            InlineNode::Keyword(_) => {}
            InlineNode::Transclusion(t) => out.push_str(t.alt.as_deref().unwrap_or(&t.url)),
            InlineNode::RawUrl(r) => out.push_str(&r.url),
            InlineNode::SoftBreak(_) => out.push(' '),
        }
    }
}

// ===== Filesystem walking =====

/// Canonicalise a wikilink target into the page name the workspace index
/// would key it by. Mirrors [`WorkspaceIndex::resolve_wiki_path`]'s preference
/// (source-relative when not explicitly absolute, root-relative otherwise) but
/// without requiring the target to already be indexed — useful for backlinks
/// where the target may be added to the index after the source.
///
/// Collapses `.` and `..` segments so `posts/../foo` → `foo`. `..` at the
/// root of the wiki is treated as a no-op (the wiki has no parent in this
/// model — interwiki crossings use the dedicated `[[wiki:Page]]` syntax).
pub fn canonical_link_name(
    path: &str,
    source_page: &str,
    is_absolute: bool,
    file_extension: Option<&str>,
) -> String {
    let normalized = strip_wiki_extension(path, file_extension);
    let joined = if !is_absolute {
        if let Some(slash) = source_page.rfind('/') {
            format!("{}/{}", &source_page[..slash], normalized)
        } else {
            normalized.to_string()
        }
    } else {
        normalized.to_string()
    };
    collapse_dots(&joined)
}

/// Collapse `.` and `..` segments in a `/`-separated path string. Doesn't
/// touch URI escapes — these are workspace-relative page names, not URLs.
pub fn collapse_dots(s: &str) -> String {
    let mut parts: Vec<&str> = Vec::new();
    for seg in s.split('/') {
        match seg {
            "" | "." => continue,
            ".." => {
                parts.pop();
            }
            _ => parts.push(seg),
        }
    }
    parts.join("/")
}

/// Strip a wiki file extension from a link target path without allocating.
///
/// `file_extension` may be `".wiki"` (with leading dot) or `"wiki"` (without).
/// Returns `path` unchanged when no extension is configured, the extension is
/// empty, or the path doesn't end with `.{ext}`.
pub fn strip_wiki_extension<'a>(path: &'a str, file_extension: Option<&str>) -> &'a str {
    let ext = match file_extension {
        Some(e) => e.trim_start_matches('.'),
        None => return path,
    };
    if ext.is_empty() {
        return path;
    }
    let dot_ext_len = ext.len() + 1; // +1 for the '.'
    if path.len() > dot_ext_len
        && path.ends_with(ext)
        && path.as_bytes()[path.len() - dot_ext_len] == b'.'
    {
        &path[..path.len() - dot_ext_len]
    } else {
        path
    }
}

/// Recursively collect every `.wiki` file under `root`, skipping dotfiles
/// and dot-directories.
pub fn walk_wiki_files(root: &Path) -> Vec<PathBuf> {
    let mut out = Vec::new();
    let mut stack = vec![root.to_path_buf()];
    while let Some(dir) = stack.pop() {
        let entries = match std::fs::read_dir(&dir) {
            Ok(e) => e,
            Err(_) => continue,
        };
        for entry in entries.flatten() {
            let path = entry.path();
            let Some(name) = path.file_name() else {
                continue;
            };
            if name.to_string_lossy().starts_with('.') {
                continue;
            }
            let Ok(meta) = entry.metadata() else {
                continue;
            };
            if meta.is_dir() {
                stack.push(path);
            } else if meta.is_file() && path.extension().and_then(|e| e.to_str()) == Some("wiki") {
                out.push(path);
            }
        }
    }
    out.sort();
    out
}