insta-toc/src/tocModel.ts

Simultaneously generate, update, and maintain a table of contents for your notes in real time.

import { TFile, htmlToMarkdown, type App } from "obsidian";
import {
    listRegex,
    localTocSettingsRegex,
    markdownLinkRegex,
    tagLinkRegex,
    wikiLinkNoAliasRegex,
    wikiLinkWithAliasRegex
} from "./constants";
import type { InstaTocSettings } from "./settings/Settings";
import { parseLocalTocSettingsYaml, resolveTocTitle } from "./settings/localTocSettings";
import type { FileKey, FoldKey, ParseLocalTocSettingsResult, TocBlockItem, TocBlockModel } from "./types";
import type UiStateManager from "./uiStateManager";

type SourceData = { sourceFilePath: FileKey; localSettings: string; };
type FoldState = { foldParentIndex: number; activeFoldKeys: Set<FoldKey>; };
type HandledLink = { contentText: string; alias: string; };
type LinkTransform = {
    pattern: RegExp;
    /** How to transform contentText (the full heading text for display). null = skip. */
    forContent: ((...args: string[]) => string) | null;
    /** How to transform alias (the cleaned link text). null = skip. */
    forAlias: ((...args: string[]) => string) | null;
};

const linkTransforms: LinkTransform[] = [ {
    pattern: wikiLinkWithAliasRegex,
    // [[wikilink|wikitext]] → "wikilink wikitext"
    forContent: (_match, refPath, refAlias) => `${refPath} ${refAlias}`,
    // [[wikilink|wikitext]] → "wikitext"
    forAlias: (_match, _refPath, refAlias) => `${refAlias}`
}, {
    pattern: wikiLinkNoAliasRegex,
    // [[path/to/note]] → "note" (last path segment — same for both)
    forContent: (_match, refPath) => String(refPath).split("/").pop() ?? String(refPath),
    forAlias: (_match, refPath) => String(refPath).split("/").pop() ?? String(refPath)
}, {
    pattern: markdownLinkRegex,
    // [Link](url) → keep full match in contentText
    forContent: (match) => match,
    // [Link](url) → keep just "Link" in alias
    forAlias: (_match, refAlias) => String(refAlias)
}, {
    pattern: tagLinkRegex,
    // #some-tag → "some-tag" (content only — alias doesn't need this)
    forContent: (_match, _symbol, tag) => String(tag),
    forAlias: null
} ];

function handleLinks(settings: InstaTocSettings, content: string): HandledLink {
    let contentText = content;
    let alias = content;

    for (const { pattern, forContent, forAlias } of linkTransforms) {
        if (forContent) contentText = contentText.replace(pattern, forContent);
        if (forAlias) alias = alias.replace(pattern, forAlias);
    }

    // Post-processing for alias: HTML → markdown + excluded chars
    alias = htmlToMarkdown(alias);
    for (const char of settings.excludedChars) {
        alias = alias.replaceAll(char, "");
    }

    return { contentText, alias };
}

export class TocModel {
    private uiStateManager: UiStateManager;
    private app: App;
    private settings: InstaTocSettings;
    private sourceData: SourceData;
    private foldStateCallback: (key: FoldKey) => boolean | undefined;
    private foldState: FoldState = { foldParentIndex: 0, activeFoldKeys: new Set<FoldKey>() };
    private _model: TocBlockModel = { title: null, items: [] };

    constructor(
        uiStateManager: UiStateManager,
        app: App,
        settings: InstaTocSettings,
        sourceData: SourceData,
        foldStateCallback: (key: FoldKey) => boolean | undefined
    ) {
        this.uiStateManager = uiStateManager;
        this.app = app;
        this.settings = settings;
        this.sourceData = sourceData;
        this.foldStateCallback = foldStateCallback;
        this.generateModel();
    }

    public get model(): TocBlockModel {
        return this._model;
    }

    private generateModel(): void {
        const fileLinkText = this.getLinkText();
        const localTocSettings = this.getLocalTocSettingsFromSource().settings;
        const resolvedTitle = resolveTocTitle(localTocSettings, this.settings);
        const lines = this.sourceData.localSettings.replace(localTocSettingsRegex, "").split("\n");
        const stack: Array<{ indent: number; item: TocBlockItem; }> = [];

        let itemIndex = 0;

        for (const line of lines) {
            const trimmedLine = line.trim();
            if (trimmedLine.length === 0) {
                continue;
            }

            if (!this._model.title) {
                this._model.title = resolvedTitle;
            }

            const match = line.match(listRegex);
            if (!match) {
                continue;
            }

            const [ , indent, _bullet, content ] = match;
            const { contentText, alias } = handleLinks(this.settings, content);
            const normalizedIndent = indent.replace(/\t/g, "    ").length;
            const item: TocBlockItem = {
                key: `${this.sourceData.sourceFilePath}:${itemIndex}`,
                text: alias || contentText,
                href: `${fileLinkText}#${contentText}`,
                children: [],
                foldKey: null,
                initialCollapsed: false
            };

            itemIndex += 1;

            while (stack.length > 0 && stack[stack.length - 1].indent >= normalizedIndent) {
                stack.pop();
            }

            const siblings = stack.length > 0 ? stack[stack.length - 1].item.children : this._model.items;

            siblings.push(item);
            stack.push({ indent: normalizedIndent, item });
        }

        this.assignInitialFoldState();
        this.uiStateManager.pruneTocFoldStateForPath(this.sourceData.sourceFilePath, {
            activeModernFoldKeys: this.foldState.activeFoldKeys
        });
    }

    public getLinkText(): string {
        const file = this.app.vault.getAbstractFileByPath(this.sourceData.sourceFilePath);
        const fallbackLinkText = this.sourceData.sourceFilePath.replace(/\.md$/u, "");
        const fileLinkText = file instanceof TFile
            ? this.app.metadataCache.fileToLinktext(file, this.sourceData.sourceFilePath, true)
            : fallbackLinkText;
        return fileLinkText;
    }

    public getLocalTocSettingsFromSource(): ParseLocalTocSettingsResult {
        const match = this.sourceData.localSettings.match(localTocSettingsRegex);

        return parseLocalTocSettingsYaml(match?.[1] ?? "");
    }

    public assignInitialFoldState(items: TocBlockItem[] = this._model.items): void {
        for (const item of items) {
            if (item.children.length === 0) continue;

            this.foldState.foldParentIndex += 1;

            const foldKey: FoldKey = `${this.sourceData.sourceFilePath}::${this.foldState.foldParentIndex}`;
            this.foldState.activeFoldKeys.add(foldKey);

            item.foldKey = foldKey;
            item.initialCollapsed = this.foldStateCallback(foldKey) ?? false;
            
            this.assignInitialFoldState(item.children);
        }
    }
}