salience-editor/frontend/src/routes/index.tsx

import {
  component$,
  useSignal,
  useVisibleTask$,
  useComputed$,
  $,
  noSerialize,
  useStore,
  type Signal,
  type NoSerialize,
} from "@builder.io/qwik";
import { type DocumentHead, routeLoader$ } from "@builder.io/qwik-city";
import { EditorState, Plugin, Transaction } from "prosemirror-state";
import { EditorView, Decoration, DecorationSet } from "prosemirror-view";
import { Schema, DOMParser, Node as PMNode } from "prosemirror-model";
import { schema as basicSchema } from "prosemirror-schema-basic";
import { keymap } from "prosemirror-keymap";
import { baseKeymap } from "prosemirror-commands";
import * as math from "mathjs";
import exampleDoc from "~/assets/example-doc-1.txt?raw";
import { DebugPanel } from "~/components/debug-panel/debug-panel";
import { autoSolveParameters } from "~/utils/autosolver";
import "~/global.css";

// API Configuration
const API_BASE_URL = import.meta.env.VITE_API_BASE_URL || 'http://127.0.0.1:5000';

// Types
interface SalienceData {
  source: string;
  intervals: [number, number][];
  adjacency: number[][];
}

interface SentenceDecoration {
  from: number;
  to: number;
  salience: number;
}

interface SalienceParameters {
  exponent: number;
  threshold: number;
  randomWalkLength: number;
}

type SyncState = "clean" | "dirty" | "loading";

// Create schema with marks
const salienceSchema = new Schema({
  nodes: basicSchema.spec.nodes,
  marks: basicSchema.spec.marks,
});

// Create decorations from sentence decorations
function createDecorations(
  doc: PMNode,
  decorations: SentenceDecoration[]
): DecorationSet {
  const decos = decorations.map((deco) => {
    return Decoration.inline(deco.from, deco.to, {
      class: "sentence",
      style: `--salience: ${deco.salience}`,
    });
  });
  return DecorationSet.create(doc, decos);
}

// Plugin to manage salience decorations
function saliencePlugin(
  initialDecorations: SentenceDecoration[]
): Plugin<DecorationSet> {
  return new Plugin<DecorationSet>({
    state: {
      init(_, { doc }) {
        return createDecorations(doc, initialDecorations);
      },
      apply(tr, old) {
        // Get updated decorations from transaction metadata
        const newDecorations = tr.getMeta("updateDecorations");
        if (newDecorations) {
          return createDecorations(tr.doc, newDecorations);
        }
        // Map existing decorations through the transaction
        return old.map(tr.mapping, tr.doc);
      },
    },
    props: {
      decorations(state) {
        return this.getState(state);
      },
    },
  });
}

// Translate decorations through a transaction
function translateDecorations(
  decorations: SentenceDecoration[],
  tr: Transaction
): SentenceDecoration[] {
  return decorations
    .map((deco) => {
      const from = tr.mapping.map(deco.from);
      const to = tr.mapping.map(deco.to);
      // Only keep decorations that still have valid ranges
      if (from < to) {
        return { ...deco, from, to };
      }
      return null;
    })
    .filter((deco): deco is SentenceDecoration => deco !== null);
}

// Load initial document
export const useInitialDocument = routeLoader$(() => {
  return exampleDoc;
});

export default component$(() => {
  const initialDocument = useInitialDocument();
  const editorRef = useSignal<HTMLDivElement>();
  const modelSelectRef = useSignal<HTMLSelectElement>();
  const models = useSignal<string[]>([]);
  const currentModel = useSignal("all-mpnet-base-v2");
  const syncState = useSignal<SyncState>("clean");
  const editorView = useSignal<EditorView | null>(null);
  const sentenceDecorations = useSignal<SentenceDecoration[]>([]);
  const salienceData = useSignal<SalienceData | null>(null);
  const debounceTimer = useSignal<number | null>(null);
  const pendingRequest = useSignal<AbortController | null>(null);
  const debugPanelOpen = useSignal(false);
  const adjacencyMatrix = useSignal<number[][] | null>(null);
  const autoSolveEnabled = useSignal(true);
  const parameters = useStore<SalienceParameters>({
    exponent: 3,
    threshold: 0.95,
    randomWalkLength: 5,
  });

  // Auto-solve parameters when adjacency matrix changes (if enabled)
  // eslint-disable-next-line qwik/no-use-visible-task
  useVisibleTask$(({ track }) => {
    track(() => adjacencyMatrix.value);
    track(() => autoSolveEnabled.value);

    if (!adjacencyMatrix.value || !autoSolveEnabled.value) {
      return;
    }

    const solved = autoSolveParameters(adjacencyMatrix.value, parameters);

    if (solved.exponent !== undefined) {
      parameters.exponent = solved.exponent;
    }
    if (solved.threshold !== undefined) {
      parameters.threshold = solved.threshold;
    }
  });

  // Compute salience scores when adjacency matrix or parameters change
  // eslint-disable-next-line qwik/no-use-visible-task
  useVisibleTask$(({ track }) => {
    track(() => adjacencyMatrix.value);
    track(() => parameters.exponent);
    track(() => parameters.threshold);
    track(() => parameters.randomWalkLength);

    if (!adjacencyMatrix.value || sentenceDecorations.value.length === 0) {
      return;
    }

    // Compute scores using random walk
    const initial = adjacencyMatrix.value.map(() => 1);
    const scores = math.multiply(
      initial,
      math.pow(adjacencyMatrix.value, parameters.randomWalkLength) as number[][]
    ) as number[];

    // Update salience values in existing decorations
    sentenceDecorations.value = sentenceDecorations.value.map((deco, i) => {
      const exponentialSpread = scores[i] ** parameters.exponent;
      const shiftedDown = exponentialSpread - parameters.threshold;
      const clamped = Math.max(0, Math.min(1, shiftedDown));

      return {
        ...deco,
        salience: clamped,
      };
    });
  });

  // Update editor when decorations change
  // eslint-disable-next-line qwik/no-use-visible-task
  useVisibleTask$(({ track }) => {
    track(() => sentenceDecorations.value);

    if (!editorView.value || sentenceDecorations.value.length === 0) {
      return;
    }

    const updateTr = editorView.value.state.tr.setMeta(
      "updateDecorations",
      sentenceDecorations.value
    );
    editorView.value.dispatch(updateTr);
  });

  // Load available models
  // eslint-disable-next-line qwik/no-use-visible-task
  useVisibleTask$(async () => {
    try {
      const res = await fetch(`${API_BASE_URL}/models`);
      const data = await res.json();
      models.value = data;
    } catch (err) {
      console.error("Failed to load models:", err);
    }
  });

  // eslint-disable-next-line qwik/no-use-visible-task
  useVisibleTask$(({ track }) => {
    track(() => editorRef.value);
    if (!editorRef.value) return;

    // Parse the initial document text into separate paragraphs
    // Split by single newlines and preserve blank lines as empty paragraphs
    const paragraphs = initialDocument.value
      .split('\n')
      .map(line => {
        // Create paragraph with text if line has content, otherwise empty paragraph
        const content = line.length > 0 ? [salienceSchema.text(line)] : [];
        return salienceSchema.node("paragraph", null, content);
      });

    const initialDoc = salienceSchema.node("doc", null, paragraphs);

    const state = EditorState.create({
      schema: salienceSchema,
      doc: initialDoc,
      plugins: [
        saliencePlugin([]),
        keymap(baseKeymap),
      ],
    });

    const view = new EditorView(editorRef.value, {
      state,
      dispatchTransaction: (tr) => {
        const newState = view.state.apply(tr);
        view.updateState(newState);

        // If this was a user edit (not a programmatic update)
        if (tr.docChanged && !tr.getMeta("fromApi")) {
          // Translate existing decorations to new positions
          const translated = translateDecorations(sentenceDecorations.value, tr);
          sentenceDecorations.value = translated;
          // The reactive task will update the editor with translated decorations

          // Mark as dirty and schedule API request
          syncState.value = "dirty";
          scheduleApiRequest(
            view,
            currentModel.value,
            syncState,
            sentenceDecorations,
            salienceData,
            debounceTimer,
            pendingRequest,
            adjacencyMatrix,
          );
        }
      },
    });

    editorView.value = view;

    fetchSalienceData(
      currentModel.value,
      view,
      syncState,
      sentenceDecorations,
      salienceData,
      pendingRequest,
      adjacencyMatrix
    );

    // Cleanup
    return () => {
      view.destroy();
    };
  });

  // Handle model change
  const handleModelChange = $((event: Event) => {
    const target = event.target as HTMLSelectElement;
    currentModel.value = target.value;
    if (editorView.value) {
      fetchSalienceData(
        target.value,
        editorView.value,
        syncState,
        sentenceDecorations,
        salienceData,
        pendingRequest,
        adjacencyMatrix,
      );
    }
  });

  return (
    <>
      <div class="container">
        <h1>
          Salience
          <span class="subtitle">
            sentence highlights based on their significance to the document
          </span>
          <a href="/about" class="about-link">How it works →</a>
        </h1>
        <div class="controls">
          <label for="model-select">Model:</label>
          <select
            id="model-select"
            ref={modelSelectRef}
            onChange$={handleModelChange}
          >
            {models.value.length === 0 ? (
              <option>Loading...</option>
            ) : (
              models.value.map((model) => (
                <option
                  key={model}
                  value={model}
                  selected={model === currentModel.value}
                >
                  {model}
                </option>
              ))
            )}
          </select>
          <label class="auto-solve-label">
            <input
              type="checkbox"
              checked={autoSolveEnabled.value}
              onChange$={(e) => {
                autoSolveEnabled.value = (e.target as HTMLInputElement).checked;
              }}
            />
            Auto highlight colors
          </label>
          <span class={`status-badge ${syncState.value}`}>
            {syncState.value === "clean"
              ? "Synchronized"
              : syncState.value === "dirty"
                ? "Modified"
                : "Processing..."}
          </span>
        </div>
        <div
          ref={editorRef}
          class="editor"
        />
      </div>

      <DebugPanel
        isOpen={debugPanelOpen}
        parameters={parameters}
        adjacencyMatrix={adjacencyMatrix}
      />
    </>
  );
});

// Shared function to fetch salience data from API
async function fetchSalienceData(
  model: string,
  view: EditorView,
  syncState: Signal<SyncState>,
  sentenceDecorations: Signal<SentenceDecoration[]>,
  salienceData: Signal<SalienceData | null>,
  pendingRequest: Signal<AbortController | null>,
  adjacencyMatrix: Signal<number[][] | null>,
) {
  // Cancel any pending request
  if (pendingRequest.value) {
    pendingRequest.value.abort();
  }

  const controller = new AbortController();
  pendingRequest.value = controller;
  syncState.value = "loading";

  // Get current document text
  const text = view.state.doc.textContent;

  try {
    const res = await fetch(`${API_BASE_URL}/salience?model=${encodeURIComponent(model)}`, {
      method: "POST",
      headers: { "Content-Type": "text/plain" },
      body: text,
      signal: controller.signal,
    });
    const data: SalienceData = await res.json();

    salienceData.value = data;
    adjacencyMatrix.value = data.adjacency;

    // Set decorations with positions and zero salience
    // The reactive task will compute and update the salience values
    sentenceDecorations.value = data.intervals.map(([start, end]) => ({
      from: start,
      to: end,
      salience: 0,
    }));

    syncState.value = "clean";
    pendingRequest.value = null;
  } catch (err: any) {
    if (err.name !== "AbortError") {
      console.error("Failed to load salience:", err);
      syncState.value = "dirty";
    }
    pendingRequest.value = null;
  }
}

// Schedule API request with debouncing
function scheduleApiRequest(
  view: EditorView,
  model: string,
  syncState: Signal<SyncState>,
  sentenceDecorations: Signal<SentenceDecoration[]>,
  salienceData: Signal<SalienceData | null>,
  debounceTimer: Signal<number | null>,
  pendingRequest: Signal<AbortController | null>,
  adjacencyMatrix: Signal<number[][] | null>,
) {
  // Clear existing timer
  if (debounceTimer.value !== null) {
    clearTimeout(debounceTimer.value);
  }

  // Set new timer
  debounceTimer.value = window.setTimeout(() => {
    // Check if we're still dirty (not loading)
    if (syncState.value === "dirty") {
      fetchSalienceData(
        model,
        view,
        syncState,
        sentenceDecorations,
        salienceData,
        pendingRequest,
        adjacencyMatrix
      );
    }
  }, 1000); // 1 second debounce
}

export const head: DocumentHead = {
  title: "Salience",
  meta: [
    {
      name: "description",
      content: "Sentence highlights based on their significance to the document",
    },
  ],
};