playchoo-manager-nextjs/src/contexts/AssessmentContext.tsx

'use client';

import React, { createContext, useContext, useState, useCallback, ReactNode, useEffect } from 'react';
import { assessmentApi } from '@/src/utils/assessmentApi';
import { TranslationsContext } from '@/src/contexts/TranslationsContext';
import { useSwissOIDAuth } from 'swissoid-front';
import type {
  Question,
  ApiQuestion,
  GradingOption,
  AssessmentResult,
  AssessmentProgress,
  ServerAnswer,
  LocalAnswer,
  ContinueSessionResponse,
  StartSessionResponse,
  SaveAnswerResponse,
  GetResultResponse
} from '@/src/types/assessment';

interface AssessmentContextType {
  sessionId: string | null;
  questions: Question[];
  answersWithLabels: Record<number, { value: number; label: string }>; // Changed key from string to number
  result: AssessmentResult | null;
  progress: AssessmentProgress | null;
  isLoading: boolean;
  error: string | null;
  countryMode: 'CH' | 'FR';
  canModify: boolean;
  isLocked: boolean;

  startAssessment: (countryMode: 'CH' | 'FR') => Promise<void>;
  submitAnswer: (questionId: number, value: number, label: string) => Promise<void>; // Changed questionId to number
  submitAssessment: () => Promise<void>;
  saveResult: (token: string | null) => Promise<void>;
  lockSession: () => Promise<void>;
  unlockSession: () => Promise<void>;
  resetAssessment: () => void;
  setCountryMode: (mode: 'CH' | 'FR') => void;
  continueExistingSession: () => Promise<boolean>;
}

const AssessmentContext = createContext<AssessmentContextType | undefined>(undefined);

export function AssessmentProvider({ children }: { children: ReactNode }) {
  const [sessionId, setSessionId] = useState<string | null>(null);
  const [questions, setQuestions] = useState<Question[]>([]);
  const [answersWithLabels, setAnswersWithLabels] = useState<Record<number, { value: number; label: string }>>({});
  const [result, setResult] = useState<AssessmentResult | null>(null);
  const [progress, setProgress] = useState<AssessmentProgress | null>(null);
  const [isLoading, setIsLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [countryMode, setCountryMode] = useState<'CH' | 'FR'>('CH');
  const [canModify, setCanModify] = useState(true);
  const [isLocked, setIsLocked] = useState(false);

  // Get locale from translations context
  const translationsCtx = useContext(TranslationsContext);
  const locale = translationsCtx?.locale;

  // Get auth state
  const { user } = useSwissOIDAuth();
  const isLoggedIn = !!user;
  const token = typeof window !== 'undefined' ? localStorage.getItem('token') : null;

  // Load session from localStorage on mount
  useEffect(() => {
    const savedSessionId = localStorage.getItem('assessment_session_id');
    if (savedSessionId) {
      setSessionId(savedSessionId);
    }
  }, []);

  // Save session to localStorage when it changes
  useEffect(() => {
    if (sessionId) {
      localStorage.setItem('assessment_session_id', sessionId);
    }
  }, [sessionId]);

  // Associate session when user logs in
  useEffect(() => {
    if (isLoggedIn && sessionId && token) {
      assessmentApi.associateSession(sessionId, token).catch(console.error);
    }
  }, [isLoggedIn, sessionId, token]);

  const continueExistingSession = useCallback(async (): Promise<boolean> => {
    const savedSessionId = localStorage.getItem('assessment_session_id');
    if (!savedSessionId) return false;

    setIsLoading(true);
    setError(null);

    try {
      const data = await assessmentApi.continueSession(savedSessionId, locale) as ContinueSessionResponse;

      if (!data) {
        // Session not found
        localStorage.removeItem('assessment_session_id');
        return false;
      }

      // Handle both nested and direct data structures
      const responseData = data.data || data;
      console.log('Continue session response:', responseData);

      if (data.status === 'success' && responseData) {
        setSessionId(responseData.session_id);

        // Process questions to add grading array from grading_types
        const processedQuestions: Question[] = responseData.questions.map((q: ApiQuestion) => ({
          ...q,
          grading: responseData.grading_types?.[q.grading_type] || []
        }));

        // Sort questions by display_order to ensure correct sequence
        processedQuestions.sort((a, b) => (a.display_order || 0) - (b.display_order || 0));

        setQuestions(processedQuestions);

        // Convert new answer format to our internal format
        // Continue session now returns existing_answers.answers array
        if (responseData.existing_answers?.answers && Array.isArray(responseData.existing_answers.answers)) {
          const answersMap: Record<number, number> = {};
          const answersWithLabelsMap: Record<number, { value: number; label: string }> = {};

          responseData.existing_answers.answers.forEach((answer: ServerAnswer) => {
            // Find the question to get its grading type
            const question = processedQuestions.find((q: Question) => q.id === answer.question_id);
            let label = '';

            // Reconstruct label from grading options since server doesn't provide it
            if (question && question.grading) {
              const gradingOption = question.grading.find((g: GradingOption) => g.value === answer.value);
              label = gradingOption ? gradingOption.label : String(answer.value);
            } else {
              // Fallback to string value if no grading found
              label = String(answer.value);
            }

            answersMap[answer.question_id] = answer.value;
            answersWithLabelsMap[answer.question_id] = {
              value: answer.value,
              label: label
            };
          });

          setAnswersWithLabels(answersWithLabelsMap);

          const newProgress = {
            level: responseData.calculated_level || 0,
            confidence: responseData.confidence || 0,
            total_points: responseData.total_points || 0,
            is_complete: responseData.is_complete || false,
            answered_count: responseData.existing_answers.total_answered || 0
          };
          console.log('Setting progress from continued session:', newProgress);
          setProgress(newProgress);
        } else {
          setAnswersWithLabels({});

          const newProgress = {
            level: responseData.calculated_level || 0,
            confidence: responseData.confidence || 0,
            total_points: responseData.total_points || 0,
            is_complete: responseData.is_complete || false,
            answered_count: 0
          };
          console.log('Setting progress from continued session (no answers):', newProgress);
          setProgress(newProgress);
        }

        setCountryMode(responseData.country_mode || 'CH');
        setCanModify(responseData.can_modify !== false); // Default to true if not specified
        setIsLocked(responseData.is_locked || false);

        // If assessment is already complete, fetch the full result
        if (responseData.is_complete) {
          console.log('Continued session is already complete! Fetching full results...');
          const resultData = await assessmentApi.getResult(responseData.session_id);

          // Handle both nested and direct response structures
          const resultInfo = resultData.data || resultData;

          if (resultData.status === 'success' && resultInfo) {
            console.log('Full result fetched for continued session:', resultInfo);

            // Check if result is all zeros and use session data as fallback
            if (resultInfo.level === 0 && resultInfo.total_points === 0 && responseData.calculated_level) {
              console.warn('Result API returned zeros, using session data');
              setResult({
                level: responseData.calculated_level,
                total_points: responseData.total_points || 0,
                axis_breakdown: responseData.axis_breakdown || {
                  technical: 0,
                  tactical: 0,
                  competition: 0
                },
                confidence: responseData.confidence || 100,
                unmet_criteria: []
              });
            } else {
              setResult({
                level: resultInfo.level,
                total_points: resultInfo.total_points,
                axis_breakdown: resultInfo.axis_breakdown || {
                  technical: 0,
                  tactical: 0,
                  competition: 0
                },
                confidence: resultInfo.confidence,
                unmet_criteria: resultInfo.unmet_criteria || []
              });
            }
          }
        }

        return true;
      }

      return false;
    } catch (err) {
      setError(err instanceof Error ? err.message : 'An error occurred');
      return false;
    } finally {
      setIsLoading(false);
    }
  }, [locale]);

  const startAssessment = useCallback(async (mode: 'CH' | 'FR') => {
    setIsLoading(true);
    setError(null);

    try {
      const data = await assessmentApi.startSession(mode, locale, token || undefined) as StartSessionResponse;

      if (data.status === 'success') {
        // API returns session_id and questions at root level
        setSessionId(data.session_id);

        // Process questions - they might have grading inline or need it from grading_types
        const processedQuestions: Question[] = data.questions.map((q: ApiQuestion) => {
          // Always get grading from grading_types based on grading_type
          return {
            ...q,
            grading: data.grading_types?.[q.grading_type] || []
          };
        });

        // Sort questions by display_order to ensure correct sequence
        processedQuestions.sort((a, b) => (a.display_order || 0) - (b.display_order || 0));

        setQuestions(processedQuestions);
        setCountryMode(mode);

        // Check for existing answers (authenticated users)
        // Note: In start response, existing_answers is just an array, not wrapped
        if (data.existing_answers && Array.isArray(data.existing_answers)) {
          const answersMap: Record<number, number> = {};
          const answersWithLabelsMap: Record<number, { value: number; label: string }> = {};

          data.existing_answers.forEach((answer: ServerAnswer) => {
            // Find the question to get its grading type
            const question = processedQuestions.find((q: Question) => q.id === answer.question_id);
            let label = '';

            // Reconstruct label from grading options since server might not provide it
            if (question && question.grading) {
              const gradingOption = question.grading.find((g: GradingOption) => g.value === answer.value);
              label = gradingOption ? gradingOption.label : String(answer.value);
            } else {
              // Fallback to string value if no grading found
              label = String(answer.value);
            }

            answersWithLabelsMap[answer.question_id] = {
              value: answer.value,
              label: label
            };
          });

          setAnswersWithLabels(answersWithLabelsMap);
          setProgress({
            level: data.calculated_level || 0,
            confidence: data.confidence || 0,
            total_points: data.total_points || 0,
            is_complete: data.is_complete || false,
            answered_count: data.existing_answers.length || 0
          });
        } else {
          setAnswersWithLabels({});
          setProgress({
            level: 0,
            confidence: 0,
            total_points: 0,
            is_complete: false,
            answered_count: 0
          });
        }
        setResult(null);
      } else {
        throw new Error(data.message || 'Failed to start assessment');
      }
    } catch (err) {
      setError(err instanceof Error ? err.message : 'An error occurred');
      throw err;
    } finally {
      setIsLoading(false);
    }
  }, [locale, token]);

  const submitAnswer = useCallback(async (questionId: number, value: number, label: string) => {
    if (!sessionId) {
      setError('No active session');
      return;
    }

    try {
      // Save answer to backend FIRST (before updating local state)
      const data = await assessmentApi.saveAnswer(sessionId, questionId, value, locale) as SaveAnswerResponse;
      console.log('Raw backend response from saveAnswer:', data);

      if (data.status !== 'success') return;
      // NOW update the state after successful API response
      setAnswersWithLabels(prev => ({
        ...prev,
        [questionId]: { value, label }
      }));

      // Backend returns data directly, not nested under data field
      const newProgress = {
        level: data.level || 0,
        confidence: data.confidence || 0,
        total_points: data.total_points || 0,
        is_complete: data.is_complete || false,
        answered_count: Object.keys(answersWithLabels).length + 1
      };

      console.log('Setting progress after answer submission:', newProgress);
      setProgress(newProgress);

      if (!data.is_complete) return;

      // If assessment is complete, fetch full result
      console.log('Assessment complete! Fetching full results...');
      // Fetch the full assessment result
      const resultData = await assessmentApi.getResult(sessionId);
      console.log('Full result response:', resultData);

      // Handle both nested and direct response structures
      const resultInfo = resultData.data || resultData;

      if (resultData.status !== 'success' || !resultInfo) return;
      console.log('Setting result from completion:', resultInfo);

      // Check if the result has actual data or just zeros
      if (resultInfo.level === 0 && resultInfo.total_points === 0) {
        console.warn('Result has all zeros - backend may not have calculated properly');
        // Use the progress data as fallback if available
        if (data.level && data.total_points) {
          console.log('Using progress data as fallback:', data);
          setResult({
            level: data.level,
            total_points: data.total_points,
            axis_breakdown: data.axis_breakdown || {
              technical: 0,
              tactical: 0,
              competition: 0
            },
            confidence: data.confidence || 100,
            unmet_criteria: []
          });
        } else {
          setResult({
            level: resultInfo.level,
            total_points: resultInfo.total_points,
            axis_breakdown: resultInfo.axis_breakdown || {
              technical: 0,
              tactical: 0,
              competition: 0
            },
            confidence: resultInfo.confidence,
            unmet_criteria: resultInfo.unmet_criteria || []
          });
        }
      } else {
        setResult({
          level: resultInfo.level,
          total_points: resultInfo.total_points,
          axis_breakdown: resultInfo.axis_breakdown || {
            technical: 0,
            tactical: 0,
            competition: 0
          },
          confidence: resultInfo.confidence,
          unmet_criteria: resultInfo.unmet_criteria || []
        });
      }
    } catch (err) {
      // Don't update state on error - nothing to revert since we didn't update yet
      setError(err instanceof Error ? err.message : 'Failed to save answer');
      throw err;
    }
  }, [sessionId, answersWithLabels, locale]);

  const submitAssessment = useCallback(async () => {
    if (!sessionId) {
      setError('No active session');
      return;
    }

    setIsLoading(true);
    setError(null);

    try {
      const data = await assessmentApi.submitAnswers(sessionId, Object.entries(answersWithLabels).reduce(
        (p, [id, { value }]) => ({ ...p, [id]: value, }),
        {}
      ));

      if (data.status === 'success') {
        // Check if result is at root level or nested
        const result = data.result || data.data?.result;
        if (result) {
          setResult(result);
        } else {
          throw new Error('No result returned from assessment');
        }
      } else {
        throw new Error(data.message || 'Failed to submit assessment');
      }
    } catch (err) {
      setError(err instanceof Error ? err.message : 'An error occurred');
      throw err;
    } finally {
      setIsLoading(false);
    }
  }, [sessionId, answersWithLabels]);

  const saveResult = useCallback(async (token: string | null) => {
    if (!sessionId) {
      setError('No active session');
      return;
    }

    setIsLoading(true);
    setError(null);

    try {
      const data = await assessmentApi.saveResult(sessionId, token);

      if (data.status === 'success' || data.level !== undefined) {
        // Successfully saved
        setIsLocked(true);
        setCanModify(false);
      } else {
        throw new Error(data.message || 'Failed to save result');
      }
    } catch (err) {
      setError(err instanceof Error ? err.message : 'An error occurred');
      throw err;
    } finally {
      setIsLoading(false);
    }
  }, [sessionId]);

  const lockSession = useCallback(async () => {
    if (!sessionId) {
      setError('No active session');
      return;
    }

    setIsLoading(true);
    setError(null);

    try {
      // Pass null for token - the API will use cookies for auth
      const data = await assessmentApi.lockSession(sessionId, null);

      if (data.status === 'success') {
        setIsLocked(true);
        setCanModify(false);
        console.log('Session locked successfully');
        return data.data;
      } else {
        throw new Error(data.message || 'Failed to lock session');
      }
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Failed to lock session');
      throw err;
    } finally {
      setIsLoading(false);
    }
  }, [sessionId]);

  const unlockSession = useCallback(async () => {
    if (!sessionId) {
      setError('No active session');
      return;
    }

    setIsLoading(true);
    setError(null);

    try {
      const data = await assessmentApi.unlockSession(sessionId, token);

      if (data.status === 'success') {
        setIsLocked(false);
        setCanModify(true);
        console.log('Session unlocked successfully');
        return data.data;
      } else {
        throw new Error(data.message || 'Failed to unlock session');
      }
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Failed to unlock session');
      throw err;
    } finally {
      setIsLoading(false);
    }
  }, [sessionId, token]);

  const resetAssessment = useCallback(() => {
    setSessionId(null);
    setQuestions([]);
    setAnswersWithLabels({});
    setResult(null);
    setProgress(null);
    setError(null);
    setCanModify(true);
    setIsLocked(false);
    localStorage.removeItem('assessment_session_id');
  }, []);

  const value: AssessmentContextType = {
    sessionId,
    questions,
    answersWithLabels,
    result,
    progress,
    isLoading,
    error,
    countryMode,
    canModify,
    isLocked,
    startAssessment,
    submitAnswer,
    submitAssessment,
    saveResult,
    lockSession,
    unlockSession,
    resetAssessment,
    setCountryMode,
    continueExistingSession
  };

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

export function useAssessment() {
  const context = useContext(AssessmentContext);
  if (context === undefined) {
    throw new Error('useAssessment must be used within an AssessmentProvider');
  }
  return context;
}