import { highlight } from "@berry/api/web.ts";
import { type Extension, type Range, StateEffect, StateField } from "@codemirror/state";
import {
  Decoration,
  type DecorationSet,
  EditorView,
  ViewPlugin,
  type ViewUpdate,
} from "@codemirror/view";

export function highlighting(path: string, render = highlight): Extension {
  const set = StateEffect.define<DecorationSet>();
  const tokens = StateField.define<DecorationSet>({
    create: () => Decoration.none,
    update(value, transaction) {
      for (const effect of transaction.effects) if (effect.is(set)) return effect.value;
      return value.map(transaction.changes);
    },
    provide: (field) => EditorView.decorations.from(field),
  });
  return [
    tokens,
    ViewPlugin.fromClass(
      class {
        timer: ReturnType<typeof setTimeout>;
        disposed = false;
        constructor(readonly view: EditorView) {
          this.timer = setTimeout(() => void this.light(), 0);
        }
        update(update: ViewUpdate) {
          if (!update.docChanged) return;
          clearTimeout(this.timer);
          this.timer = setTimeout(() => void this.light(), 150);
        }
        async light() {
          const doc = this.view.state.doc;
          const lines = await render(doc.toString(), { path }).catch(() => null);
          if (!lines || this.disposed || this.view.state.doc !== doc) return;
          const marks: Range<Decoration>[] = [];
          for (let i = 0; i < Math.min(lines.length, doc.lines); i++) {
            let at = doc.line(i + 1).from;
            const visit = (node: Node) => {
              const start = at;
              if (node.nodeType === 3) at += node.textContent!.length;
              else {
                for (const child of node.childNodes) visit(child);
                if (node instanceof Element && /^a-[a-z]+$/.test(node.localName) && start < at) {
                  marks.push(Decoration.mark({ tagName: node.localName }).range(start, at));
                }
              }
            };
            for (const node of lines[i]) visit(node);
          }
          this.view.dispatch({ effects: set.of(Decoration.set(marks, true)) });
        }
        destroy() {
          this.disposed = true;
          clearTimeout(this.timer);
        }
      },
    ),
  ];
}