FabledScribe/frontend/src/composables/useAssist.ts

import { ref, computed, watch, type Ref } from "vue";
import { apiPost, apiSSEStream, type SSEStreamHandle } from "@/api/client";
import { useChatStore } from "@/stores/chat";
import {
  parseMarkdownSections,
  type MarkdownSection,
} from "@/utils/sectionParser";

export type AssistState = "idle" | "streaming" | "review";

export interface AssistTarget {
  text: string;
  startOffset: number;
  endOffset: number;
}

export interface DiffLine {
  type: 'equal' | 'delete' | 'insert';
  text: string;
}

function computeDiff(a: string, b: string): DiffLine[] {
  const aLines = a.split('\n');
  const bLines = b.split('\n');
  const m = aLines.length, n = bLines.length;
  const dp: number[][] = Array.from({ length: m + 1 }, () => new Array(n + 1).fill(0));
  for (let i = m - 1; i >= 0; i--)
    for (let j = n - 1; j >= 0; j--)
      dp[i][j] = aLines[i] === bLines[j]
        ? dp[i+1][j+1] + 1
        : Math.max(dp[i+1][j], dp[i][j+1]);
  const result: DiffLine[] = [];
  let i = 0, j = 0;
  while (i < m && j < n) {
    if (aLines[i] === bLines[j]) { result.push({ type: 'equal', text: aLines[i++] }); j++; }
    else if (dp[i+1][j] >= dp[i][j+1]) result.push({ type: 'delete', text: aLines[i++] });
    else result.push({ type: 'insert', text: bLines[j++] });
  }
  while (i < m) result.push({ type: 'delete', text: aLines[i++] });
  while (j < n) result.push({ type: 'insert', text: bLines[j++] });
  return result;
}

export function useAssist(body: Ref<string>) {
  const chatStore = useChatStore();

  const state = ref<AssistState>("idle");
  const sections = ref<MarkdownSection[]>([]);
  const selectedSection = ref<MarkdownSection | null>(null);
  const customSelection = ref<{ start: number; end: number; text: string } | null>(null);
  const instruction = ref("");
  const streamingText = ref("");
  const proposedText = ref("");
  const error = ref("");
  const isProofreading = ref(false);

  // Snapshot of body at the time a section/selection was chosen
  let bodySnapshot = "";
  let streamHandle: SSEStreamHandle | null = null;

  const target = computed<AssistTarget | null>(() => {
    if (customSelection.value) {
      return {
        text: customSelection.value.text,
        startOffset: customSelection.value.start,
        endOffset: customSelection.value.end,
      };
    }
    if (selectedSection.value) {
      return {
        text: selectedSection.value.content,
        startOffset: selectedSection.value.startOffset,
        endOffset: selectedSection.value.endOffset,
      };
    }
    return null;
  });

  const canSubmit = computed(
    () =>
      target.value !== null &&
      instruction.value.trim().length > 0 &&
      chatStore.chatReady &&
      state.value !== "streaming"
  );

  const diff = computed<DiffLine[]>(() => {
    if (state.value !== 'review' || !target.value || !proposedText.value) return [];
    return computeDiff(target.value.text, proposedText.value);
  });

  function refreshSections() {
    sections.value = parseMarkdownSections(body.value);
  }

  function selectSection(section: MarkdownSection) {
    customSelection.value = null;
    selectedSection.value = section;
    bodySnapshot = body.value;
    error.value = "";
  }

  function selectTextRange(start: number, end: number) {
    if (start === end) return;
    selectedSection.value = null;
    customSelection.value = {
      start,
      end,
      text: body.value.slice(start, end),
    };
    bodySnapshot = body.value;
    error.value = "";
  }

  function clearSelection() {
    if (streamHandle) {
      streamHandle.close();
      streamHandle = null;
    }
    selectedSection.value = null;
    customSelection.value = null;
    instruction.value = "";
    streamingText.value = "";
    proposedText.value = "";
    error.value = "";
    state.value = "idle";
  }

  async function submit() {
    if (!canSubmit.value || !target.value) return;

    state.value = "streaming";
    streamingText.value = "";
    proposedText.value = "";
    error.value = "";

    try {
      // Step 1: Launch background generation
      await apiPost("/api/notes/assist", {
        body: body.value,
        target_section: target.value.text,
        instruction: instruction.value,
      });

      // Step 2: Tail the SSE buffer
      streamHandle = apiSSEStream("/api/notes/assist/stream", (evt) => {
        if (evt.event === "chunk") {
          streamingText.value += evt.data.chunk as string;
        }
        if (evt.event === "done") {
          proposedText.value = (evt.data.full_text as string) || streamingText.value;
          state.value = "review";
          streamHandle = null;
        }
        if (evt.event === "error") {
          error.value = evt.data.error as string;
          state.value = "idle";
          streamHandle = null;
        }
      });

      await streamHandle.done;

      // Fallback: if stream ended without done/error, use accumulated text
      if (state.value === "streaming") {
        if (streamingText.value) {
          proposedText.value = streamingText.value;
          state.value = "review";
        } else {
          state.value = "idle";
        }
        streamHandle = null;
      }
    } catch (e) {
      error.value = e instanceof Error ? e.message : "Request failed";
      state.value = "idle";
      streamHandle = null;
    }
  }

  async function proofread() {
    if (state.value === 'streaming') return;
    isProofreading.value = true;
    selectedSection.value = null;
    customSelection.value = { start: 0, end: body.value.length, text: body.value };
    bodySnapshot = body.value;
    error.value = '';
    instruction.value =
      "Proofread for clarity, grammar, spelling, and flow. " +
      "Preserve the author's voice and all factual content. " +
      "Return only the improved text without explanation.";
    await submit();
  }

  function accept(): string {
    if (!target.value || !proposedText.value) return body.value;

    const t = target.value;

    // Validate that the body hasn't changed at the target offsets
    if (body.value !== bodySnapshot) {
      const currentSlice = body.value.slice(t.startOffset, t.endOffset);
      if (currentSlice !== t.text) {
        error.value = "The document changed since this suggestion was made. Please clear and try again.";
        return body.value;
      }
    }

    // Ensure proposed text ends with a newline for clean separation
    let text = proposedText.value;
    if (!text.endsWith("\n")) {
      text += "\n";
    }

    const newBody =
      body.value.slice(0, t.startOffset) +
      text +
      body.value.slice(t.endOffset);

    // Reset state
    selectedSection.value = null;
    customSelection.value = null;
    streamingText.value = "";
    proposedText.value = "";
    error.value = "";
    state.value = "idle";
    isProofreading.value = false;

    return newBody;
  }

  function reject() {
    proposedText.value = "";
    streamingText.value = "";
    error.value = "";
    state.value = "idle";
    isProofreading.value = false;
    // Keep selection + instruction for retry
  }

  // Keep sections in sync with body automatically
  watch(body, () => {
    refreshSections();
  }, { immediate: true });

  return {
    state,
    sections,
    selectedSection,
    customSelection,
    target,
    instruction,
    streamingText,
    proposedText,
    error,
    canSubmit,
    diff,
    isProofreading,
    refreshSections,
    selectSection,
    selectTextRange,
    clearSelection,
    submit,
    proofread,
    accept,
    reject,
  };
}