freezmo/components/homework/HomeworkScreen.tsx

"use client";

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

import { AiDisclaimer } from "@/components/AiDisclaimer";
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 { wizardBusy } from "@/lib/wizard-busy";

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",
  });
  /* Mode "check" has its own intake: photos/scans of the attempt itself. */
  const solutionIntake = useNotesIntake({
    tooManyPhotosKey: "homework.solution.tooManyPhotos",
    tooManyDocsKey: "homework.solution.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]);

  /* A loading run locks the mode switch (see WizardModeSwitch). */
  const viewKind = view.kind;
  useEffect(() => {
    wizardBusy.set(viewKind === "loading");
  }, [viewKind]);
  useEffect(() => () => wizardBusy.set(false), []);

  /* Leaving the page mid-run stops the request instead of letting it finish
     invisibly and save a result the user never sees. */
  useEffect(() => () => abortRef.current?.abort(), []);

  /** 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);

    /* Mode "check": the attempt can be typed and/or attached. Documents are
       already text in the browser, so their text joins what was typed; photos
       go to the server to be transcribed (see /api/homework). */
    const solutionReady = solutionIntake.documents.filter((entry) => entry.status === "ready" && entry.extracted);
    const solutionDocumentText = documentsToText(solutionReady.map((entry) => entry.extracted!));
    const mergedSolution = [solution.trim(), solutionDocumentText.trim()].filter(Boolean).join("\n\n");
    const solutionImages = [
      ...solutionIntake.images.map((image) => image.dataUrl),
      ...solutionReady.flatMap((entry) => entry.extracted!.images).map((image) => image.dataUrl),
    ].slice(0, MAX_IMAGES_TO_SEND);

    return {
      mode: modeValue,
      images,
      text: mergedText,
      solution: modeValue === "check" ? mergedSolution : undefined,
      solutionImages: modeValue === "check" && solutionImages.length > 0 ? solutionImages : undefined,
    };
  };

  const run = useCallback(
    async (input: {
      mode: HomeworkMode;
      images: string[];
      text: string;
      solution?: string;
      solutionImages?: 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, meta } = await generateHomework(
          {
            mode: input.mode,
            images: input.images.length > 0 ? input.images : undefined,
            text: input.text || undefined,
            solution: input.solution,
            solutionImages: input.solutionImages?.length ? input.solutionImages : undefined,
          },
          { 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,
                  /* What was actually checked: typed text plus the server's
                     transcription of any attached photos, so "Redo" can
                     re-check the same work (the photos are never stored). */
                  attempt: (typeof meta.solutionText === "string" && meta.solutionText.trim()) || 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;
    const payload = buildPayload(target);
    /* Check mode needs an attempt — typed, attached, or both. */
    if (target === "check" && !payload.solution?.trim() && !payload.solutionImages?.length) {
      setError({ message: t("homework.solution.required") });
      return;
    }
    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();
    solutionIntake.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") {
    const recordKey = `${view.record.id}-${view.record.updatedAt}`;
    return (
      <div className="space-y-5">
        {view.record.mode === "make" ? (
          <HomeworkResult
            key={recordKey}
            record={view.record}
            onAgain={() => again(view.record)}
            onNew={newAssignment}
            onToDeck={() => turnIntoDeck(view.record)}
            againLoading={againLoading}
          />
        ) : view.record.mode === "check" ? (
          <HomeworkCheckResult
            key={recordKey}
            record={view.record}
            onAgain={() => again(view.record)}
            onNew={newAssignment}
            onToDeck={() => turnIntoDeck(view.record)}
            againLoading={againLoading}
          />
        ) : (
          <HomeworkGuide
            key={recordKey}
            record={view.record}
            onClose={() => setView({ kind: "start" })}
            onAgain={() => again(view.record)}
            onNew={newAssignment}
            onToDeck={() => turnIntoDeck(view.record)}
            againLoading={againLoading}
          />
        )}
        {/* The guide is a full-screen overlay and carries its own copy; this
            page-flow one would ghost through the transparent overlay. */}
        {view.record.mode === "guide" ? null : <AiDisclaimer />}
      </div>
    );
  }

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