freezmo/components/homework/HomeworkScreen.tsx

"use client";

import { useCallback, useMemo, useRef, useState, useSyncExternalStore } from "react";

import { documentsToText } from "@/lib/documents";
import { ApiError, generateHomework, type HomeworkProgress } from "@/lib/client-api";
import {
  deleteHomework,
  homeworkStore,
  saveHomework,
  sortHomework,
  updateHomework,
  type NewHomeworkResult,
} from "@/lib/homework";
import { useI18n } from "@/lib/i18n";
import type { HomeworkMode, HomeworkRecord } from "@/lib/schemas";
import { useNotesIntake } from "@/lib/use-notes-intake";

import { HomeworkGuide } from "./HomeworkGuide";
import { HomeworkInput } from "./HomeworkInput";
import { HomeworkLoading } from "./HomeworkLoading";
import { HomeworkResult } from "./HomeworkResult";

/** Photos + rendered PDF pages + images embedded in documents. */
const MAX_IMAGES_TO_SEND = 12;
const MAX_TEXT_TO_SEND = 100_000;

type View = { kind: "start" } | { kind: "loading" } | { kind: "record"; record: HomeworkRecord };

/** The /homework screen: give it an assignment, pick a mode, get a result. */
export function HomeworkScreen() {
  const { t } = useI18n();
  const records = useSyncExternalStore(
    homeworkStore.subscribe,
    homeworkStore.getSnapshot,
    homeworkStore.getServerSnapshot,
  );
  const history = useMemo(() => sortHomework(records), [records]);
  const intake = useNotesIntake({
    tooManyPhotosKey: "homework.source.tooManyPhotos",
    tooManyDocsKey: "homework.source.tooManyDocs",
  });

  const [mode, setMode] = useState<HomeworkMode>("make");
  const [text, setText] = useState("");
  const [view, setView] = useState<View>({ kind: "start" });
  const [progress, setProgress] = useState<HomeworkProgress>({ stage: "vision", message: "" });
  const [error, setError] = useState<{ message: string; kind?: string } | null>(null);
  const [againLoading, setAgainLoading] = useState(false);
  const abortRef = useRef<AbortController | null>(null);

  /** Everything the API needs, built from the current input. */
  const buildPayload = () => {
    const readyDocuments = intake.documents.filter((entry) => entry.status === "ready" && entry.extracted);
    const documentImages = readyDocuments.flatMap((entry) => entry.extracted!.images);
    const documentText = documentsToText(readyDocuments.map((entry) => entry.extracted!));
    const mergedText = [text.trim(), documentText.trim()].filter(Boolean).join("\n\n").slice(0, MAX_TEXT_TO_SEND);
    const images = [
      ...intake.images.map((image) => image.dataUrl),
      ...documentImages.map((image) => image.dataUrl),
    ].slice(0, MAX_IMAGES_TO_SEND);
    return { mode, images, text: mergedText };
  };

  const run = useCallback(
    async (input: { mode: HomeworkMode; images: string[]; text: string; recordId?: string }) => {
      setError(null);
      setView({ kind: "loading" });
      const vision = input.images.length > 0;
      setProgress({
        stage: vision ? "vision" : "generation",
        message: vision
          ? t("homework.loading.reading")
          : input.mode === "make"
            ? t("homework.loading.writing")
            : t("homework.loading.planning"),
      });

      const controller = new AbortController();
      abortRef.current = controller;

      try {
        const { homework } = await generateHomework(
          {
            mode: input.mode,
            images: input.images.length > 0 ? input.images : undefined,
            text: input.text || undefined,
          },
          { onProgress: setProgress, signal: controller.signal },
        );

        // The server hands back the assignment as text (the transcription when
        // photos were used), so “Make it again” works without the files.
        const sourceText = homework.sourceText ?? input.text;
        const result: NewHomeworkResult =
          homework.mode === "make"
            ? {
                mode: "make",
                title: homework.title,
                language: homework.language,
                sourceText,
                answer: homework.answer,
              }
            : {
                mode: "guide",
                title: homework.title,
                language: homework.language,
                sourceText,
                steps: homework.steps,
              };

        const record = input.recordId ? updateHomework(input.recordId, result) : saveHomework(result);
        if (!record) throw new ApiError(t("homework.error.body"), "bad_response");
        setView({ kind: "record", record });
      } catch (generateError) {
        if (generateError instanceof ApiError && generateError.kind === "aborted") {
          setView({ kind: "start" });
          return;
        }
        setError({
          message: generateError instanceof ApiError ? generateError.message : t("homework.error.body"),
          kind: generateError instanceof ApiError ? generateError.kind : undefined,
        });
        setView({ kind: "start" });
      } finally {
        abortRef.current = null;
        setAgainLoading(false);
      }
    },
    [t],
  );

  const submit = () => {
    const payload = buildPayload();
    void run(payload);
  };

  /** Re-run from the stored assignment text (files were never stored). */
  const again = (record: HomeworkRecord) => {
    setAgainLoading(true);
    void run({ mode: record.mode, images: [], text: record.sourceText, recordId: record.id });
  };

  const newAssignment = () => {
    intake.clear();
    setText("");
    setError(null);
    setView({ kind: "start" });
  };

  if (view.kind === "loading") {
    return <HomeworkLoading progress={progress} mode={mode} onCancel={() => abortRef.current?.abort()} />;
  }

  if (view.kind === "record") {
    return view.record.mode === "make" ? (
      <HomeworkResult
        record={view.record}
        onAgain={() => again(view.record)}
        onNew={newAssignment}
        againLoading={againLoading}
      />
    ) : (
      <HomeworkGuide
        record={view.record}
        onClose={() => setView({ kind: "start" })}
        onAgain={() => again(view.record)}
        onNew={newAssignment}
        againLoading={againLoading}
      />
    );
  }

  return (
    <HomeworkInput
      mode={mode}
      onModeChange={setMode}
      text={text}
      onTextChange={setText}
      intake={intake}
      onSubmit={submit}
      error={error}
      history={history}
      onOpen={(record) => setView({ kind: "record", record })}
      onDelete={deleteHomework}
    />
  );
}