// components/QuizBlock.jsx
import { useState, useEffect, useRef, useCallback } from "react";
import { z } from "zod";
import { Button } from "@/components/ui/button";
import { Skeleton } from "@/components/ui/skeleton";
import {
    ChevronLeft, ChevronRight,
    Circle, CheckCircle2,
    Square, CheckSquare2,
    Clock, AlertTriangle, Info, ArrowRight,
} from "lucide-react";
import { useClientNotifications } from "@/contexts/ClientNotificationContext";
import { useDateFormat } from "@/hooks/useDateFormat";

// Zod schema: every question must have a non-empty answer before submitting.
function buildAnswerSchema(questions = []) {
    const shape = {};
    for (const q of questions) {
        const key = String(q.question_id);
        shape[key] = q.type === "multi_select"
            ? z.array(z.unknown()).min(1, "This question requires at least one selection.")
            : z.union([z.string(), z.number()]).refine((v) => v !== "" && v != null, "This question requires an answer.");
    }
    return z.object(shape);
}

function QuizSkeleton() {
    return (
        <div className="max-w-2xl mx-auto space-y-5">
            <Skeleton className="h-5 w-1/3" />
            <Skeleton className="h-1.5 w-full rounded-full" />
            <Skeleton className="h-16 w-full rounded-xl" />
            <div className="space-y-2">
                <Skeleton className="h-11 w-full rounded-lg" />
                <Skeleton className="h-11 w-full rounded-lg" />
                <Skeleton className="h-11 w-full rounded-lg" />
                <Skeleton className="h-11 w-full rounded-lg" />
            </div>
        </div>
    );
}

function formatTime(secs) {
    if (secs === null || secs === undefined) return null;
    const m = Math.floor(secs / 60).toString().padStart(2, "0");
    const s = (secs % 60).toString().padStart(2, "0");
    return `${m}:${s}`;
}

function isQuestionAnswered(answer) {
    return Array.isArray(answer) ? answer.length > 0 : answer !== undefined;
}

/**
 * Props:
 *   quiz     — assessment data including active_session, time_limit_minutes, etc.
 *   loading  — true while fetch is in-flight
 *   onStart  — async () => { session_id, expires_at, remaining_seconds, draft_answers }
 *   onDraft          — (answers) => void — called every 25s
 *   onRefreshSession — async () => { expires_at, remaining_seconds }
 *   onSubmit — async (answers, sessionId?) => result|null
 *   onRetake — () => void
 *   onActiveChange   — (isActive: boolean) => void — fires when session starts/ends
 *   label    — "Quiz" or "Assessment"
 */
const QuizBlock = ({ quiz, loading = false, onStart, onDraft, onRefreshSession, onSubmit, onRetake, onActiveChange, onNextContent, nextLabel, label = "Quiz" }) => {
    const { fmtDateTime } = useDateFormat();
    const questions = quiz?.questions ?? [];
    const total = questions.length;

    // ── Core stage state ──────────────────────────────────────────────────────
    const [stage, setStage]               = useState("intro");   // 'intro' | 'taking' | 'review' | 'result'
    const [currentIndex, setCurrentIndex] = useState(0);
    const [answers, setAnswers]           = useState({});
    const [submitting, setSubmitting]     = useState(false);
    const [starting, setStarting]         = useState(false);
    const [result, setResult]             = useState(null);
    const [submitError, setSubmitError]   = useState(null);

    // ── Timer state ───────────────────────────────────────────────────────────
    const [remainingSeconds, setRemainingSeconds] = useState(null);
    const [timeExpired, setTimeExpired]           = useState(false);

    // ── Notifications ─────────────────────────────────────────────────────────
    const { notifications, fetchNotifications, accelerate, decelerate } = useClientNotifications();
    const [assessmentUpdatedAlert, setAssessmentUpdatedAlert] = useState(false);
    const seenNotifRef = useRef(new Set());

    // ── Session refs ──────────────────────────────────────────────────────────
    const sessionRef    = useRef({ sessionId: null, expiresAt: null });
    const answersRef    = useRef({});
    const timerFiredRef = useRef(false);

    useEffect(() => { answersRef.current = answers; }, [answers]);

    // Notify parent when an active session starts or ends
    useEffect(() => {
        onActiveChange?.(stage === "taking" || stage === "review");
    }, [stage]); // eslint-disable-line react-hooks/exhaustive-deps

    // Always deactivate when component unmounts (user navigated away mid-session)
    useEffect(() => () => { onActiveChange?.(false); }, []); // eslint-disable-line react-hooks/exhaustive-deps

    useEffect(() => {
        setStage("intro");
        setResult(null);
        setRemainingSeconds(null);
        setTimeExpired(false);
        sessionRef.current = { sessionId: null, expiresAt: null };
        timerFiredRef.current = false;
        const draft = quiz?.active_session?.draft_answers ?? {};
        const hasDraft = Object.keys(draft).length > 0;
        setAnswers(hasDraft ? draft : {});
        answersRef.current = hasDraft ? draft : {};
        if (hasDraft) {
            const qs = quiz?.questions ?? [];
            const firstUnanswered = qs.findIndex(q => !isQuestionAnswered(draft[q.question_id]));
            setCurrentIndex(firstUnanswered > -1 ? firstUnanswered : 0);
        } else {
            setCurrentIndex(0);
        }
    }, [quiz?.quiz_id]);

    useEffect(() => {
        if (stage !== "taking") { decelerate(); return; }
        accelerate();
        fetchNotifications();
        return () => decelerate();
    }, [stage]); // eslint-disable-line react-hooks/exhaustive-deps

    useEffect(() => {
        if (stage !== "taking") return;
        notifications.forEach((n) => {
            if (n.type === "assessment" && n.title === "Assessment Updated" && !seenNotifRef.current.has(n.notification_id)) {
                seenNotifRef.current.add(n.notification_id);
                setAssessmentUpdatedAlert(true);
                if (onRefreshSession) {
                    onRefreshSession().then(updated => {
                        if (updated?.remaining_seconds != null) {
                            setRemainingSeconds(updated.remaining_seconds);
                            sessionRef.current.expiresAt = updated.expires_at ? new Date(updated.expires_at) : null;
                        }
                    });
                }
            }
        });
    }, [notifications, stage]); // eslint-disable-line react-hooks/exhaustive-deps

    useEffect(() => {
        if (stage !== "taking" || !onDraft) return;
        onDraft(answersRef.current);
        const id = setInterval(() => onDraft(answersRef.current), 25_000);
        return () => clearInterval(id);
    }, [stage, onDraft]); // eslint-disable-line react-hooks/exhaustive-deps

    // Save immediately on every answer change — ensures no answer is lost if the tab closes
    // before the 25s interval fires. Guards on stage and non-empty answers to skip resets.
    useEffect(() => {
        if (stage !== "taking" || !onDraft || Object.keys(answers).length === 0) return;
        onDraft(answers);
    }, [answers]); // eslint-disable-line react-hooks/exhaustive-deps

    useEffect(() => {
        if (stage !== "taking" || !sessionRef.current.expiresAt) return;
        timerFiredRef.current = false;
        const tick = () => {
            const secs = Math.max(0, Math.ceil((sessionRef.current.expiresAt.getTime() - Date.now()) / 1000));
            setRemainingSeconds(secs);
            if (secs <= 0 && !timerFiredRef.current) {
                timerFiredRef.current = true;
                setTimeExpired(true);
                clearInterval(id);
            }
        };
        tick();
        const id = setInterval(tick, 1000);
        return () => clearInterval(id);
    }, [stage]);

    useEffect(() => {
        if (!timeExpired || submitting) return;
        const autoSubmit = async () => {
            setSubmitting(true);
            const res = await onSubmit?.(answersRef.current, sessionRef.current.sessionId);
            setSubmitting(false);
            if (res) { setResult(res); setStage("result"); }
        };
        autoSubmit();
    }, [timeExpired]); // eslint-disable-line react-hooks/exhaustive-deps

    // ── Handlers ─────────────────────────────────────────────────────────────
    const handleStart = useCallback(async () => {
        const isAssessment = label === "Assessment";
        const hasTimeLimit = isAssessment && (quiz?.time_limit_minutes ?? 0) > 0;

        if (isAssessment && onStart) {
            setStarting(true);
            const session = await onStart();
            setStarting(false);
            if (!session) return;
            sessionRef.current = {
                sessionId: session.session_id,
                expiresAt: session.expires_at ? new Date(session.expires_at) : null,
            };
            if (hasTimeLimit && session.remaining_seconds != null) setRemainingSeconds(session.remaining_seconds);
            if (session.draft_answers && Object.keys(session.draft_answers).length > 0) {
                setAnswers(session.draft_answers);
                answersRef.current = session.draft_answers;
                const qs = quiz?.questions ?? [];
                const firstUnanswered = qs.findIndex(q => !isQuestionAnswered(session.draft_answers[q.question_id]));
                setCurrentIndex(firstUnanswered > -1 ? firstUnanswered : 0);
            }
        }
        setStage("taking");
    }, [label, quiz?.time_limit_minutes, onStart]);

    const handleRetake = () => {
        setCurrentIndex(0);
        setAnswers({});
        setResult(null);
        setRemainingSeconds(null);
        setTimeExpired(false);
        sessionRef.current = { sessionId: null, expiresAt: null };
        timerFiredRef.current = false;
        setStage("intro");
        onRetake?.();
    };

    const handleSubmit = useCallback(async () => {
        // Zod: all questions must be answered before submission is allowed
        const schema = buildAnswerSchema(questions);
        const stringifiedAnswers = Object.fromEntries(
            Object.entries(answers).map(([k, v]) => [String(k), v])
        );
        const parsed = schema.safeParse(stringifiedAnswers);
        if (!parsed.success) {
            const unanswered = questions.filter(
                (q) => !isQuestionAnswered(answers[q.question_id])
            ).length;
            setSubmitError(`Answer all questions before submitting — ${unanswered} still unanswered.`);
            return;
        }
        setSubmitError(null);
        setSubmitting(true);
        const res = await onSubmit?.(answers, sessionRef.current.sessionId);
        setSubmitting(false);
        if (res) { setResult(res); setStage("result"); }
    }, [answers, questions, onSubmit]);

    const handleOptionClick = (optionId) => {
        const question = questions[currentIndex];
        const isMulti  = question.type === "multi_select";
        const multiLimit = isMulti ? (question.correct_count ?? null) : null;
        setAnswers((prev) => {
            if (!isMulti) return { ...prev, [question.question_id]: optionId };
            const current = prev[question.question_id] ?? [];
            if (!current.includes(optionId) && multiLimit !== null && current.length >= multiLimit) return prev;
            const next = current.includes(optionId)
                ? current.filter((id) => id !== optionId)
                : [...current, optionId];
            return { ...prev, [question.question_id]: next };
        });
    };

    const handleGoToReview = () => {
        setStage("review");
    };

    const handleNext = () => {
        if (isLast) { handleGoToReview(); return; }
        setCurrentIndex((i) => Math.min(i + 1, total - 1));
    };

    const handlePrev = () => {
        if (isFirst) {
            if (label !== "Assessment") setStage("intro");
            return;
        }
        setCurrentIndex((i) => Math.max(i - 1, 0));
    };

    // ── Guards ────────────────────────────────────────────────────────────────
    if (!quiz && !loading) {
        return (
            <div className="h-full w-full flex items-center justify-center text-muted-foreground py-20">
                <p className="text-sm">No {label.toLowerCase()} available yet.</p>
            </div>
        );
    }
    if (loading) return <QuizSkeleton />;
    if (!total) {
        return (
            <div className="h-full w-full flex items-center justify-center text-muted-foreground py-20">
                <p className="text-sm">This {label.toLowerCase()} doesn't have any questions yet.</p>
            </div>
        );
    }

    // ── Derived values ────────────────────────────────────────────────────────
    const isAssessment  = label === "Assessment";
    const hasTimeLimit  = isAssessment && (quiz?.time_limit_minutes ?? 0) > 0;
    const activeSession = quiz?.active_session ?? null;

    // ── ASSESSMENT INTRO ──────────────────────────────────────────────────────
    if (stage === "intro" && isAssessment) {
        const attempts          = quiz.attempt_count ?? 0;
        const attemptsRemaining = quiz.attempts_remaining ?? null;
        const cooldownUntil     = quiz.cooldown_until ? new Date(quiz.cooldown_until) : null;
        const canAttempt        = quiz.can_attempt ?? true;

        const attemptsDisplay = attemptsRemaining !== null && quiz.max_attempts
            ? `${attemptsRemaining} / ${quiz.max_attempts}`
            : attemptsRemaining !== null
            ? `${attemptsRemaining}`
            : attempts > 0 ? `${attempts}` : "—";

        return (
            <div className="max-w-2xl mx-auto space-y-6">
                {/* Header */}
                <div className="space-y-1">
                    <p className="text-xs font-semibold uppercase tracking-widest text-muted-foreground">
                        {label}
                    </p>
                    <h1 className="text-2xl font-bold sm:text-3xl">{quiz.title || label}</h1>
                    {quiz.description && (
                        <p className="text-sm text-muted-foreground">{quiz.description}</p>
                    )}
                    {quiz.is_required && !quiz.description && (
                        <p className="text-sm text-muted-foreground">
                            Required to complete this course.
                        </p>
                    )}
                </div>

                {/* Stats grid */}
                <div className="grid grid-cols-2 sm:grid-cols-4 gap-3">
                    <div className="rounded-xl border bg-card p-4 space-y-0.5">
                        <p className="text-2xl font-bold">{total}</p>
                        <p className="text-xs text-muted-foreground">Questions</p>
                    </div>
                    <div className="rounded-xl border bg-card p-4 space-y-0.5">
                        <p className="text-2xl font-bold">
                            {hasTimeLimit ? `${quiz.time_limit_minutes} min` : "—"}
                        </p>
                        <p className="text-xs text-muted-foreground">Time limit</p>
                    </div>
                    <div className="rounded-xl border bg-card p-4 space-y-0.5">
                        <p className="text-2xl font-bold">{quiz.passing_score}%</p>
                        <p className="text-xs text-muted-foreground">To pass</p>
                    </div>
                    <div className="rounded-xl border bg-card p-4 space-y-0.5">
                        <p className="text-2xl font-bold">{attemptsDisplay}</p>
                        <p className="text-xs text-muted-foreground">Attempts left</p>
                    </div>
                </div>

                {/* Status banners */}
                {quiz.has_passed && (
                    <div className="rounded-lg border border-green-500/30 bg-green-500/5 px-4 py-3">
                        <p className="text-sm font-medium text-green-700 dark:text-green-400">
                            You've already passed this assessment
                        </p>
                        <p className="text-sm">
                            Best score: {quiz.best_attempt?.score}%
                            {quiz.best_attempt?.passing_score != null && (
                                <span className="text-muted-foreground"> · passing score {quiz.best_attempt.passing_score}%</span>
                            )}
                        </p>
                    </div>
                )}
                {!canAttempt && cooldownUntil && (
                    <div className="rounded-lg border border-amber-500/30 bg-amber-500/5 px-4 py-3">
                        <p className="text-sm font-medium text-amber-700 dark:text-amber-400">You're on cooldown</p>
                        <p className="text-sm">You can retake this assessment after {fmtDateTime(cooldownUntil)}</p>
                    </div>
                )}
                {!canAttempt && !cooldownUntil && attemptsRemaining === 0 && (
                    <div className="rounded-lg border border-red-500/30 bg-red-500/5 px-4 py-3">
                        <p className="text-sm font-medium text-red-700 dark:text-red-400">No attempts remaining</p>
                        <p className="text-sm">
                            {quiz.window_reset_at
                                ? `You can try again after ${fmtDateTime(quiz.window_reset_at)}`
                                : "You've used all available attempts for this assessment"}
                        </p>
                    </div>
                )}
                {activeSession && canAttempt && (
                    <div className="rounded-lg border border-blue-500/30 bg-blue-500/5 px-4 py-3">
                        <p className="text-sm font-medium text-blue-700 dark:text-blue-400 flex items-center gap-1.5">
                            <Clock className="size-4" /> Session in progress
                        </p>
                        <p className="text-xs text-muted-foreground mt-1">
                            You have an unfinished attempt.{" "}
                            {activeSession.remaining_seconds != null
                                ? `${formatTime(activeSession.remaining_seconds)} remaining — resume before time runs out.`
                                : "Resume where you left off."}
                        </p>
                    </div>
                )}

                {/* Before you begin */}
                <div className="rounded-xl border bg-card p-5 space-y-3">
                    <p className="font-semibold text-sm">Before you begin</p>
                    <ul className="space-y-2.5 text-sm text-foreground/80">
                        {hasTimeLimit && (
                            <li className="flex items-start gap-2.5">
                                <span className="mt-2 size-1.5 rounded-full bg-primary shrink-0" />
                                <span>
                                    The timer starts the moment you begin and{" "}
                                    <strong>cannot be paused</strong>. Your assessment auto-submits when time runs out.
                                </span>
                            </li>
                        )}
                        <li className="flex items-start gap-2.5">
                            <span className="mt-2 size-1.5 rounded-full bg-primary shrink-0" />
                            <span>You can skip questions and return to them later. All questions must be answered before you can submit.</span>
                        </li>
                        <li className="flex items-start gap-2.5">
                            <span className="mt-2 size-1.5 rounded-full bg-primary shrink-0" />
                            <span>
                                Your progress is <strong>saved automatically</strong> — you can safely resume if you lose connection.
                            </span>
                        </li>
                        <li className="flex items-start gap-2.5">
                            <span className="mt-2 size-1.5 rounded-full bg-primary shrink-0" />
                            <span>Review all answers on the summary screen before final submission.</span>
                        </li>
                    </ul>
                </div>

                <Button
                    size="lg"
                    onClick={handleStart}
                    disabled={!canAttempt || starting}
                >
                    {starting
                        ? "Starting…"
                        : activeSession
                        ? "Resume Assessment"
                        : attempts > 0
                        ? "Retake Assessment"
                        : "Begin assessment"}
                </Button>
            </div>
        );
    }

    // ── QUIZ INTRO ────────────────────────────────────────────────────────────
    if (stage === "intro") {
        const attempts          = quiz.attempt_count ?? 0;
        const attemptsRemaining = quiz.attempts_remaining ?? null;
        const cooldownUntil     = quiz.cooldown_until ? new Date(quiz.cooldown_until) : null;
        const canAttempt        = quiz.can_attempt ?? true;
        const hasActiveSession  = !!quiz.active_session;

        return (
            <div className="max-w-2xl mx-auto">
                <div className="rounded-xl border bg-card p-6 space-y-6 text-center sm:p-10 shadow-lg">
                    <div className="space-y-1.5">
                        <h2 className="text-xl font-semibold sm:text-2xl">{quiz.title || label}</h2>
                        {quiz.is_required && (
                            <span className="inline-block rounded-full bg-amber-500/10 px-2.5 py-0.5 text-xs font-medium text-amber-600">
                                Required to complete this unit
                            </span>
                        )}
                    </div>

                    {quiz.has_passed && !hasActiveSession && (
                        <div className="rounded-lg border border-green-500/30 bg-green-500/5 px-4 py-3">
                            <p className="text-md font-medium text-green-700 dark:text-green-400">You've already passed this quiz</p>
                            <p className="text-sm">Best score: {quiz.best_attempt?.score}%</p>
                        </div>
                    )}

                    {hasActiveSession && (
                        <div className="rounded-lg border border-amber-500/30 bg-amber-500/10 px-4 py-3 flex items-start gap-3 text-left">
                            <Clock className="size-4 text-amber-500 shrink-0 mt-0.5" />
                            <div>
                                <p className="text-sm font-semibold text-amber-700 dark:text-amber-400">You have an unfinished session</p>
                                <p className="text-xs text-amber-600 dark:text-amber-500 mt-0.5 leading-relaxed">
                                    You left this quiz before completing it. Resume to continue from where you left off.
                                </p>
                            </div>
                        </div>
                    )}

                    <div className="flex items-center justify-center gap-6 sm:gap-10">
                        <div className="space-y-0.5">
                            <p className="text-2xl font-bold sm:text-3xl">{total}</p>
                            <p className="text-xs text-muted-foreground sm:text-sm">Question{total === 1 ? "" : "s"}</p>
                        </div>
                        <div className="h-10 w-px bg-border" />
                        <div className="space-y-0.5">
                            <p className="text-2xl font-bold sm:text-3xl">{quiz.passing_score}%</p>
                            <p className="text-xs text-muted-foreground sm:text-sm">To pass</p>
                        </div>
                    </div>

                    <Button
                        size="lg"
                        className="w-full sm:w-auto"
                        onClick={handleStart}
                        disabled={!canAttempt || starting}
                    >
                        {starting ? "Starting…" : hasActiveSession ? "Resume Quiz" : attempts > 0 ? "Retake Quiz" : "Start Quiz"}
                    </Button>
                </div>
            </div>
        );
    }

    // ── RESULT ────────────────────────────────────────────────────────────────
    if (stage === "result" && result) {
        return (
            <div className="max-w-2xl mx-auto space-y-5">
                <div className={`rounded-xl border p-5 text-center space-y-1 ${result.passed ? "border-green-500/30 bg-green-500/5" : "border-red-500/30 bg-red-500/5"}`}>
                    <p className="text-sm text-muted-foreground">{result.passed ? "You passed!" : "You did not pass"}</p>
                    <p className="text-3xl font-bold">{result.score}%</p>
                    <p className="text-xs text-muted-foreground">
                        {result.earned_points} / {result.total_points} points · passing score {result.passing_score}%
                    </p>
                    {result.attempt_number && (
                        <p className="text-xs text-muted-foreground">Attempt #{result.attempt_number}</p>
                    )}
                </div>

                {result.passed && onNextContent && nextLabel && (
                    <div
                        onClick={onNextContent}
                        className="flex items-center gap-3 bg-card border rounded-xl px-4 py-3 shadow-sm hover:bg-muted transition-colors text-sm font-medium cursor-pointer select-none"
                    >
                        <div className="flex flex-col items-start flex-1">
                            <span className="text-xs text-muted-foreground font-normal">Up next</span>
                            <span>{nextLabel}</span>
                        </div>
                        <ArrowRight className="size-4 text-muted-foreground shrink-0" />
                    </div>
                )}

                {!result.passed && (
                    <div className="flex justify-center">
                        <Button variant="outline" onClick={handleRetake}>Retake {label}</Button>
                    </div>
                )}
            </div>
        );
    }

    // ── ASSESSMENT REVIEW ─────────────────────────────────────────────────────
    if (stage === "review" && isAssessment) {
        const answeredCount   = questions.filter(q => isQuestionAnswered(answers[q.question_id])).length;
        const unansweredCount = total - answeredCount;

        return (
            <div className="max-w-2xl mx-auto space-y-5">
                <div className="space-y-1">
                    <p className="text-xs font-semibold uppercase tracking-widest text-muted-foreground">Final Step</p>
                    <h2 className="text-2xl font-bold">Review your answers</h2>
                    <p className="text-sm text-muted-foreground">
                        Once you submit, your answers are final and your score is calculated.
                    </p>
                </div>

                {/* Stats row */}
                <div className="grid grid-cols-3 gap-3">
                    <div className="rounded-xl border bg-card p-4 space-y-0.5">
                        <p className="text-2xl font-bold text-green-600 dark:text-green-400">{answeredCount}</p>
                        <p className="text-xs text-muted-foreground">Answered</p>
                    </div>
                    <div className="rounded-xl border bg-card p-4 space-y-0.5">
                        <p className={`text-2xl font-bold ${unansweredCount > 0 ? "text-amber-600 dark:text-amber-400" : "text-muted-foreground"}`}>
                            {unansweredCount}
                        </p>
                        <p className="text-xs text-muted-foreground">Unanswered</p>
                    </div>
                    <div className="rounded-xl border bg-card p-4 space-y-0.5">
                        <p className="text-2xl font-bold font-mono tabular-nums">
                            {remainingSeconds !== null ? formatTime(remainingSeconds) : "—"}
                        </p>
                        <p className="text-xs text-muted-foreground">Time left</p>
                    </div>
                </div>

                {/* Unanswered warning */}
                {unansweredCount > 0 && (
                    <div className="flex items-start gap-2.5 rounded-lg border border-amber-400/50 bg-amber-50 dark:bg-amber-500/10 px-4 py-3 text-sm text-amber-700 dark:text-amber-400">
                        <AlertTriangle className="size-4 shrink-0 mt-0.5" />
                        <span>
                            <strong>{unansweredCount} question{unansweredCount !== 1 ? "s are" : " is"} still unanswered.</strong>{" "}
                            You can still submit. Tap a highlighted number below to jump back.
                        </span>
                    </div>
                )}

                {/* All questions grid */}
                <div className="rounded-xl border bg-card p-5 space-y-4">
                    <p className="text-sm font-semibold">All questions</p>
                    <div className="grid xs:grid-cols-8 lg:grid-cols-12 gap-2">
                        {questions.map((q, i) => {
                            const answered = isQuestionAnswered(answers[q.question_id]);
                            return (
                                <button
                                    key={q.question_id}
                                    type="button"
                                    onClick={() => { setCurrentIndex(i); setStage("taking"); }}
                                    className={`aspect-square rounded-lg text-sm font-medium transition-colors
                                        ${answered
                                            ? "bg-primary text-background hover:opacity-80"
                                            : "border-2 border-amber-400 text-amber-600 dark:text-amber-400 hover:bg-amber-50 dark:hover:bg-amber-500/10"
                                        }`}
                                >
                                    {i + 1}
                                </button>
                            );
                        })}
                    </div>
                    <div className="flex gap-4 text-xs text-muted-foreground">
                        <span className="flex items-center gap-1.5">
                            <span className="size-3 rounded-sm bg-primary inline-block" /> Answered
                        </span>
                        <span className="flex items-center gap-1.5">
                            <span className="size-3 rounded-sm border-2 border-amber-400 inline-block" /> Unanswered
                        </span>
                    </div>
                </div>

                <div className="space-y-2">
                    <div className="flex items-center justify-between">
                        <Button variant="outline" onClick={() => setStage("taking")} disabled={submitting}>
                            <ChevronLeft className="size-4" />
                            Back to assessment
                        </Button>
                        <Button onClick={handleSubmit} disabled={submitting}>
                            {submitting ? "Submitting…" : "Submit assessment"}
                        </Button>
                    </div>
                    {submitError && (
                        <p className="text-xs text-red-600 dark:text-red-400 text-right flex items-center justify-end gap-1">
                            <AlertTriangle className="size-3 shrink-0" /> {submitError}
                        </p>
                    )}
                </div>
            </div>
        );
    }

    // ── QUIZ REVIEW ───────────────────────────────────────────────────────────
    if (stage === "review") {
        return (
            <div className="max-w-2xl mx-auto space-y-5">
                <div>
                    <h2 className="text-xl font-bold sm:text-2xl pb-3 border-b border-border">
                        {quiz.title || label}
                    </h2>
                </div>

                {/* Question pills */}
                <div className="flex flex-wrap gap-2">
                    {questions.map((q, i) => {
                        const answered = isQuestionAnswered(answers[q.question_id]);
                        return (
                            <button
                                key={q.question_id}
                                type="button"
                                onClick={() => { setCurrentIndex(i); setStage("taking"); }}
                                className={`size-9 rounded-full text-sm font-semibold transition-colors
                                    ${answered
                                        ? "bg-primary text-background hover:opacity-80"
                                        : "bg-muted text-muted-foreground border border-border hover:bg-muted/80"
                                    }`}
                            >
                                {i + 1}
                            </button>
                        );
                    })}
                </div>

                <div className="space-y-1">
                    <h3 className="text-lg font-bold">Review your answers</h3>
                    <p className="text-sm text-muted-foreground">
                        Check your responses before submitting. Tap any question to go back and change your answer.
                    </p>
                </div>

                <div className="space-y-3">
                    {questions.map((q, i) => {
                        const ans      = answers[q.question_id];
                        const answered = isQuestionAnswered(ans);
                        const isMultiQ = q.type === "multi_select";

                        let answerLabel = null;
                        if (answered) {
                            if (isMultiQ && Array.isArray(ans)) {
                                answerLabel = ans
                                    .map(id => {
                                        const idx = (q.options ?? []).findIndex(o => o.option_id === id);
                                        return idx >= 0 ? String.fromCharCode(65 + idx) : null;
                                    })
                                    .filter(Boolean)
                                    .join(", ");
                            } else {
                                const idx = (q.options ?? []).findIndex(o => o.option_id === ans);
                                if (idx >= 0) {
                                    answerLabel = `${String.fromCharCode(65 + idx)}. ${q.options[idx].text}`;
                                }
                            }
                        }

                        return (
                            <button
                                key={q.question_id}
                                type="button"
                                onClick={() => { setCurrentIndex(i); setStage("taking"); }}
                                className="w-full flex items-center gap-4 rounded-xl border bg-card p-4 text-left hover:bg-muted/30 transition-colors"
                            >
                                <span className="shrink-0 size-8 rounded-full bg-primary text-background text-sm font-semibold flex items-center justify-center">
                                    {i + 1}
                                </span>
                                <div className="flex-1 min-w-0">
                                    <p className="text-sm font-semibold leading-snug line-clamp-2">{q.question}</p>
                                    {answerLabel && (
                                        <p className="text-xs text-muted-foreground mt-0.5 truncate">{answerLabel}</p>
                                    )}
                                </div>
                                <span className={`shrink-0 rounded-full px-2.5 py-0.5 text-xs font-medium ${
                                    answered
                                        ? "bg-green-500/10 text-green-700 dark:text-green-400"
                                        : "bg-muted text-muted-foreground"
                                }`}>
                                    {answered ? "Answered" : "Not answered"}
                                </span>
                            </button>
                        );
                    })}
                </div>

                <div className="space-y-2">
                    <div className="flex items-center justify-between">
                        <Button variant="outline" onClick={() => setStage("taking")} disabled={submitting}>
                            <ChevronLeft className="size-4" />
                            Back to questions
                        </Button>
                        <Button onClick={handleSubmit} disabled={submitting}>
                            {submitting ? "Submitting…" : "Submit quiz"}
                        </Button>
                    </div>
                    {submitError && (
                        <p className="text-xs text-red-600 dark:text-red-400 text-right flex items-center justify-end gap-1">
                            <AlertTriangle className="size-3 shrink-0" /> {submitError}
                        </p>
                    )}
                </div>
            </div>
        );
    }

    // ── QUESTION STEPPER (shared derived values) ──────────────────────────────
    const question   = questions[currentIndex];
    const isFirst    = currentIndex === 0;
    const isLast     = currentIndex === total - 1;
    const isMulti    = question.type === "multi_select";
    const selected   = answers[question.question_id];
    const progress   = Math.round(((currentIndex + 1) / total) * 100);

    const multiLimit         = isMulti ? (question.correct_count ?? null) : null;
    const selectedCount      = isMulti ? (selected ?? []).length : 0;
    const limitReached       = multiLimit !== null && selectedCount >= multiLimit;
    const isCurrentAnswered  = isQuestionAnswered(selected);
    const unansweredCount    = questions.filter(q => !isQuestionAnswered(answers[q.question_id])).length;

    const timerColor = remainingSeconds !== null
        ? remainingSeconds <= 60
            ? "text-red-500 dark:text-red-400"
            : remainingSeconds <= 300
            ? "text-amber-500 dark:text-amber-400"
            : "text-muted-foreground"
        : null;

    // ── ASSESSMENT TAKING (2-column layout) ───────────────────────────────────
    if (isAssessment) {
        const answeredCount = questions.filter(q => isQuestionAnswered(answers[q.question_id])).length;

        const optionRows = (question.options ?? []).map((option, i) => {
            const letter     = String.fromCharCode(65 + i);
            const isSelected = isMulti
                ? (selected ?? []).includes(option.option_id)
                : selected === option.option_id;
            const isDisabled = (isMulti && limitReached && !isSelected) || (timeExpired && submitting);
            const Icon = isMulti
                ? (isSelected ? CheckSquare2 : Square)
                : (isSelected ? CheckCircle2 : Circle);

            return (
                <button
                    key={option.option_id}
                    type="button"
                    onClick={() => handleOptionClick(option.option_id)}
                    disabled={isDisabled}
                    className={`flex w-full items-center gap-3 rounded-lg border px-3 py-3 text-left text-sm transition-colors sm:text-base
                        ${isSelected ? "border-primary bg-primary/5" : "border-border"}
                        ${isDisabled ? "opacity-40 cursor-not-allowed" : "hover:bg-muted-foreground/5"}`}
                >
                    <Icon className={`size-4 shrink-0 ${isSelected ? "text-primary" : "text-muted-foreground"}`} />
                    <span className="font-medium text-muted-foreground w-4 shrink-0">{letter}</span>
                    <span>{option.text}</span>
                </button>
            );
        });

        return (
            <div className="max-w-5xl mx-auto">
                {/* Assessment updated alert */}
                {assessmentUpdatedAlert && (
                    <div className="flex items-start gap-3 rounded-lg border border-blue-500/30 bg-blue-500/5 px-4 py-3 mb-5">
                        <Info className="h-4 w-4 text-blue-500 shrink-0 mt-0.5" />
                        <div className="flex-1 min-w-0">
                            <p className="text-sm font-medium text-blue-700 dark:text-blue-400">Assessment Updated</p>
                            <p className="text-xs text-blue-600/80 dark:text-blue-300/70 mt-0.5">
                                Your administrator has made changes to this assessment. Your current session and saved answers are unaffected.
                            </p>
                        </div>
                        <button
                            onClick={() => setAssessmentUpdatedAlert(false)}
                            className="text-blue-400 hover:text-blue-600 shrink-0 text-xs leading-none mt-0.5"
                        >✕</button>
                    </div>
                )}

                <div className="grid grid-cols-1 lg:grid-cols-[1fr_300px] gap-4 items-start">
                    {/* ── Right: sidebar — order-first on mobile so timer/progress sit above the question ── */}
                    <div className="space-y-2 order-first lg:order-last">
                        {/* Timer */}
                        {remainingSeconds !== null && (
                            <div className={`rounded-xl flex items-center justify-between gap-3 px-4 py-3 lg:block lg:p-4 ${
                                timeExpired
                                    ? "bg-red-600 text-white"
                                    : remainingSeconds <= 60
                                    ? "bg-red-700 text-white"
                                    : "bg-primary text-background"
                            }`}>
                                <p className="text-xs uppercase tracking-widest opacity-60">Time Remaining</p>
                                <p className="text-2xl lg:text-4xl font-bold font-mono tabular-nums lg:mt-1">
                                    {timeExpired ? "00:00" : formatTime(remainingSeconds)}
                                </p>
                            </div>
                        )}

                        {/* Progress panel */}
                        <div className="rounded-xl border bg-card p-3 lg:p-4 space-y-2 lg:space-y-3">
                            <div className="flex items-center justify-between text-sm">
                                <span className="font-medium">Progress</span>
                                <span className="text-muted-foreground">{answeredCount} / {total} answered</span>
                            </div>

                            {/* Number grid — scrollable after ~30 questions (3 rows on mobile, 4 on desktop) */}
                            <div className="overflow-y-auto max-h-24 lg:max-h-40">
                                <div className="grid grid-cols-10 lg:grid-cols-8 gap-1 lg:gap-1.5">
                                    {questions.map((q, i) => {
                                        const ans      = answers[q.question_id];
                                        const answered  = isQuestionAnswered(ans);
                                        const isCurrent = i === currentIndex;
                                        return (
                                            <button
                                                key={q.question_id}
                                                type="button"
                                                onClick={() => setCurrentIndex(i)}
                                                disabled={timeExpired && submitting}
                                                className={`aspect-square rounded text-xs font-medium transition-colors
                                                    ${isCurrent
                                                        ? "border-2 border-foreground bg-background text-foreground"
                                                        : answered
                                                        ? "bg-primary text-background hover:opacity-80"
                                                        : "border border-border bg-background text-muted-foreground hover:bg-muted"
                                                    }`}
                                            >
                                                {i + 1}
                                            </button>
                                        );
                                    })}
                                </div>
                            </div>

                            {/* Legend — desktop only */}
                            <div className="hidden lg:flex flex-col gap-1.5 text-xs text-muted-foreground">
                                <span className="flex items-center gap-1.5">
                                    <span className="size-3 rounded-sm bg-primary inline-block shrink-0" />
                                    Answered
                                </span>
                                <span className="flex items-center gap-1.5">
                                    <span className="size-3 rounded-sm border border-border inline-block shrink-0" />
                                    Not answered
                                </span>
                                <span className="flex items-center gap-1.5">
                                    <span className="size-3 rounded-sm border-2 border-foreground inline-block shrink-0" />
                                    Current
                                </span>
                            </div>

                            {/* Auto-save indicator */}
                            <div className="flex items-center gap-1.5 text-xs text-green-600 dark:text-green-400 rounded-lg bg-green-500/5 border border-green-500/20 px-3 py-2">
                                <CheckCircle2 className="size-3.5 shrink-0" />
                                All answers saved automatically
                            </div>

                            {/* Review button */}
                            <Button
                                className="w-full"
                                onClick={handleGoToReview}
                                disabled={submitting || unansweredCount > 0}
                            >
                                Review &amp; submit
                            </Button>
                            {unansweredCount > 0 && (
                                <p className="text-xs text-red-600 dark:text-red-400 text-center">
                                    Answer all {total} questions first — {unansweredCount} remaining.
                                </p>
                            )}
                        </div>
                    </div>

                    {/* ── Left: question — order-last on mobile so it appears below the compact sidebar ── */}
                    <div className="space-y-4 order-last lg:order-first">
                        {/* Question header */}
                        <div className="flex items-center gap-2.5 text-sm">
                            <span className="text-muted-foreground">Question {currentIndex + 1} of {total}</span>
                            <span className="rounded-full bg-muted px-2.5 py-0.5 text-xs font-medium text-muted-foreground">
                                {isMulti ? "Multiple select" : "Multiple choice"}
                            </span>
                        </div>

                        {/* Question card */}
                        <div className="rounded-xl border bg-card p-5 sm:p-6 space-y-4">
                            <div className="space-y-1.5">
                                <p className="font-bold text-base leading-relaxed sm:text-lg">
                                    {currentIndex + 1}. {question.question}
                                </p>
                                <p className="text-sm text-muted-foreground">
                                    {isMulti
                                        ? multiLimit
                                            ? `Select ${multiLimit} answers`
                                            : "Select all that apply"
                                        : "Select one answer"}
                                </p>
                                {isMulti && multiLimit !== null && (
                                    <p className={`text-xs font-medium ${limitReached ? "text-amber-600 dark:text-amber-400" : "text-blue-600 dark:text-blue-400"}`}>
                                        {limitReached
                                            ? `${selectedCount} / ${multiLimit} selected — limit reached`
                                            : `${selectedCount} / ${multiLimit} selected`}
                                    </p>
                                )}
                            </div>
                            <div className="space-y-2">{optionRows}</div>
                        </div>

                        {timeExpired && submitting && (
                            <div className="rounded-lg border border-amber-500/30 bg-amber-500/5 px-4 py-2 text-xs text-amber-700 dark:text-amber-400 text-center">
                                Time's up — submitting your answers…
                            </div>
                        )}

                        {/* Navigation */}
                        <div className="flex items-center justify-between">
                            <Button variant="outline" onClick={handlePrev} disabled={submitting || isFirst}>
                                <ChevronLeft className="size-4" />
                                Previous
                            </Button>
                            <Button onClick={handleNext} disabled={submitting || (isLast && unansweredCount > 0)}>
                                {submitting ? "Submitting…" : isLast ? "Review & submit" : "Next"}
                                {!submitting && <ChevronRight className="size-4" />}
                            </Button>
                        </div>
                    </div>
                </div>
            </div>
        );
    }

    // ── QUIZ TAKING (1-column with pills) ─────────────────────────────────────
    return (
        <div className="max-w-2xl mx-auto space-y-5">
            {/* Header: title + progress */}
            <div className="space-y-2">
                <div className="flex items-center justify-between">
                    {quiz.title && <h2 className="text-lg font-bold sm:text-xl truncate">{quiz.title}</h2>}
                </div>
                <div className="flex items-center justify-between text-xs text-muted-foreground sm:text-sm">
                    <span>Question {currentIndex + 1} of {total}</span>
                    <span>{progress}%</span>
                </div>
                <div className="h-1.5 w-full rounded-full bg-border overflow-hidden">
                    <div className="h-full bg-primary transition-all duration-300 ease-out" style={{ width: `${progress}%` }} />
                </div>
            </div>

            {/* Question number pills — all freely clickable; skipped questions show as muted */}
            <div className="flex flex-wrap gap-2">
                {questions.map((q, i) => {
                    const ans      = answers[q.question_id];
                    const answered  = isQuestionAnswered(ans);
                    const isCurrent = i === currentIndex;
                    return (
                        <button
                            key={q.question_id}
                            type="button"
                            onClick={() => setCurrentIndex(i)}
                            disabled={timeExpired && submitting}
                            className={`size-9 rounded-full text-sm font-semibold transition-colors
                                ${isCurrent
                                    ? "border-2 border-primary bg-background text-foreground"
                                    : answered
                                    ? "bg-primary text-background hover:opacity-80"
                                    : "bg-muted text-muted-foreground border border-border hover:bg-muted/60"
                                }`}
                        >
                            {i + 1}
                        </button>
                    );
                })}
            </div>

            {timeExpired && submitting && (
                <div className="rounded-lg border border-amber-500/30 bg-amber-500/5 px-4 py-2 text-xs text-amber-700 dark:text-amber-400 text-center">
                    Time's up — submitting your answers…
                </div>
            )}

            {/* Question card */}
            <div className="rounded-xl border bg-card p-4 space-y-4 sm:p-6">
                <div className="space-y-1">
                    <p className="font-bold text-base leading-relaxed sm:text-lg">
                        {currentIndex + 1}. {question.question}
                    </p>
                    {isMulti && multiLimit !== null && (
                        <p className={`text-xs font-medium ${limitReached ? "text-amber-600 dark:text-amber-400" : "text-blue-600 dark:text-blue-400"}`}>
                            {limitReached
                                ? `${selectedCount} / ${multiLimit} selected — limit reached`
                                : `Select ${multiLimit} answer${multiLimit !== 1 ? "s" : ""} · ${selectedCount} / ${multiLimit} selected`}
                        </p>
                    )}
                </div>
                <div className="space-y-2">
                    {(question.options ?? []).map((option, i) => {
                        const letter     = String.fromCharCode(65 + i);
                        const isSelected = isMulti
                            ? (selected ?? []).includes(option.option_id)
                            : selected === option.option_id;
                        const isDisabled = (isMulti && limitReached && !isSelected) || (timeExpired && submitting);
                        const Icon = isMulti
                            ? (isSelected ? CheckSquare2 : Square)
                            : (isSelected ? CheckCircle2 : Circle);

                        return (
                            <button
                                key={option.option_id}
                                type="button"
                                onClick={() => handleOptionClick(option.option_id)}
                                disabled={isDisabled}
                                className={`flex w-full items-center gap-3 rounded-lg border px-3 py-2.5 text-left text-sm transition-colors sm:text-base
                                    ${isSelected ? "border-primary bg-primary/5" : "border-border"}
                                    ${isDisabled ? "opacity-40 cursor-not-allowed" : "hover:bg-muted-foreground/5"}`}
                            >
                                <Icon className={`size-4 shrink-0 ${isSelected ? "text-primary" : "text-muted-foreground"}`} />
                                <span className="font-medium text-muted-foreground w-4 shrink-0">{letter}</span>
                                <span>{option.text}</span>
                            </button>
                        );
                    })}
                </div>
            </div>

            {/* Navigation */}
            <div className="space-y-2">
                <div className="flex items-center justify-between">
                    <Button variant="outline" onClick={handlePrev} disabled={submitting}>
                        <ChevronLeft className="size-4" />
                        Previous
                    </Button>
                    <Button onClick={handleNext} disabled={submitting || (isLast && unansweredCount > 0)}>
                        {submitting ? "Submitting…" : isLast ? "Review answers" : "Next"}
                        {!submitting && <ChevronRight className="size-4" />}
                    </Button>
                </div>
                {isLast && unansweredCount > 0 && (
                    <p className="text-xs text-red-600 dark:text-red-400 text-center">
                        Answer all {total} questions first — {unansweredCount} remaining.
                    </p>
                )}
            </div>
        </div>
    );
};

export default QuizBlock;