"use client"

import { useEffect, useState, use, useRef } from "react"
import { Button } from "@/components/ui/button"
import { Badge } from "@/components/ui/badge"
import { Clock, Flag, ChevronLeft, ChevronRight, Pause as PauseIcon, MoreVertical, Calculator as CalculatorIcon, FlaskConical, Edit3, XCircle, Play, Shield, AlertTriangle, Moon, Sun } from "lucide-react"
import Image from "next/image"
import { useCBTStore, QuestionStatus } from "@/store/cbtStore"
import { useExamTimer } from "@/hooks/useExamTimer"
import ReactMarkdown from "react-markdown"
import remarkMath from "remark-math"
import rehypeKatex from "rehype-katex"
import rehypeRaw from "rehype-raw"
import "katex/dist/katex.min.css"
import { Calculator } from "@/components/CBT/Calculator"
import { Scratchpad } from "@/components/CBT/Scratchpad"
import { PeriodicTable } from "@/components/CBT/PeriodicTable"
import { ExamResult } from "@/components/CBT/ExamResult"

// Dummy data for testing the CBT client engine
const DUMMY_QUESTIONS = Array.from({ length: 20 }).map((_, i) => ({
  id: `q_${i+1}`,
  number: i + 1,
  subject: "Physics",
  text_en: `A solid cylinder of mass M and radius R is rolling without slipping down an inclined plane of angle θ. Question ${i+1}. What is the minimum coefficient of static friction μ_s required to prevent slipping?`,
  text_hi: `द्रव्यमान M और त्रिज्या R का एक ठोस बेलन कोण θ के एक आनत तल पर बिना फिसले लुढ़क रहा है। प्रश्न ${i+1}...`,
  options: ['(1/3) tan θ', '(2/3) tan θ', '(1/2) tan θ', 'tan θ']
}));

export default function CBTExamPage({ params }: { params: Promise<{ examId: string }> }) {
  const resolvedParams = use(params);
  const examId = resolvedParams.examId;
  const [mounted, setMounted] = useState(false);
  const [examType, setExamType] = useState<string>('Test');
  const [showCalc, setShowCalc] = useState(false);
  const [showScratchpad, setShowScratchpad] = useState(false);
  const [showPeriodicTable, setShowPeriodicTable] = useState(false);
  const [isPaused, setIsPaused] = useState(false);
  const [showSubmitConfirm, setShowSubmitConfirm] = useState(false);
  const [accessError, setAccessError] = useState<string | null>(null);
  const [localTheme, setLocalTheme] = useState<string>('light');
  
  useEffect(() => {
    if (typeof window !== 'undefined') {
      if (document.documentElement.classList.contains('dark')) {
        setLocalTheme('dark');
      } else {
        setLocalTheme('light');
      }
    }
  }, []);

  const handleThemeToggle = () => {
    if (document.documentElement.classList.contains('dark')) {
      document.documentElement.classList.remove('dark');
      localStorage.setItem('theme', 'light');
      setLocalTheme('light');
    } else {
      document.documentElement.classList.add('dark');
      localStorage.setItem('theme', 'dark');
      setLocalTheme('dark');
    }
  };
  
  // Anti-cheat state
  const [violationCount, setViolationCount] = useState(0);
  const [showWarning, setShowWarning] = useState(false);
  const hasShownWarning = useRef(false);
  const hasAutoSubmitted = useRef(false);
  
  const { 
    initExam, 
    startExam,
    questions, 
    answers, 
    currentQuestionIndex, 
    selectAnswer, 
    clearResponse, 
    markForReview, 
    goToQuestion,
    addTimeSpent,
    submitExam,
    getAttemptPayload,
    submitAttemptToServer,
    resetExamState,
    status,
    attemptId,
  } = useCBTStore();

  const { formattedTime, isTimedOut } = useExamTimer();
  const [questionStartTime, setQuestionStartTime] = useState<number>(Date.now());

  // Initialize Exam on mount
  useEffect(() => {
    async function fetchExamData() {
      try {
        const apiUrl = process.env.NEXT_PUBLIC_API_URL || 'http://localhost:8000';
        const testRes = await fetch(`${apiUrl}/api/tests/${examId}?user_id=b0e17378-864d-439e-8806-708dc48a891b`);
        if (!testRes.ok) throw new Error("Failed to fetch test structure");
        const testData = await testRes.json();
        
        // Strict boundary checking
        const now = new Date();
        if (testData.scheduled_date && new Date(testData.scheduled_date) > now) {
          setAccessError(`This exam is scheduled for ${new Date(testData.scheduled_date).toLocaleString()}. You cannot start it yet.`);
          setMounted(true);
          return;
        }
        if (testData.end_date && new Date(testData.end_date) < now) {
          setAccessError(`This exam expired on ${new Date(testData.end_date).toLocaleString()} and is no longer available.`);
          setMounted(true);
          return;
        }
        
        setExamType(testData.type);
        
        // Extract all question IDs from sections
        const questionIds: string[] = [];
        Object.values(testData.sections).forEach((section: any) => {
          questionIds.push(...section.questions);
        });

        const qRes = await fetch(`${apiUrl}/api/tests/${examId}/questions`);
        if (!qRes.ok) throw new Error("Failed to fetch questions in bulk");
        const allQuestionsData = await qRes.json();
        
        // Map questions by ID for fast lookup
        const questionsMap = new Map();
        allQuestionsData.forEach((q: any) => {
          questionsMap.set(q.id, q);
        });

        // Construct fetchedQuestions retaining the exact order of questionIds
        const fetchedQuestions = questionIds.map((id, idx) => {
          const qData = questionsMap.get(id);
          if (!qData) return null; // Fallback in case of mismatch
          return {
            id: qData.id,
            number: idx + 1,
            subject: qData.subject,
            type: qData.type,
            text_en: qData.question_text,
            text_hi: qData.question_text_hi || qData.question_text,
            options: Object.values(qData.options || {}),
            options_hi: Object.values(qData.options_hi || qData.options || {}),
          };
        }).filter(Boolean) as any[];

        initExam(examId, fetchedQuestions, testData.total_time * 60);
        setMounted(true);
      } catch (err) {
        console.error("Failed to load exam from network, falling back to dummy data:", err);
        initExam(examId, DUMMY_QUESTIONS, 10800);
        setMounted(true);
      }
    }
    fetchExamData();
  }, [examId, initExam]);

  // Track time spent per question
  useEffect(() => {
    if (status === 'IN_PROGRESS') {
      setQuestionStartTime(Date.now());
    }
  }, [currentQuestionIndex, status]);

  // Anti-Cheat: Tab switching listener
  useEffect(() => {
    if (status !== 'IN_PROGRESS') return;

    const handleVisibilityChange = () => {
      if (document.visibilityState === 'hidden') {
        setViolationCount((prev) => prev + 1);
      }
    };

    document.addEventListener("visibilitychange", handleVisibilityChange);
    return () => document.removeEventListener("visibilitychange", handleVisibilityChange);
  }, [status]);

  // Handle side-effects of violation count changing
  useEffect(() => {
    if (violationCount === 1 && !hasShownWarning.current) {
      hasShownWarning.current = true;
      setShowWarning(true);
    } else if (violationCount >= 2 && status !== 'SUBMITTED' && status !== 'TIMED_OUT' && !hasAutoSubmitted.current) {
      hasAutoSubmitted.current = true;
      submitExam('SUBMITTED');
      submitAttemptToServer('b0e17378-864d-439e-8806-708dc48a891b')
        .then(() => {
          window.alert("Test Auto-Submitted due to multiple tab switching violations.");
        })
        .catch(err => {
          console.error("Auto-submit fetch error:", err);
          window.alert("Test Auto-Submitted due to multiple tab switching violations.");
        });
    }
  }, [violationCount, status, submitExam, submitAttemptToServer]);

  const startTest = async () => {
    try {
      if (document.documentElement.requestFullscreen) {
        await document.documentElement.requestFullscreen();
      }
    } catch (e) {
      console.warn("Fullscreen API failed", e);
    }
    startExam();
  };

  const handleNext = () => {
    if (!currentQuestion) return;
    const timeSpent = Math.floor((Date.now() - questionStartTime) / 1000);
    addTimeSpent(currentQuestion.id, timeSpent);
    if (currentQuestionIndex < questions.length - 1) {
      goToQuestion(currentQuestionIndex + 1);
    } else {
      goToQuestion(0);
    }
  }

  const handlePrev = () => {
    if (!currentQuestion) return;
    const timeSpent = Math.floor((Date.now() - questionStartTime) / 1000);
    addTimeSpent(currentQuestion.id, timeSpent);
    if (currentQuestionIndex > 0) {
      goToQuestion(currentQuestionIndex - 1);
    }
  }

  const handleGridClick = (index: number) => {
    if (!currentQuestion) return;
    const timeSpent = Math.floor((Date.now() - questionStartTime) / 1000);
    addTimeSpent(currentQuestion.id, timeSpent);
    goToQuestion(index);
  }

  const handleSubmit = () => {
    setShowSubmitConfirm(true);
  }

  const confirmSubmit = async () => {
    setShowSubmitConfirm(false);
    submitExam('SUBMITTED');
    const payload = getAttemptPayload('b0e17378-864d-439e-8806-708dc48a891b');
    console.log("EXAM SUBMITTED. PAYLOAD:", JSON.stringify(payload, null, 2));
    await submitAttemptToServer('b0e17378-864d-439e-8806-708dc48a891b');
  }

  if (!mounted || questions.length === 0) {
    return <div className="min-h-screen flex items-center justify-center bg-background"><div className="animate-pulse flex items-center gap-2"><Clock className="animate-spin text-muted-foreground" size={20}/><span className="font-bold text-muted-foreground tracking-widest uppercase text-sm">Initializing CBT Engine</span></div></div>
  }

  if (accessError) {
    return (
      <div className="min-h-screen flex items-center justify-center bg-background p-4">
        <div className="max-w-md w-full bg-secondary/20 p-8 rounded-2xl border border-border text-center space-y-4">
          <div className="mx-auto size-12 bg-red-500/10 text-red-500 rounded-full flex items-center justify-center">
            <AlertTriangle size={24} />
          </div>
          <h2 className="text-xl font-bold">Access Denied</h2>
          <p className="text-muted-foreground font-medium">{accessError}</p>
          <Button variant="outline" className="mt-4" onClick={() => window.history.back()}>
            Go Back
          </Button>
        </div>
      </div>
    );
  }

  if (status === 'NOT_STARTED') {
    return (
      <div className="h-screen flex flex-col items-center justify-center bg-background p-8 text-center animate-in fade-in duration-500">
        <div className="max-w-2xl bg-secondary/10 border border-border p-8 rounded-2xl shadow-2xl">
          <h1 className="text-3xl font-bold mb-4 flex items-center justify-center gap-3">
            <Shield className="text-red-500" /> Secure Exam Environment
          </h1>
          <div className="text-left space-y-4 text-muted-foreground mb-8">
            <p><strong>Target:</strong> {questions.some(q => q.subject.toLowerCase() === 'biology') ? 'NEET' : 'JEE'} Pattern Test</p>
            <ul className="list-disc pl-5 space-y-2">
              <li>The test will open in <strong>Full Screen Mode</strong>.</li>
              <li><strong>Do not switch tabs or exit full screen.</strong></li>
              <li>First violation will result in a strict warning.</li>
              <li>Second violation will <strong>automatically submit</strong> your exam to prevent cheating.</li>
              <li>Ensure your internet connection is stable before starting.</li>
            </ul>
          </div>
          
          <div className="mb-8 p-4 bg-yellow-500/10 border border-yellow-500/30 rounded-xl text-yellow-600 dark:text-yellow-500 text-sm text-left">
            <strong>Disclaimer:</strong> Only the score obtained in your <em>first attempt</em> will be considered for growth tracking and performance evaluation. Any subsequent attempts are intended solely for practice and will not be included in the official assessment.
          </div>
          
          <div className="mb-8 p-4 bg-muted rounded-xl flex items-center justify-between">
            <div>
              <p className="font-semibold">Preferred Theme</p>
              <p className="text-sm text-muted-foreground">Choose dark or light mode for the exam</p>
            </div>
            <Button variant="outline" size="icon" onClick={handleThemeToggle}>
              {localTheme === 'dark' ? <Sun className="h-5 w-5" /> : <Moon className="h-5 w-5" />}
            </Button>
          </div>

          <Button size="lg" className="w-full text-lg font-bold h-14 shadow-[0_0_20px_-5px_var(--color-primary)]" onClick={startTest}>
            I Understand, Start Test
          </Button>
        </div>
      </div>
    );
  }

  if (status === 'SUBMITTED' || status === 'TIMED_OUT') {
    const payload = getAttemptPayload('b0e17378-864d-439e-8806-708dc48a891b');
    if (!payload) return <div className="p-8 text-center mt-20">Loading results...</div>;
    
    return (
      <ExamResult 
        payload={payload} 
        attemptId={attemptId}
        examType={examType}
        onRetake={examType === 'DPS' ? () => {
          resetExamState();
          window.location.reload();
        } : () => {}}
      />
    )
  }

  const currentQuestion = questions[currentQuestionIndex];
  const currentState = currentQuestion ? answers[currentQuestion.id] : null

  // Helper to ensure math options missing $ delimiters get rendered properly
  const renderLatex = (text: string) => {
    if (typeof text !== 'string') return text;
    // If it contains backslash but no dollar signs, we might want to wrap the whole thing in $ $
    if (text.includes('\\') && !text.includes('$')) {
      return `$${text}$`;
    }
    return text;
  };

  const handleOptionSelect = (qId: string, letter: string, type: string) => {
    if (type === 'MCQ_Multiple' || type === 'Multiple_Choice') {
      const currentSelected = currentState?.selected_answer || "";
      const selectedArray = currentSelected.split(',').filter(Boolean);
      if (selectedArray.includes(letter)) {
        const newArray = selectedArray.filter((l: string) => l !== letter).sort();
        if (newArray.length > 0) {
          selectAnswer(qId, newArray.join(','));
        } else {
          clearResponse(qId);
        }
      } else {
        selectedArray.push(letter);
        selectAnswer(qId, selectedArray.sort().join(','));
      }
    } else {
      selectAnswer(qId, letter);
    }
  }

  // Effect to close overlay tools on Escape
  const paletteStats = Object.values(answers || {}).reduce(
    (acc, val) => {
      acc[val.status]++;
      return acc;
    },
    { UNVISITED: 0, NOT_ANSWERED: 0, ANSWERED: 0, MARKED: 0, ANSWERED_MARKED: 0 } as Record<QuestionStatus, number>
  );

  // Group questions by subject for the side panel
  const questionsBySubject = questions.reduce((acc, q, index) => {
    if (!acc[q.subject]) {
      acc[q.subject] = [];
    }
    acc[q.subject].push({ ...q, originalIndex: index });
    return acc;
  }, {} as Record<string, any[]>);

  return (
    <div className="h-screen flex flex-col overflow-hidden">
      {/* Warning Modal */}
      {showWarning && (
        <div className="fixed inset-0 z-[100] flex items-center justify-center bg-black/80 backdrop-blur-sm animate-in fade-in duration-200">
          <div className="bg-background border border-red-500 shadow-[0_0_50px_-10px_rgba(239,68,68,0.5)] rounded-2xl w-[400px] overflow-hidden animate-in zoom-in-95 duration-200">
            <div className="p-6 text-center space-y-4">
              <AlertTriangle size={64} className="text-red-500 mx-auto animate-pulse" />
              <h3 className="text-2xl font-bold text-red-500">Warning!</h3>
              <p className="text-foreground font-medium">You switched tabs or left the secure window.</p>
              <p className="text-sm text-muted-foreground">This is your final warning. Doing this again will automatically submit your exam and your score will be evaluated as is.</p>
              <Button 
                className="w-full mt-4 bg-red-500 hover:bg-red-600 text-white font-bold" 
                onClick={() => {
                  setShowWarning(false);
                  startTest(); // Try to re-enter fullscreen
                }}
              >
                Return to Exam
              </Button>
            </div>
          </div>
        </div>
      )}

      <header className="border-b border-border bg-background px-6 py-4 flex items-center justify-between sticky top-0 z-20 shadow-sm shrink-0">
        <div className="flex items-center gap-4">
          <div className="relative h-10 w-40 shrink-0">
            <Image 
              src="/PragatiVeer_Logo/SVG_Logos/Rectangular_light_background.svg" 
              alt="PragatiVeer" fill className="object-contain object-left scale-125 origin-left dark:hidden" priority
            />
            <Image 
              src="/PragatiVeer_Logo/SVG_Logos/Rectangular_dark_background.svg" 
              alt="PragatiVeer" fill className="object-contain object-left scale-125 origin-left hidden dark:block" priority
            />
          </div>
          <div className="h-6 w-px bg-border hidden md:block"></div>
          <div className="hidden md:block">
            <h1 className="font-semibold text-sm">AITS Full Syllabus Test</h1>
            <p className="text-xs text-muted-foreground font-medium mt-0.5">{currentQuestion.subject}</p>
          </div>
        </div>
        
        <div className="flex items-center gap-3 md:gap-6">
          <div className="hidden lg:flex items-center gap-2 mr-4">
            <Button variant="ghost" size="sm" onClick={() => setShowCalc(true)} className="gap-2 text-xs font-semibold text-muted-foreground hover:bg-[#00F0FF] hover:text-black transition-colors"><CalculatorIcon size={14}/> Calculator</Button>
            <Button variant="ghost" size="sm" onClick={() => setShowPeriodicTable(true)} className="gap-2 text-xs font-semibold text-muted-foreground hover:bg-[#00F0FF] hover:text-black transition-colors"><FlaskConical size={14}/> Periodic Table</Button>
            <Button variant="ghost" size="sm" onClick={() => setShowScratchpad(true)} className="gap-2 text-xs font-semibold text-muted-foreground hover:bg-[#00F0FF] hover:text-black transition-colors"><Edit3 size={14}/> Scratchpad</Button>
          </div>
          
          <div className={`flex items-center gap-2 px-3 py-1.5 md:px-4 md:py-2 rounded-xl font-mono text-base md:text-lg tracking-wider font-medium border border-border ${isTimedOut ? 'bg-red-500/10 text-red-500' : 'bg-secondary/50 text-foreground'}`}>
            <Clock size={18} className="text-muted-foreground" />
            <span className={isTimedOut ? "font-bold animate-pulse text-red-500" : "font-bold text-primary"}>{formattedTime}</span>
          </div>
          <Button variant="outline" size="sm" onClick={() => setIsPaused(!isPaused)} className="gap-2 rounded-xl hidden md:flex border-border">
            {isPaused ? <><Play size={14} /> Resume</> : <><PauseIcon size={14} /> Pause</>}
          </Button>
          <Button size="sm" onClick={handleSubmit} className="rounded-xl shadow-[0_0_15px_-3px_var(--color-primary)]">Submit Exam</Button>
        </div>
      </header>

      <div className="flex flex-1 overflow-hidden relative">
        {/* Main Question Area - Configured for Bilingual Split Pane */}
        <main className="flex-1 overflow-y-auto p-0 flex flex-col xl:flex-row pb-24 lg:pb-0">
          <div className="flex-1 p-6 lg:p-10 border-b xl:border-b-0 xl:border-r border-border flex flex-col">
            <div className="max-w-2xl mx-auto w-full space-y-8 flex-1">
              <div className="flex items-center justify-between border-b border-border pb-4">
                <div className="flex items-center gap-2">
                  <Badge variant="secondary" className="px-3 py-1 text-sm bg-secondary text-secondary-foreground font-medium rounded-lg">Question {currentQuestion.number}</Badge>
                  <span className="text-[10px] uppercase font-bold text-muted-foreground bg-secondary/30 px-2 py-0.5 rounded-md border border-border">English</span>
                </div>
                <div className="flex items-center gap-3">
                  <span className="text-xs font-bold uppercase tracking-wider text-muted-foreground bg-secondary/30 px-2 py-1 rounded-md">+4 / -1</span>
                  <Button 
                    variant="ghost" 
                    size="icon" 
                    onClick={() => markForReview(currentQuestion.id)}
                    className={`h-8 w-8 rounded-lg transition-colors ${currentState?.status.includes('MARKED') ? 'text-amber-500 bg-amber-500/10' : 'text-muted-foreground hover:text-amber-500'}`}
                  >
                    <Flag size={16} />
                  </Button>
                </div>
              </div>

              <div className="prose prose-neutral dark:prose-invert max-w-none text-xl md:text-2xl leading-relaxed text-foreground font-black">
                <ReactMarkdown remarkPlugins={[remarkMath]} rehypePlugins={[rehypeKatex, rehypeRaw]}>
                  {currentQuestion.text_en}
                </ReactMarkdown>
              </div>
              
              {(currentQuestion.type === 'MCQ_Multiple' || currentQuestion.type === 'Multiple_Choice') && (
                <p className="text-sm font-bold text-muted-foreground mt-6 mb-[-1rem] bg-secondary/50 p-2.5 rounded-md border border-border inline-flex items-center gap-2">
                  <AlertTriangle size={16} className="text-amber-500" /> 
                  This question has MULTIPLE CORRECT OPTIONS. Please select all that apply.
                </p>
              )}

              <div className="space-y-3 mt-8">
                {currentQuestion.type === 'Numerical' ? (
                  <input
                    type="text"
                    className="w-full p-4 rounded-xl border border-border bg-background focus:ring-2 focus:ring-primary focus:border-transparent text-lg md:text-xl font-bold"
                    placeholder="Type your numerical answer here..."
                    value={currentState?.selected_answer || ""}
                    onChange={(e) => selectAnswer(currentQuestion.id, e.target.value)}
                  />
                ) : currentQuestion.type === 'Subjective' ? (
                  <textarea
                     className="w-full min-h-[150px] p-4 rounded-xl border border-border bg-background focus:ring-2 focus:ring-primary focus:border-transparent resize-y text-lg md:text-xl font-bold"
                     placeholder="Type your answer or explanation here..."
                     value={currentState?.selected_answer || ""}
                     onChange={(e) => selectAnswer(currentQuestion.id, e.target.value)}
                  />
                ) : (
                  currentQuestion.options?.map((option: string, idx: number) => {
                  const letter = String.fromCharCode(65 + idx);
                  const isMultiple = currentQuestion.type === 'MCQ_Multiple' || currentQuestion.type === 'Multiple_Choice';
                  const isSelected = isMultiple 
                    ? (currentState?.selected_answer || "").split(',').includes(letter)
                    : currentState?.selected_answer === letter;
                    
                  return (
                    <button 
                      key={idx}
                      onClick={() => handleOptionSelect(currentQuestion.id, letter, currentQuestion.type)}
                      className={`w-full text-left p-4.5 rounded-xl border transition-all duration-200 flex items-center gap-4 group focus:outline-none focus:ring-2 focus:ring-primary focus:border-transparent
                        ${isSelected ? 'border-primary bg-primary/5 shadow-[0_0_15px_-5px_var(--color-primary)]' : 'border-border bg-background hover:border-primary/50 hover:bg-secondary/20'}
                      `}
                    >
                      <div className={`flex-shrink-0 size-10 ${isMultiple ? 'rounded-md' : 'rounded-full'} border flex items-center justify-center font-black text-base transition-colors
                        ${isSelected ? 'border-primary bg-primary text-primary-foreground' : 'border-border text-muted-foreground group-hover:border-primary/50 group-hover:text-primary group-hover:bg-primary/5'}
                      `}>
                        {letter}
                      </div>
                      <span className={`text-lg md:text-xl font-bold flex-1 ${isSelected ? 'text-primary' : ''}`}>
                        <ReactMarkdown remarkPlugins={[remarkMath]} rehypePlugins={[rehypeKatex, rehypeRaw]}>
                          {renderLatex(option)}
                        </ReactMarkdown>
                      </span>
                    </button>
                  )
                })
                )}
              </div>
            </div>
          </div>
          
          {/* Optional Bilingual Hindi Pane */}
          <div className="flex-1 p-6 lg:p-10 hidden lg:flex bg-secondary/5 border-b xl:border-b-0">
            <div className="max-w-2xl mx-auto space-y-8 w-full">
               <div className="flex items-center justify-between border-b border-border pb-4">
                <div className="flex items-center gap-2">
                  <Badge variant="outline" className="px-3 py-1 text-sm text-muted-foreground font-medium rounded-lg">प्रश्न {currentQuestion.number}</Badge>
                  <span className="text-[10px] uppercase font-bold text-muted-foreground bg-secondary/30 px-2 py-0.5 rounded-md border border-border">हिंदी</span>
                </div>
                <div className="flex items-center gap-3">
                  <span className="text-xs font-bold uppercase tracking-wider text-muted-foreground bg-secondary/30 px-2 py-1 rounded-md">+4 / -1</span>
                  <Button 
                    variant="ghost" 
                    size="icon" 
                    onClick={() => markForReview(currentQuestion.id)}
                    className={`h-8 w-8 rounded-lg transition-colors ${currentState?.status.includes('MARKED') ? 'text-amber-500 bg-amber-500/10' : 'text-muted-foreground hover:text-amber-500'}`}
                  >
                    <Flag size={16} />
                  </Button>
                </div>
              </div>
              <div className="prose prose-neutral dark:prose-invert max-w-none text-xl md:text-2xl leading-relaxed text-foreground font-black">
                <ReactMarkdown remarkPlugins={[remarkMath]} rehypePlugins={[rehypeKatex, rehypeRaw]}>
                  {currentQuestion.text_hi}
                </ReactMarkdown>
              </div>
              
              {(currentQuestion.type === 'MCQ_Multiple' || currentQuestion.type === 'Multiple_Choice') && (
                <p className="text-sm font-bold text-muted-foreground mt-6 mb-[-1rem] bg-secondary/50 p-2.5 rounded-md border border-border inline-flex items-center gap-2">
                  <AlertTriangle size={16} className="text-amber-500" /> 
                  This question has MULTIPLE CORRECT OPTIONS. Please select all that apply.
                </p>
              )}

              <div className="space-y-3 mt-8">
                {currentQuestion.type === 'Numerical' ? (
                  <input
                    type="text"
                    className="w-full p-4 rounded-xl border border-border bg-background focus:ring-2 focus:ring-primary focus:border-transparent text-lg md:text-xl font-bold"
                    placeholder="अपना संख्यात्मक उत्तर यहाँ टाइप करें..."
                    value={currentState?.selected_answer || ""}
                    onChange={(e) => selectAnswer(currentQuestion.id, e.target.value)}
                  />
                ) : currentQuestion.type === 'Subjective' ? (
                  <textarea
                     className="w-full min-h-[150px] p-4 rounded-xl border border-border bg-background focus:ring-2 focus:ring-primary focus:border-transparent resize-y text-lg md:text-xl font-bold"
                     placeholder="अपना उत्तर या स्पष्टीकरण यहाँ टाइप करें..."
                     value={currentState?.selected_answer || ""}
                     onChange={(e) => selectAnswer(currentQuestion.id, e.target.value)}
                  />
                ) : (
                  currentQuestion.options_hi?.map((option: string, idx: number) => {
                  const letter = String.fromCharCode(65 + idx);
                  const isMultiple = currentQuestion.type === 'MCQ_Multiple' || currentQuestion.type === 'Multiple_Choice';
                  const isSelected = isMultiple 
                    ? (currentState?.selected_answer || "").split(',').includes(letter)
                    : currentState?.selected_answer === letter;
                    
                  return (
                    <button 
                      key={idx}
                      onClick={() => handleOptionSelect(currentQuestion.id, letter, currentQuestion.type)}
                      className={`w-full text-left p-4.5 rounded-xl border transition-all duration-200 flex items-center gap-4 group focus:outline-none focus:ring-2 focus:ring-primary focus:border-transparent
                        ${isSelected ? 'border-primary bg-primary/5 shadow-[0_0_15px_-5px_var(--color-primary)]' : 'border-border bg-background hover:border-primary/50 hover:bg-secondary/20'}
                      `}
                    >
                      <div className={`flex-shrink-0 size-10 ${isMultiple ? 'rounded-md' : 'rounded-full'} border flex items-center justify-center font-black text-base transition-colors
                        ${isSelected ? 'border-primary bg-primary text-primary-foreground' : 'border-border text-muted-foreground group-hover:border-primary/50 group-hover:text-primary group-hover:bg-primary/5'}
                      `}>
                        {letter}
                      </div>
                      <span className={`text-lg md:text-xl font-bold flex-1 ${isSelected ? 'text-primary' : ''}`}>
                        <ReactMarkdown remarkPlugins={[remarkMath]} rehypePlugins={[rehypeKatex, rehypeRaw]}>
                          {renderLatex(option)}
                        </ReactMarkdown>
                      </span>
                    </button>
                  )
                })
                )}
              </div>
            </div>
          </div>
        </main>

        {/* Side Panel for Question Grid */}
        <aside className="w-80 border-l border-border bg-secondary/10 flex flex-col hidden lg:flex shrink-0">
          <div className="p-5 border-b border-border flex items-center justify-between bg-background">
            <h3 className="font-bold text-sm tracking-tight uppercase text-muted-foreground">Question Palette</h3>
            <Button variant="ghost" size="icon" className="h-8 w-8 rounded-lg"><MoreVertical size={16}/></Button>
          </div>
          <div className="flex-1 overflow-y-auto p-5">
            <div className="mb-6 grid grid-cols-2 gap-2 bg-background border border-border p-3 rounded-xl">
              <div className="flex items-center gap-2 text-[11px] font-bold text-muted-foreground">
                <div className="size-4 rounded border border-transparent bg-green-500 shadow-sm text-white flex items-center justify-center">{paletteStats.ANSWERED}</div> <span className="truncate">Answered</span>
              </div>
              <div className="flex items-center gap-2 text-[11px] font-bold text-muted-foreground">
                <div className="size-4 rounded border border-transparent bg-red-500 shadow-sm text-white flex items-center justify-center">{paletteStats.NOT_ANSWERED}</div> <span className="truncate">Not Answered</span>
              </div>
              <div className="flex items-center gap-2 text-[11px] font-bold text-muted-foreground">
                <div className="size-4 rounded border border-border bg-background shadow-sm text-foreground flex items-center justify-center">{paletteStats.UNVISITED}</div> <span className="truncate">Not Visited</span>
              </div>
              <div className="flex items-center gap-2 text-[11px] font-bold text-muted-foreground">
                <div className="size-4 rounded-full border border-transparent bg-amber-500 shadow-sm text-white flex items-center justify-center">{paletteStats.MARKED}</div> <span className="truncate">Marked</span>
              </div>
              <div className="flex items-center gap-2 text-[11px] font-bold text-muted-foreground col-span-2 mt-1">
                <div className="size-4 rounded-full border border-transparent bg-amber-500 shadow-sm text-white flex items-center justify-center relative">
                  {paletteStats.ANSWERED_MARKED}<div className="absolute -bottom-0.5 -right-0.5 size-1.5 rounded-full bg-green-500 border border-white"></div>
                </div> 
                <span className="truncate">Answered & Marked</span>
              </div>
            </div>

            <div className="space-y-6">
              {Object.entries(questionsBySubject).map(([subject, subQuestions]: [string, any]) => (
                <div key={subject}>
                  <h4 className="text-xs font-bold uppercase tracking-wider text-muted-foreground mb-4 flex items-center gap-2">
                    <span className="h-px bg-border flex-1"></span> {subject} <span className="h-px bg-border flex-1"></span>
                  </h4>
                  <div className="grid grid-cols-5 gap-2.5">
                    {subQuestions.map((q: any) => {
                      const i = q.originalIndex;
                      const qState = answers[q.id];
                      const isActive = i === currentQuestionIndex;
                      
                      let btnClass = "bg-background border-border text-muted-foreground hover:border-primary/50 hover:text-foreground";
                      let shapeClass = "rounded"; // Default shape
                      
                      if (qState) {
                        if (qState.status === 'ANSWERED') {
                          btnClass = "bg-green-500 border-green-500 text-white shadow-sm";
                        } else if (qState.status === 'NOT_ANSWERED') {
                          btnClass = "bg-red-500 border-red-500 text-white shadow-sm";
                        } else if (qState.status === 'MARKED') {
                          btnClass = "bg-amber-500 border-amber-500 text-white shadow-sm";
                          shapeClass = "rounded-full";
                        } else if (qState.status === 'ANSWERED_MARKED') {
                          btnClass = "bg-amber-500 border-amber-500 text-white shadow-sm relative";
                          shapeClass = "rounded-full";
                        }
                      }

                      if (isActive) {
                        btnClass = "ring-2 ring-primary ring-offset-2 ring-offset-background border-transparent bg-background shadow-sm text-foreground";
                      }

                      return (
                        <button 
                          key={i} 
                          onClick={() => handleGridClick(i)}
                          className={`size-11 flex items-center justify-center text-sm font-bold transition-all duration-200 border focus:outline-none focus:ring-2 focus:ring-primary ${btnClass} ${shapeClass}`}
                        >
                          {i + 1}
                          {qState?.status === 'ANSWERED_MARKED' && <div className="absolute -bottom-1 -right-1 size-3 rounded-full bg-green-500 border-2 border-background"></div>}
                        </button>
                      )
                    })}
                  </div>
                </div>
              ))}
            </div>
          </div>
          <div className="p-4 border-t border-border bg-background flex flex-col gap-3">
             <div className="flex gap-2">
                <Button 
                  variant="outline" 
                  onClick={() => markForReview(currentQuestion.id)} 
                  className={`flex-1 text-xs font-bold h-9 transition-colors ${currentState?.status.includes('MARKED') ? 'bg-amber-500/10 text-amber-500 border-amber-500/50' : 'bg-secondary/50 hover:bg-amber-500/10 hover:text-amber-500 hover:border-amber-500/50'}`}
                >
                  <Flag size={14} className="mr-1.5"/> {currentState?.status.includes('MARKED') ? 'Unmark' : 'Mark'}
                </Button>
                <Button variant="outline" onClick={() => clearResponse(currentQuestion.id)} className="flex-1 text-xs font-bold h-9 bg-secondary/50 hover:bg-destructive/10 hover:text-destructive hover:border-destructive/50"><XCircle size={14} className="mr-1.5"/> Clear</Button>
             </div>
             <div className="flex gap-2">
                <Button variant="outline" onClick={handlePrev} className="flex-1 gap-1 h-12 rounded-xl font-bold"><ChevronLeft size={16}/> Back</Button>
                <Button onClick={handleNext} className="flex-[2] gap-1 h-12 rounded-xl font-bold shadow-[0_0_15px_-5px_var(--color-primary)]">Save & Next <ChevronRight size={16}/></Button>
             </div>
          </div>
        </aside>

        {/* Mobile bottom nav */}
        <div className="lg:hidden absolute bottom-0 left-0 right-0 p-4 border-t border-border bg-background/95 backdrop-blur-sm grid grid-cols-2 gap-3 z-30">
            <Button variant="outline" onClick={handlePrev} className="rounded-xl w-full gap-1.5 h-12 font-bold"><ChevronLeft size={16}/> Previous</Button>
            <Button onClick={handleNext} className="rounded-xl w-full gap-1.5 h-12 font-bold shadow-sm">Save & Next <ChevronRight size={16}/></Button>
        </div>

        {/* Modals and Overlays */}
        {isPaused && (
          <div className="absolute inset-0 z-40 bg-background/95 backdrop-blur-md flex flex-col items-center justify-center p-8">
            <PauseIcon size={64} className="text-muted-foreground mb-6" />
            <h2 className="text-3xl font-bold mb-4">Exam Paused</h2>
            <p className="text-muted-foreground max-w-md text-center mb-8">
              The exam timer is currently running in the background. Pausing only hides the questions to prevent cheating while taking a break.
            </p>
            <Button size="lg" onClick={() => setIsPaused(false)} className="rounded-xl h-14 px-10 text-lg font-bold gap-2">
              <Play size={20} /> Resume Exam
            </Button>
          </div>
        )}
        
        {showSubmitConfirm && (
          <div className="absolute inset-0 z-50 bg-background/80 backdrop-blur-sm flex items-center justify-center p-4">
            <div className="bg-background border border-border shadow-2xl rounded-2xl p-6 max-w-md w-full animate-in fade-in zoom-in duration-200">
              <h2 className="text-2xl font-bold mb-2">Submit Exam?</h2>
              <p className="text-muted-foreground mb-6">Are you sure you want to submit your exam? You will not be able to change your answers after submission.</p>
              
              <div className="grid grid-cols-2 gap-4 mb-8 p-4 bg-secondary/20 rounded-xl border border-border">
                <div className="text-center">
                  <div className="text-2xl font-bold text-green-500">{paletteStats.ANSWERED + paletteStats.ANSWERED_MARKED}</div>
                  <div className="text-xs font-semibold text-muted-foreground uppercase">Answered</div>
                </div>
                <div className="text-center">
                  <div className="text-2xl font-bold text-red-500">{paletteStats.NOT_ANSWERED + paletteStats.UNVISITED}</div>
                  <div className="text-xs font-semibold text-muted-foreground uppercase">Unanswered</div>
                </div>
                <div className="text-center col-span-2 mt-2 pt-2 border-t border-border/50">
                  <div className="text-2xl font-bold text-amber-500">{paletteStats.MARKED + paletteStats.ANSWERED_MARKED}</div>
                  <div className="text-xs font-semibold text-muted-foreground uppercase">Marked for Review</div>
                </div>
              </div>

              <div className="flex gap-3">
                <Button variant="outline" className="flex-1 rounded-xl h-12 font-bold" onClick={() => setShowSubmitConfirm(false)}>Cancel</Button>
                <Button className="flex-1 rounded-xl h-12 font-bold shadow-[0_0_15px_-3px_var(--color-primary)]" onClick={confirmSubmit}>Yes, Submit</Button>
              </div>
            </div>
          </div>
        )}
        
        <Calculator isOpen={showCalc} onClose={() => setShowCalc(false)} />
        <Scratchpad isOpen={showScratchpad} onClose={() => setShowScratchpad(false)} />
        <PeriodicTable isOpen={showPeriodicTable} onClose={() => setShowPeriodicTable(false)} />
      </div>
    </div>
  )
}