import { fetchWithAuth } from "@/lib/apiClient";
import { create } from 'zustand'
import { persist, StateStorage, createJSONStorage } from 'zustand/middleware'
import { get, set, del } from 'idb-keyval'

// Custom IndexedDB storage engine for Zustand
const idbStorage: StateStorage = {
  getItem: async (name: string): Promise<string | null> => {
    return (await get(name)) || null
  },
  setItem: async (name: string, value: string): Promise<void> => {
    await set(name, value)
  },
  removeItem: async (name: string): Promise<void> => {
    await del(name)
  },
}

export type QuestionStatus = 'UNVISITED' | 'NOT_ANSWERED' | 'ANSWERED' | 'MARKED' | 'ANSWERED_MARKED'

export interface AnswerState {
  question_id: string
  selected_answer: string | null
  status: QuestionStatus
  time_spent: number // in seconds
}

export interface ExamAttemptPayload {
  id: string
  user_id: string
  source_type: 'DPS' | 'Test' | 'MicroTest'
  source_id: string
  started_at: string
  submitted_at: string
  time_taken: number
  answers: Array<{
    question_id: string
    selected_answer: string | null
    is_correct: boolean // Typically evaluated on server, but requested in payload structure
    is_subjective: boolean
    time_spent: number
  }>
  score: number
  total_questions: number
  accuracy: number
  status: 'In_Progress' | 'Submitted' | 'Timed_Out'
}

interface CBTState {
  userId: string | null
  examId: string | null
  attemptId: string | null
  status: 'NOT_STARTED' | 'IN_PROGRESS' | 'SUBMITTED' | 'TIMED_OUT'
  startTime: number | null
  totalDuration: number // in seconds
  questions: any[]
  answers: Record<string, AnswerState>
  currentQuestionIndex: number
  
  // Actions
  initExam: (examId: string, questions: any[], duration: number, userId: string) => void
  startExam: () => void
  selectAnswer: (questionId: string, answer: string) => void
  clearResponse: (questionId: string) => void
  markForReview: (questionId: string) => void
  goToQuestion: (index: number) => void
  submitExam: (reason?: 'SUBMITTED' | 'TIMED_OUT') => void
  addTimeSpent: (questionId: string, seconds: number) => void
  getAttemptPayload: (userId: string) => ExamAttemptPayload | null
  submitAttemptToServer: (userId: string) => Promise<string | null>
  resetExamState: () => void
}

export const useCBTStore = create<CBTState>()(
  persist(
    (set, get) => ({
      userId: null,
      examId: null,
      attemptId: null,
      status: 'NOT_STARTED',
      startTime: null,
      totalDuration: 0,
      questions: [],
      answers: {},
      currentQuestionIndex: 0,

      initExam: (examId, questions, duration, userId) => {
        const currentExamId = get().examId;
        const currentUserId = get().userId;
        const currentQuestions = get().questions;
        
        // Force reset if it's a new exam, new user, or the number of questions has changed
        const needsReset = currentExamId !== examId || currentUserId !== userId || currentQuestions.length !== questions.length;

        if (needsReset) {
          const initialAnswers: Record<string, AnswerState> = {};
          questions.forEach((q, idx) => {
            initialAnswers[q.id] = {
              question_id: q.id,
              selected_answer: null,
              status: idx === 0 ? 'NOT_ANSWERED' : 'UNVISITED',
              time_spent: 0
            }
          })

          set({
            userId,
            examId,
            attemptId: null,
            status: 'NOT_STARTED',
            startTime: null,
            totalDuration: duration,
            questions,
            answers: initialAnswers,
            currentQuestionIndex: 0
          })
        } else {
          // Just update questions array in case text/options changed
          set({ questions, totalDuration: duration })
        }
      },

      startExam: () => {
        set({ status: 'IN_PROGRESS', startTime: Date.now() })
      },

      selectAnswer: (questionId, answer) => set((state) => {
        const current = state.answers[questionId]
        if (!current) return state;
        
        let newStatus: QuestionStatus = 'ANSWERED';
        if (current.status === 'MARKED' || current.status === 'ANSWERED_MARKED') {
          newStatus = 'ANSWERED_MARKED';
        }
        
        return {
          answers: {
            ...state.answers,
            [questionId]: { ...current, selected_answer: answer, status: newStatus }
          }
        }
      }),

      clearResponse: (questionId) => set((state) => {
        const current = state.answers[questionId]
        if (!current) return state;

        let newStatus: QuestionStatus = 'NOT_ANSWERED';
        if (current.status === 'ANSWERED_MARKED' || current.status === 'MARKED') {
          newStatus = 'MARKED';
        }

        return {
          answers: {
            ...state.answers,
            [questionId]: { ...current, selected_answer: null, status: newStatus }
          }
        }
      }),

      markForReview: (questionId) => set((state) => {
        const current = state.answers[questionId]
        if (!current) return state;

        let newStatus: QuestionStatus = current.selected_answer ? 'ANSWERED_MARKED' : 'MARKED';
        
        // If already marked, unmark it
        if (current.status === 'MARKED') newStatus = 'NOT_ANSWERED';
        if (current.status === 'ANSWERED_MARKED') newStatus = 'ANSWERED';

        return {
          answers: {
            ...state.answers,
            [questionId]: { ...current, status: newStatus }
          }
        }
      }),

      goToQuestion: (index) => set((state) => {
        if (index < 0 || index >= state.questions.length) return state;
        
        const targetQId = state.questions[index].id;
        const targetState = state.answers[targetQId];
        
        // Update status of the target question to NOT_ANSWERED if it was UNVISITED
        const newAnswers = { ...state.answers };
        if (targetState.status === 'UNVISITED') {
          newAnswers[targetQId] = { ...targetState, status: 'NOT_ANSWERED' };
        }

        return {
          currentQuestionIndex: index,
          answers: newAnswers
        }
      }),

      addTimeSpent: (questionId, seconds) => set((state) => {
        const current = state.answers[questionId]
        if (!current) return state;
        return {
          answers: {
            ...state.answers,
            [questionId]: { ...current, time_spent: current.time_spent + seconds }
          }
        }
      }),

      submitExam: (reason = 'SUBMITTED') => set({ status: reason }),

      getAttemptPayload: (userId: string) => {
        const state = get();
        if (!state.examId || !state.startTime) return null;

        const timeTaken = Math.floor((Date.now() - state.startTime) / 1000);
        const answerList = Object.values(state.answers);
        
        // Backend will perform secure evaluation
        let correctCount = 0; // Removing mock count
        const finalAnswers = answerList.map(a => {
          const q = state.questions.find(quest => quest.id === a.question_id);
          const isSubjective = q?.type === 'Subjective' || q?.type === 'Numerical';
          
          return {
            question_id: a.question_id,
            selected_answer: a.selected_answer,
            is_correct: false,
            is_subjective: isSubjective,
            time_spent: a.time_spent
          }
        });

        const totalQ = state.questions.length;
        const payload: ExamAttemptPayload = {
          id: crypto.randomUUID(), // Mock Attempt ID
          user_id: userId,
          source_type: 'Test', // Or DPS based on exam metadata
          source_id: state.examId,
          started_at: new Date(state.startTime).toISOString(),
          submitted_at: new Date().toISOString(),
          time_taken: timeTaken,
          answers: finalAnswers,
          score: correctCount * 4, // Assuming +4 logic
          total_questions: totalQ,
          accuracy: totalQ > 0 ? (correctCount / totalQ) * 100 : 0,
          status: state.status === 'IN_PROGRESS' ? 'Submitted' : (state.status === 'TIMED_OUT' ? 'Timed_Out' : 'Submitted')
        }

        return payload;
      },

      submitAttemptToServer: async (userId: string) => {
        const payload = get().getAttemptPayload(userId);
        if (!payload) return null;

        try {
          const res = await fetchWithAuth(`/api/attempts/submit`, {
            method: 'POST',
            headers: { 'Content-Type': 'application/json' },
            body: JSON.stringify(payload),
          });
          
          if (!res.ok) {
            console.error("Failed to submit attempt to server", await res.text());
            return null;
          } else {
            const data = await res.json();
            console.log("Successfully submitted attempt to server", data.id);
            set({ attemptId: data.id });
            return data.id;
          }
        } catch (error) {
          console.error("Network error during submission:", error);
          return null;
        }
      },

      resetExamState: () => {
        set({
          userId: null,
          examId: null,
          attemptId: null,
          status: 'NOT_STARTED',
          startTime: null,
          totalDuration: 0,
          questions: [],
          answers: {},
          currentQuestionIndex: 0
        });
      }
    }),
    {
      name: 'cbt-storage', // unique name
      storage: createJSONStorage(() => idbStorage), // use custom IndexedDB storage
      partialize: (state) => ({
        userId: state.userId,
        examId: state.examId,
        status: state.status,
        startTime: state.startTime,
        totalDuration: state.totalDuration,
        questions: state.questions,
        answers: state.answers,
        currentQuestionIndex: state.currentQuestionIndex,
      }),
    }
  )
)