FabledScribe/frontend/src/composables/useAssist.ts

import { ref, computed, watch, type Ref } from "vue";
import { computeDiff, type DiffLine } from "@/utils/diff";
import { apiPost, apiPut, apiDelete, apiSSEStream, type SSEStreamHandle } from "@/api/client";
import { useToastStore } from "@/stores/toast";
import {
  parseMarkdownSections,
  type MarkdownSection,
} from "@/utils/sectionParser";

export type AssistState = "idle" | "streaming" | "review";
export type ScopeMode = "document" | "section";

// Re-exported: this composable was where DiffLine lived before the diff
// moved to a shared util, and every consumer still imports the type from here.
export type { DiffLine };

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

export interface NoteDraft {
  id: number;
  note_id: number;
  proposed_body: string;
  original_body: string;
  instruction: string;
  scope: string;
  created_at: string;
  updated_at: string;
}


export function useAssist(body: Ref<string>, noteId?: Ref<number | null>, projectId?: Ref<number | null>) {
  const toast = useToastStore();

  const state = ref<AssistState>("idle");
  const scopeMode = ref<ScopeMode>("document");
  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("");
  // Full proposed document body (for section mode: body with section replaced)
  const proposedFullBody = ref("");
  const error = ref("");
  const isProofreading = ref(false);

  // Snapshot of body at the time generation was started
  let bodySnapshot = "";
  let streamHandle: SSEStreamHandle | null = null;

  const target = computed<AssistTarget | null>(() => {
    if (scopeMode.value === 'document') return 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(() =>
    instruction.value.trim().length > 0 &&
    state.value !== "streaming" &&
    (scopeMode.value === 'document' || target.value !== null)
  );

  // Diff always compares full document: original snapshot → proposed
  const diff = computed<DiffLine[]>(() => {
    if (state.value !== 'review' || !proposedFullBody.value) return [];
    return computeDiff(bodySnapshot, proposedFullBody.value);
  });

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

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

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

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

  async function _saveDraft() {
    if (!noteId?.value) return;
    try {
      await apiPut(`/api/notes/${noteId.value}/draft`, {
        proposed_body: proposedFullBody.value,
        original_body: bodySnapshot,
        instruction: instruction.value,
        scope: scopeMode.value,
      });
    } catch {
      // Non-critical — ignore draft save errors
    }
  }

  async function _deleteDraft() {
    if (!noteId?.value) return;
    try {
      await apiDelete(`/api/notes/${noteId.value}/draft`);
    } catch {
      // Ignore — draft may not exist
    }
  }

  function _buildProposedFullBody(fullText: string, wholeDoc: boolean): string {
    if (wholeDoc) return fullText;
    const t = target.value;
    if (t) {
      let text = fullText;
      if (!text.endsWith("\n")) text += "\n";
      return bodySnapshot.slice(0, t.startOffset) + text + bodySnapshot.slice(t.endOffset);
    }
    return fullText;
  }

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

    bodySnapshot = body.value;
    state.value = "streaming";
    streamingText.value = "";
    proposedText.value = "";
    proposedFullBody.value = "";
    error.value = "";

    const wholeDoc = scopeMode.value === 'document';
    const targetSection = wholeDoc ? body.value : (target.value?.text ?? "");

    try {
      await apiPost("/api/notes/assist", {
        body: body.value,
        target_section: targetSection,
        instruction: instruction.value,
        whole_doc: wholeDoc,
        ...(noteId?.value ? { note_id: noteId.value } : {}),
        ...(projectId?.value ? { project_id: projectId.value } : {}),
      });

      // Stream with automatic reconnection on dropped connections
      const MAX_RETRIES = 3;
      let lastEventId = -1;
      let gotDone = false;

      for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
        if (attempt > 0) {
          if (state.value !== "streaming") break;
          await new Promise<void>((r) => setTimeout(r, Math.min(1000 * 2 ** (attempt - 1), 5000)));
          if (state.value !== "streaming") break;
        }

        gotDone = false;
        streamHandle = apiSSEStream(
          `/api/notes/assist/stream${lastEventId >= 0 ? `?last_event_id=${lastEventId}` : ""}`,
          (evt) => {
            lastEventId = evt.id;
            if (evt.event === "chunk") {
              streamingText.value += evt.data.chunk as string;
            }
            if (evt.event === "done") {
              gotDone = true;
              const fullText = (evt.data.full_text as string) || streamingText.value;
              proposedText.value = fullText;
              proposedFullBody.value = _buildProposedFullBody(fullText, wholeDoc);
              state.value = "review";
              streamHandle = null;
              _saveDraft();
            }
            if (evt.event === "error") {
              gotDone = true;
              error.value = evt.data.error as string;
              toast.show("Assist failed: " + error.value, "error");
              state.value = "idle";
              streamHandle = null;
            }
          },
        );

        await streamHandle.done;
        streamHandle = null;
        if (gotDone) break;
      }

      // If still streaming after all retries, treat accumulated text as result
      if (state.value === "streaming") {
        if (streamingText.value) {
          proposedText.value = streamingText.value;
          proposedFullBody.value = _buildProposedFullBody(streamingText.value, wholeDoc);
          state.value = "review";
          await _saveDraft();
        } else {
          state.value = "idle";
        }
      }
    } catch (e) {
      error.value = e instanceof Error ? e.message : "Request failed";
      toast.show("Assist failed: " + error.value, "error");
      state.value = "idle";
      streamHandle = null;
    }
  }

  async function proofread() {
    if (state.value === 'streaming') return;
    isProofreading.value = true;
    scopeMode.value = 'document';
    selectedSection.value = null;
    customSelection.value = null;
    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 {
    const result = proposedFullBody.value || body.value;

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

    _deleteDraft();
    return result;
  }

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

  function loadDraft(draft: NoteDraft) {
    bodySnapshot = draft.original_body;
    proposedText.value = draft.proposed_body;
    proposedFullBody.value = draft.proposed_body;
    instruction.value = draft.instruction;
    scopeMode.value = draft.scope as ScopeMode;
    state.value = "review";
  }

  watch(body, () => {
    refreshSections();
  }, { immediate: true });

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