freezmo/components/homework/HomeworkScreen.tsx

"use client";

import { useCallback, useEffect, useMemo, useRef, useState, useSyncExternalStore } from "react";
import { useRouter } from "next/navigation";

import { documentsToText } from "@/lib/documents";
import { ApiError, generateHomework, type HomeworkProgress } from "@/lib/client-api";
import { setDeckSeed } from "@/lib/deck-seed";
import {
  assembleSolution,
  deleteHomework,
  getHomework,
  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 { clearWizardDraft } from "@/lib/wizard-draft";

import { HomeworkCheckResult } from "./HomeworkCheckResult";
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({ initialRecordId }: { initialRecordId?: string }) {
  const { t } = useI18n();
  const router = useRouter();
  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("");
  /** Mode "check": the student's own attempt at the assignment. */
  const [solution, setSolution] = 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);

  // Deep link from the library's Homework tab: open the saved result directly.
  useEffect(() => {
    if (!initialRecordId) return;
    const record = getHomework(initialRecordId);
    if (record) setView({ kind: "record", record });
  }, [initialRecordId]);

  /** Everything the API needs, built from the current input. */
  const buildPayload = (modeValue: HomeworkMode) => {
    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: modeValue,
      images,
      text: mergedText,
      solution: modeValue === "check" ? solution.trim() : undefined,
    };
  };

  const run = useCallback(
    async (input: { mode: HomeworkMode; images: string[]; text: string; solution?: 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")
            : input.mode === "check"
              ? t("homework.loading.checking")
              : 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,
            solution: input.solution,
          },
          { onProgress: setProgress, signal: controller.signal },
        );

        // The server hands back the assignment as text (the transcription when
        // photos were used), so “Redo” 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,
              }
            : homework.mode === "check"
              ? {
                  mode: "check",
                  title: homework.title,
                  language: homework.language,
                  sourceText,
                  verdict: homework.verdict,
                  correct: homework.correct,
                  wrong: homework.wrong,
                  corrected: homework.corrected,
                  attempt: input.solution,
                }
              : {
                  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],
  );

  /** `modeOverride` lets the phone flow start straight from a tapped mode card. */
  const submit = (modeOverride?: HomeworkMode) => {
    const target = modeOverride ?? mode;
    if (target === "check" && !solution.trim()) {
      setError({ message: t("homework.solution.required") });
      return;
    }
    const payload = buildPayload(target);
    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,
      solution: record.mode === "check" ? record.attempt : undefined,
      recordId: record.id,
    });
  };

  /** Hand the finished content to the New deck flow with the options pre-filled. */
  const turnIntoDeck = (record: HomeworkRecord) => {
    const content =
      record.mode === "make"
        ? record.answer
        : record.mode === "check"
          ? record.corrected
          : assembleSolution(record);
    setDeckSeed({ text: content, difficulty: "medium", count: 10 });
    router.push("/create");
  };

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

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

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

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