import { create } from 'zustand'
import { persist } from 'zustand/middleware'
import type { SessionPart } from '@/types'

interface StartExamConfig {
  testId: number | null
  testName: string
  parts: SessionPart[]
  totalSeconds: number  // 0 = no limit
}

interface TestStore {
  testId: number | null
  testName: string
  parts: SessionPart[]
  currentPartIndex: number
  answers: Record<number, number | null>  // questionId → answerIndex (0-3), null=unanswered
  isSubmitted: boolean
  timeUsed: number     // seconds elapsed when submitted
  totalSeconds: number // time limit (0 = no limit)

  startExam: (config: StartExamConfig) => void
  setAnswer: (questionId: number, answerIndex: number) => void
  setCurrentPart: (partIndex: number) => void
  submitExam: (timeUsed: number) => void
  reset: () => void
}

const INITIAL_STATE = {
  testId: null,
  testName: '',
  parts: [],
  currentPartIndex: 0,
  answers: {},
  isSubmitted: false,
  timeUsed: 0,
  totalSeconds: 0,
}

export const useTestStore = create<TestStore>()(
  persist(
    (set) => ({
      ...INITIAL_STATE,

      startExam: ({ testId, testName, parts, totalSeconds }) => {
        // Pre-fill all question IDs with null (unanswered)
        const answers: Record<number, number | null> = {}
        for (const part of parts) {
          for (const q of part.questions) answers[q.id] = null
        }
        set({ testId, testName, parts, currentPartIndex: 0, answers, isSubmitted: false, timeUsed: 0, totalSeconds })
      },

      setAnswer: (questionId, answerIndex) =>
        set((state) => ({
          answers: { ...state.answers, [questionId]: answerIndex },
        })),

      setCurrentPart: (partIndex) => set({ currentPartIndex: partIndex }),

      submitExam: (timeUsed) => set({ isSubmitted: true, timeUsed }),

      reset: () => set(INITIAL_STATE),
    }),
    { name: 'test-store', version: 2 },
  ),
)