examjam-frontend/context/ExamContext.tsx

"use client";

import React, {
  createContext,
  useContext,
  useState,
  useEffect,
  ReactNode,
} from "react";
import { useRouter } from "next/navigation";

import { Test, TestAttempt, TestContextType, Answer } from "@/types/exam";
import { getFromStorage, removeFromStorage, setToStorage } from "@/lib/utils";
import { useAuth } from "./AuthContext";

const ExamContext = createContext<TestContextType | undefined>(undefined);

const STORAGE_KEYS = {
  CURRENT_EXAM: "current-exam",
  CURRENT_ATTEMPT: "current-attempt",
} as const;

export const ExamProvider: React.FC<{ children: ReactNode }> = ({
  children,
}) => {
  const router = useRouter();
  const [currentExam, setCurrentExamState] = useState<Test | null>(null);
  const [currentAttempt, setCurrentAttemptState] = useState<TestAttempt | null>(
    null
  );
  const [isHydrated, setIsHydrated] = useState(false);
  const [isInitialized, setIsInitialized] = useState(false);
  const { user } = useAuth();

  // Hydrate from storage
  useEffect(() => {
    const savedExam = getFromStorage<Test>(STORAGE_KEYS.CURRENT_EXAM);
    const savedAttempt = getFromStorage<TestAttempt>(
      STORAGE_KEYS.CURRENT_ATTEMPT
    );

    if (savedExam) setCurrentExamState(savedExam);
    if (savedAttempt) setCurrentAttemptState(savedAttempt);

    setIsHydrated(true);
  }, []);

  // Persist exam
  useEffect(() => {
    if (!isHydrated) return;
    if (currentExam) {
      setToStorage(STORAGE_KEYS.CURRENT_EXAM, currentExam);
    } else {
      removeFromStorage(STORAGE_KEYS.CURRENT_EXAM);
    }
  }, [currentExam, isHydrated]);

  // Persist attempt
  useEffect(() => {
    if (!isHydrated) return;
    if (currentAttempt) {
      setToStorage(STORAGE_KEYS.CURRENT_ATTEMPT, currentAttempt);
    } else {
      removeFromStorage(STORAGE_KEYS.CURRENT_ATTEMPT);
    }
  }, [currentAttempt, isHydrated]);

  const setCurrentExam = (exam: Test) => {
    setCurrentExamState(exam);
    setCurrentAttemptState(null);
  };

  const startExam = (exam?: Test): void => {
    const examToUse = exam || currentExam;
    const user_id = user?.user_id;

    if (!examToUse) {
      console.warn("No exam selected, redirecting to /unit");
      router.push("/unit");
      return;
    }

    const attempt: TestAttempt = {
      user_id: user_id, // TODO: inject from auth
      test_id: examToUse.metadata.test_id,
      subject_id: "temp-subject", // TODO: might delete
      topic_id: "temp-topic", // TODO: might delete
      test_type: "Mock", // or "Subject"/"Topic" depending on flow
      attempt_id: crypto.randomUUID(),
      start_time: new Date().toISOString(),
      end_time: "",
      user_questions: examToUse.questions,
      user_answers: Array(examToUse.questions.length).fill(null),
      correct_answers: [],
      correct_answers_count: 0,
      wrong_answers_count: 0,
      skipped_questions_count: 0,
    };

    setCurrentAttemptState(attempt);
    setIsInitialized(true);
  };

  const setAnswer = (questionId: string, answer: Answer) => {
    if (!currentAttempt) {
      console.warn("No exam attempt started, redirecting to /unit");
      router.push("/unit");
      return;
    }

    const questionIndex = currentAttempt.user_questions.findIndex(
      (q) => q.question_id === questionId
    );
    if (questionIndex === -1) return;

    setCurrentAttemptState((prev) => {
      if (!prev) return null;
      const updatedAnswers = [...prev.user_answers];
      updatedAnswers[questionIndex] = answer;
      return { ...prev, user_answers: updatedAnswers };
    });
  };

  const setApiResponse = (response: any) => {
    // If you want to store API response in attempt
    setCurrentAttemptState((prev) =>
      prev ? { ...prev, apiResponse: response } : null
    );
  };

  const submitExam = (): TestAttempt | null => {
    if (!currentAttempt) {
      console.warn("No exam attempt to submit, redirecting to /unit");
      router.push("/unit");
      return null;
    }

    const totalQuestions = currentAttempt.user_questions.length;

    // Example evaluation (assumes correct answers in `correct_answers`)
    const wrong = currentAttempt.user_answers.filter(
      (ans, idx) => ans !== null && ans !== currentAttempt.correct_answers[idx]
    ).length;

    const skipped = currentAttempt.user_answers.filter(
      (ans) => ans === null
    ).length;
    const correct = totalQuestions - wrong - skipped;

    const completedAttempt: TestAttempt = {
      ...currentAttempt,
      end_time: new Date().toISOString(),
      correct_answers_count: correct,
      wrong_answers_count: wrong,
      skipped_questions_count: skipped,
    };

    setCurrentAttemptState(completedAttempt);
    return completedAttempt;
  };

  const clearExam = (): void => {
    setCurrentExamState(null);
    setCurrentAttemptState(null);
  };

  const getAnswer = (questionId: string): Answer => {
    if (!currentAttempt) return null;
    const index = currentAttempt.user_questions.findIndex(
      (q) => q.question_id === questionId
    );
    return index >= 0 ? currentAttempt.user_answers[index] : null;
  };

  const getApiResponse = (): any => {
    return (currentAttempt as any)?.apiResponse;
  };

  const getProgress = (): number => {
    if (!currentAttempt) return 0;
    const totalQuestions = currentAttempt.user_questions.length;
    const answered = currentAttempt.user_answers.filter(
      (a) => a !== null
    ).length;
    return totalQuestions > 0 ? (answered / totalQuestions) * 100 : 0;
  };

  const isExamStarted = () => !!currentExam && !!currentAttempt;

  const isExamCompleted = (): boolean => {
    if (!isHydrated) return false;
    return !!currentAttempt?.end_time;
  };

  const contextValue: TestContextType = {
    currentExam,
    currentAttempt,
    setCurrentExam,
    startExam,
    setAnswer,
    submitExam,
    clearExam,
    setApiResponse,
    getAnswer,
    getProgress,
    isExamStarted,
    isExamCompleted,
    getApiResponse,
    isHydrated,
    isInitialized,
  };

  return (
    <ExamContext.Provider value={contextValue}>{children}</ExamContext.Provider>
  );
};

export const useExam = (): TestContextType => {
  const context = useContext(ExamContext);
  if (!context) {
    throw new Error("useExam must be used within an ExamProvider");
  }
  return context;
};

export const useExamResults = (): TestAttempt | null => {
  const { currentAttempt, isExamCompleted, isHydrated } = useExam();
  if (!isHydrated) return null;
  if (!isExamCompleted() || !currentAttempt) return null;
  return currentAttempt;
};