chatkit/frontend/src/App.tsx

import { useState, useEffect, useCallback } from "react";
import { Header } from "./components/Header";
import { InviteGate } from "./components/InviteGate";
import { LoginForm } from "./components/LoginForm";
import { StepIndicator } from "./components/StepIndicator";
import { FileUploadPanel } from "./components/step1/FileUploadPanel";
import { ParsedDataSummary } from "./components/step2/ParsedDataSummary";
import { StudentSelector } from "./components/step2/StudentSelector";
import { PromptEditor } from "./components/step2/PromptEditor";
import { ReportViewer } from "./components/step3/ReportViewer";
import { AdminDashboard } from "./components/admin/AdminDashboard";
import { TeacherDashboard } from "./components/TeacherDashboard";
import { StudentQuizPage } from "./components/StudentQuizPage";
import { StudentLoginPage } from "./components/StudentLoginPage";
import { StudentDashboardPage } from "./components/StudentDashboardPage";
import { apiGet, apiPost, getToken, getRole, clearToken } from "./lib/api";
import type { User } from "./types";

interface StudentSelf {
  id: number;
  name: string;
  username: string;
}

const INVITE_CODE = "astraea";

interface StudentInfo {
  index: number;
  name: string;
  id: string;
}

export interface StudentReport {
  studentName: string;
  studentIndex: number;
  html: string;
  status: "pending" | "generating" | "done" | "error";
  error?: string;
}

// Student quiz route — checked before any of App's hooks run at all
// (separate component), so the teacher-only auth/invite flow is never
// touched for this path.
const QUIZ_PATH_RE = /^\/quiz\/([^/]+)\/?$/;

export default function App() {
  const path = window.location.pathname;
  const quizMatch = path.match(QUIZ_PATH_RE);
  if (quizMatch) {
    return <StudentQuizGate shareToken={quizMatch[1]} />;
  }
  return <RootApp />;
}

// A shared quiz link is student-only and requires the exact student it was
// assigned to be logged in first, so the result actually reflects who took
// it — a bare share_token used to be the sole credential, but now it's
// combined with the student's own login (see the backend's
// _check_quiz_owner). Any existing teacher session is ignored here since
// this route always needs a student identity.
function StudentQuizGate({ shareToken }: { shareToken: string }) {
  const [student, setStudent] = useState<StudentSelf | null>(null);
  const [checked, setChecked] = useState(false);

  useEffect(() => {
    if (getToken() && getRole() === "student") {
      apiGet<StudentSelf>("/api/student/me")
        .then(setStudent)
        .catch(() => clearToken())
        .finally(() => setChecked(true));
    } else {
      setChecked(true);
    }
  }, []);

  if (!checked) {
    return (
      <div className="min-h-screen flex items-center justify-center bg-[var(--color-background)]">
        <div className="w-8 h-8 border-2 border-[var(--color-primary)] border-t-transparent rounded-full animate-spin" />
      </div>
    );
  }

  if (!student) {
    return (
      <StudentLoginPage
        onLogin={setStudent}
        onSwitchToTeacher={() => { window.location.href = "/"; }}
      />
    );
  }

  return <StudentQuizPage shareToken={shareToken} />;
}

// Student and teacher logins share the root page via a toggle, rather than
// separate paths — so a logged-out visitor's mode is just local UI state,
// while a logged-in one is pinned to whichever role their stored token
// belongs to (checked before their token is verified, so we mount the right
// "am I logged in" check on the first render instead of guessing wrong and
// bouncing them to a login screen).
function RootApp() {
  const [mode, setMode] = useState<"student" | "teacher">(
    () => (getToken() ? getRole() ?? "teacher" : "teacher")
  );

  if (mode === "teacher") {
    return <TeacherApp onSwitchToStudent={() => setMode("student")} />;
  }
  return <StudentApp onSwitchToTeacher={() => setMode("teacher")} />;
}

function StudentApp({ onSwitchToTeacher }: { onSwitchToTeacher: () => void }) {
  const [student, setStudent] = useState<StudentSelf | null>(null);
  const [loading, setLoading] = useState(true);

  useEffect(() => {
    const token = getToken();
    if (token) {
      apiGet<StudentSelf>("/api/student/me")
        .then(setStudent)
        .catch(() => clearToken())
        .finally(() => setLoading(false));
    } else {
      setLoading(false);
    }
  }, []);

  if (loading) {
    return (
      <div className="min-h-screen flex items-center justify-center bg-[var(--color-background)]">
        <div className="w-8 h-8 border-2 border-[var(--color-primary)] border-t-transparent rounded-full animate-spin" />
      </div>
    );
  }

  if (!student) {
    return <StudentLoginPage onLogin={setStudent} onSwitchToTeacher={onSwitchToTeacher} />;
  }

  return <StudentDashboardPage student={student} onLogout={() => setStudent(null)} />;
}

function TeacherApp({ onSwitchToStudent }: { onSwitchToStudent: () => void }) {
  const [invited, setInvited] = useState(() => localStorage.getItem("invite_code") === INVITE_CODE);
  const [user, setUser] = useState<User | null>(null);
  const [loading, setLoading] = useState(true);
  const [currentStep, setCurrentStep] = useState<1 | 2 | 3>(1);
  const [maxStep, setMaxStep] = useState<1 | 2 | 3>(1);
  const [sessionId, setSessionId] = useState<number | null>(null);
  const [parsedData, setParsedData] = useState<Record<string, unknown>>({});
  const [students, setStudents] = useState<StudentInfo[]>([]);
  const [selectedIndices, setSelectedIndices] = useState<number[]>([]);
  const [reports, setReports] = useState<StudentReport[]>([]);
  const [isGenerating, setIsGenerating] = useState(false);
  // Tracks whether the current `reports` batch has been persisted via
  // save-to-db — reset whenever a new generation run starts, flipped once
  // the save succeeds. Lets confirmLeaveStep3 warn about completed-but-
  // unsaved reports, not just an in-progress generation.
  const [reportsSaved, setReportsSaved] = useState(false);
  const [showAdmin, setShowAdmin] = useState(false);
  const [showDashboard, setShowDashboard] = useState(true);

  // TeacherDashboard/AdminDashboard own a lot of internal fetch state (student
  // lists, sub-tab data, etc). Toggling showDashboard/showAdmin used to
  // early-return a completely different JSX tree, which unmounted whichever
  // view you left — so switching back to it always refetched from scratch.
  // Keep each view mounted (just hidden) once it's been opened.
  const currentView: "main" | "dashboard" | "admin" = showDashboard
    ? "dashboard"
    : showAdmin
    ? "admin"
    : "main";
  const [visitedViews, setVisitedViews] = useState<Set<"dashboard" | "admin">>(
    () => new Set(showDashboard ? (["dashboard"] as const) : [])
  );
  useEffect(() => {
    if (currentView === "dashboard" || currentView === "admin") {
      setVisitedViews((prev) => (prev.has(currentView) ? prev : new Set(prev).add(currentView)));
    }
  }, [currentView]);

  // Check for saved JWT on mount
  useEffect(() => {
    const token = getToken();
    if (token) {
      apiGet<User>("/api/auth/me")
        .then(setUser)
        .catch(() => clearToken())
        .finally(() => setLoading(false));
    } else {
      setLoading(false);
    }
  }, []);

  // Create session when user is authenticated and no session exists
  useEffect(() => {
    if (user && !sessionId) {
      apiPost<{ id: number }>("/api/sessions", {})
        .then((session) => setSessionId(session.id))
        .catch(() => {});
    }
    // Intentionally omits sessionId — this effect is what sets sessionId, so
    // including it would risk re-running as soon as it's created.
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [user]);

  // Extract student list from parsed data
  useEffect(() => {
    const sa = parsedData.student_answers as { students?: { name?: string; id?: string }[] } | undefined;
    if (sa?.students) {
      setStudents(
        sa.students.map((s, i) => ({
          index: i,
          name: s.name || `Student ${i + 1}`,
          id: s.id || "",
        }))
      );
    }
  }, [parsedData]);

  const handleLogout = () => {
    setUser(null);
    setSessionId(null);
    setParsedData({});
    setStudents([]);
    setSelectedIndices([]);
    setReports([]);
    setCurrentStep(1);
    setMaxStep(1);
  };

  const handleParsedDataUpdate = useCallback((dataType: string, data: unknown) => {
    setParsedData((prev) => ({ ...prev, [dataType]: data }));
  }, []);

  // If the current session was deleted server-side (e.g. a bulk wipe from the
  // dashboard) mid-upload, create a fresh one so the user can just retry
  // instead of hitting a raw "Session not found" error.
  const handleSessionExpired = useCallback(async () => {
    const session = await apiPost<{ id: number }>("/api/sessions", {});
    setSessionId(session.id);
    setParsedData({});
    setStudents([]);
    return session.id;
  }, []);

  const handleGoToStep2 = useCallback(() => {
    setCurrentStep(2);
    setMaxStep((prev) => (prev < 2 ? 2 : prev));
  }, []);

  // Sequential per-student report generation
  const handleGenerateReports = useCallback(async (model: string) => {
    if (!sessionId || selectedIndices.length === 0) return;

    const initial: StudentReport[] = selectedIndices.map((idx) => {
      const s = students.find((s) => s.index === idx);
      return {
        studentName: s?.name || `Student ${idx + 1}`,
        studentIndex: idx,
        html: "",
        status: "pending",
      };
    });
    setReports(initial);
    setIsGenerating(true);
    setReportsSaved(false);
    setCurrentStep(3);
    setMaxStep(3);

    for (let i = 0; i < selectedIndices.length; i++) {
      const idx = selectedIndices[i];

      setReports((prev) =>
        prev.map((r, j) => (j === i ? { ...r, status: "generating" } : r))
      );

      try {
        const res = await apiPost<{
          report_id: number;
          html_content: string;
          student_name: string;
        }>(`/api/sessions/${sessionId}/generate-student-report`, {
          student_index: idx,
          model,
        });

        setReports((prev) =>
          prev.map((r, j) =>
            j === i ? { ...r, html: res.html_content, status: "done" } : r
          )
        );
      } catch (err) {
        const msg = err instanceof Error ? err.message : "Generation failed";
        setReports((prev) =>
          prev.map((r, j) =>
            j === i ? { ...r, status: "error", error: msg } : r
          )
        );
      }
    }

    setIsGenerating(false);
  }, [sessionId, selectedIndices, students]);

  const handleExportAllHtml = useCallback(() => {
    const doneReports = reports.filter((r) => r.status === "done");
    if (doneReports.length === 0) return;

    const combined = `<!DOCTYPE html>
<html lang="zh-TW"><head><meta charset="UTF-8"><title>ClassLens - Student Reports</title>
<style>
  .page-break { page-break-after: always; break-after: page; }
  .report-container { margin-bottom: 2rem; }
</style></head><body>
${doneReports
  .map(
    (r, i) =>
      `<div class="report-container${i < doneReports.length - 1 ? " page-break" : ""}">${r.html}</div>`
  )
  .join("\n")}
</body></html>`;

    const blob = new Blob([combined], { type: "text/html" });
    const url = URL.createObjectURL(blob);
    const a = document.createElement("a");
    a.href = url;
    a.download = "classlens-all-reports.html";
    a.click();
    URL.revokeObjectURL(url);
  }, [reports]);

  // Warn before leaving step 3 while reports are either still generating, or
  // done but not yet saved to the database — either way, leaving silently
  // abandons work: mid-generation reports are lost outright, and completed-
  // but-unsaved reports never make it into the database or teacher dashboard.
  const confirmLeaveStep3 = (): boolean => {
    if (currentStep !== 3) return true;
    if (isGenerating) {
      return window.confirm(
        "報告尚未生成完成，離開此頁面將會中斷生成，尚未完成的報告會遺失且不會進入資料庫。\n\n確定要離開嗎？"
      );
    }
    const hasDoneReports = reports.some((r) => r.status === "done");
    if (hasDoneReports && !reportsSaved) {
      return window.confirm(
        "報告尚未儲存至資料庫，離開此頁面將會遺失這些報告。\n\n確定要離開嗎？"
      );
    }
    return true;
  };

  // Warn on tab close/refresh for the same reasons.
  useEffect(() => {
    const handler = (e: BeforeUnloadEvent) => {
      if (currentStep !== 3) return;
      const hasDoneReports = reports.some((r) => r.status === "done");
      if (isGenerating || (hasDoneReports && !reportsSaved)) {
        e.preventDefault();
        e.returnValue = "";
      }
    };
    window.addEventListener("beforeunload", handler);
    return () => window.removeEventListener("beforeunload", handler);
  }, [currentStep, isGenerating, reports, reportsSaved]);

  const handleStepClick = (step: 1 | 2 | 3) => {
    if (step > maxStep) return;
    if (currentStep === 3 && step !== 3) {
      if (!confirmLeaveStep3()) return;
      if (step === 1) {
        const ok = window.confirm(
          "回到步驟一將會清除目前的報告資料。\n\n未儲存的報告不會進入資料庫或教師儀表板。\n\n確定要繼續嗎？"
        );
        if (!ok) return;
        setReports([]);
        setMaxStep(1);
      }
    }
    setCurrentStep(step);
  };

  if (loading) {
    return (
      <div className="min-h-screen flex items-center justify-center bg-[var(--color-background)]">
        <div className="fixed top-0 left-0 right-0 h-[3px]" style={{ background: "var(--color-primary)" }} />
        <div className="flex flex-col items-center gap-6">
          <img src="/plane.svg" alt="" width="220" height="184" className="plane-float" />
          <p className="text-sm text-[var(--color-text-muted)] tracking-widest" style={{ fontFamily: "Georgia, serif", letterSpacing: "0.1em" }}>
            ClassLens
          </p>
        </div>
      </div>
    );
  }

  if (!invited) {
    return (
      <InviteGate
        onSuccess={() => {
          localStorage.setItem("invite_code", INVITE_CODE);
          setInvited(true);
        }}
        correctCode={INVITE_CODE}
        onSwitchToStudent={onSwitchToStudent}
      />
    );
  }

  if (!user) {
    return <LoginForm onLogin={setUser} onSwitchToStudent={onSwitchToStudent} />;
  }

  const headerProps = {
    isAdmin: user.is_admin,
    showAdmin,
    onToggleAdmin: () => {
      if (!confirmLeaveStep3()) return;
      setShowAdmin((v: boolean) => !v);
      setShowDashboard(false);
    },
    showDashboard,
    onToggleDashboard: () => {
      if (!confirmLeaveStep3()) return;
      setShowDashboard((v: boolean) => !v);
      setShowAdmin(false);
    },
    onLogout: () => {
      if (!confirmLeaveStep3()) return;
      handleLogout();
    },
  };

  return (
    <div className="min-h-screen">
      <Header {...headerProps} />

      {visitedViews.has("dashboard") && (
        <div className="pt-20" style={{ display: currentView === "dashboard" ? "block" : "none" }}>
          <TeacherDashboard active={currentView === "dashboard"} />
        </div>
      )}

      {visitedViews.has("admin") && (
        <div className="pt-20" style={{ display: currentView === "admin" ? "block" : "none" }}>
          <AdminDashboard isAdmin={!!user.is_admin} userEmail={user.email} userName={user.display_name || user.full_name || ""} />
        </div>
      )}

      <div style={{ display: currentView === "main" ? "block" : "none" }}>
      <div className="pt-20">
        <StepIndicator currentStep={currentStep} maxStep={maxStep} onStepClick={handleStepClick} />

        <main className="max-w-6xl mx-auto px-4 pb-12">
          {currentStep === 1 && (
            <div className="space-y-6 animate-fade-in-up">
              <div className="text-center mb-8">
                <h2 className="font-display text-2xl font-bold text-[var(--color-text)]">
                  上傳考試資料
                </h2>
              </div>
              {sessionId && (
                <FileUploadPanel
                  sessionId={sessionId}
                  parsedData={parsedData}
                  onParsedDataUpdate={handleParsedDataUpdate}
                  onGoToStep2={handleGoToStep2}
                  onSessionExpired={handleSessionExpired}
                />
              )}
            </div>
          )}

          {currentStep === 2 && (
            <div className="space-y-6 animate-fade-in-up">
              <div className="text-center mb-4">
                <h2 className="font-display text-2xl font-bold text-[var(--color-text)]">
                  選擇學生 & 編輯提示詞
                </h2>
                <p className="text-[var(--color-text-muted)] mt-2">
                  勾選要生成報告的學生，編輯提示詞，預覽資料後生成報告
                </p>
              </div>

              <ParsedDataSummary parsedData={parsedData} />

              {students.length > 0 ? (
                <>
                  <StudentSelector
                    students={students}
                    onSelectionChange={setSelectedIndices}
                  />
                  {sessionId && (
                    <PromptEditor
                      sessionId={sessionId}
                      onGenerate={(model) => { void handleGenerateReports(model); }}
                      isGenerating={isGenerating}
                      selectedCount={selectedIndices.length}
                      selectedIndices={selectedIndices}
                      students={students}
                    />
                  )}
                </>
              ) : (
                <div className="card p-8 text-center">
                  <p className="text-[var(--color-text-muted)]">
                    尚未上傳學生答案資料，請先回到上一步上傳檔案
                  </p>
                </div>
              )}
            </div>
          )}

          {currentStep === 3 && (
            <div className="animate-fade-in-up">
              <div className="text-center mb-4">
                <h2 className="font-display text-2xl font-bold text-[var(--color-text)]">
                  學生分析報告
                </h2>
              </div>
              <ReportViewer
                reports={reports}
                isGenerating={isGenerating}
                onBack={() => handleStepClick(2)}
                onExportAllHtml={handleExportAllHtml}
                onSaveToDb={
                  sessionId
                    ? async () => {
                        const res = await apiPost<{ students_saved: number }>(
                          `/api/sessions/${sessionId}/save-to-db`
                        );
                        setReportsSaved(true);
                        return res;
                      }
                    : undefined
                }
              />
            </div>
          )}
        </main>

        <footer className="py-8 text-center text-sm text-[var(--color-text-muted)]">
          <p>&copy; 2026 ClassLens &bull; AI-Powered Teaching Assistant</p>
        </footer>
      </div>
      </div>
    </div>
  );
}