"use client";

import React, { useState, useEffect, useCallback } from 'react';
import { useRouter } from 'next/navigation';
import { notFound } from 'next/navigation';
import type { TestSession, UserAnswer, Question, TestResult } from '@/lib/types';
import { Button } from '@/components/ui/button';
import { Card, CardContent, CardHeader, CardTitle, CardFooter } from '@/components/ui/card';
import { RadioGroup, RadioGroupItem } from '@/components/ui/radio-group';
import { Label } from '@/components/ui/label';
import { AlertDialog, AlertDialogAction, AlertDialogCancel, AlertDialogContent, AlertDialogDescription, AlertDialogFooter, AlertDialogHeader, AlertDialogTitle, AlertDialogTrigger } from '@/components/ui/alert-dialog';
import { ScrollArea, ScrollBar } from '@/components/ui/scroll-area';
import { useToast } from '@/hooks/use-toast';
import { ChevronLeft, ChevronRight, Bookmark, Circle, CheckCircle, XCircle, Loader2 } from 'lucide-react';
import { cn } from '@/lib/utils';
import { useUser, useFirestore } from '@/firebase';
import { doc, writeBatch } from 'firebase/firestore';
import ReactMarkdown from 'react-markdown';
import remarkMath from 'remark-math';
import rehypeKatex from 'rehype-katex';


const Palette = ({
  session,
  currentQuestionIndex,
  onQuestionSelect,
}: {
  session: TestSession;
  currentQuestionIndex: number;
  onQuestionSelect: (index: number) => void;
}) => {
    return (
        <div className='w-full'>
            <ScrollArea className="w-full whitespace-nowrap rounded-md border border-white/10">
                <div className="flex w-max space-x-2 p-4">
                    {session.questions.map((_, index) => {
                        const status = session.userAnswers[index].status;
                        let statusClass = '';
                        let isCurrent = index === currentQuestionIndex;

                        if (isCurrent) {
                            statusClass = 'bg-primary text-primary-foreground scale-110';
                        } else {
                            switch (status) {
                            case 'answered': statusClass = 'bg-green-500/20 border-green-700 text-green-400'; break;
                            case 'not-answered': statusClass = 'bg-red-500/20 border-red-700 text-red-400'; break;
                            case 'marked-for-review': statusClass = 'bg-purple-500/20 border-purple-700 text-purple-400'; break;
                            case 'answered-and-marked-for-review': statusClass = 'bg-yellow-500/20 border-yellow-700 text-yellow-400'; break;
                            default: statusClass = 'bg-muted/40 hover:bg-muted/80 border-white/20';
                            }
                        }
                        
                        return (
                        <Button
                            key={index}
                            variant="outline"
                            size="icon"
                            onClick={() => onQuestionSelect(index)}
                            className={cn(
                            'h-12 w-12 text-lg font-semibold transition-all flex-shrink-0 border',
                            statusClass,
                            isCurrent ? 'ring-2 ring-primary-foreground ring-offset-2 ring-offset-background' : 'hover:scale-105'
                            )}
                        >
                            {index + 1}
                        </Button>
                        );
                    })}
                </div>
                <ScrollBar orientation="horizontal" />
            </ScrollArea>
             <div className="mt-4 grid grid-cols-2 md:grid-cols-3 lg:grid-cols-5 gap-x-4 gap-y-2 text-sm text-muted-foreground">
                <div className="flex items-center gap-2"><CheckCircle className="h-4 w-4 text-green-500" /> Answered</div>
                <div className="flex items-center gap-2"><XCircle className="h-4 w-4 text-red-500" /> Not Answered</div>
                <div className="flex items-center gap-2"><Bookmark className="h-4 w-4 text-purple-500" /> Marked</div>
                <div className="flex items-center gap-2"><Bookmark className="h-4 w-4 text-yellow-500 fill-current" /> Answered & Marked</div>
                <div className="flex items-center gap-2"><Circle className="h-4 w-4 text-muted-foreground" /> Not Visited</div>
            </div>
        </div>
    )
}


export function TestInterface({ testId }: { testId: string }) {
  const router = useRouter();
  const { toast } = useToast();
  const [session, setSession] = useState<TestSession | null>(null);
  const [currentQuestionIndex, setCurrentQuestionIndex] = useState(0);
  const [timeLeft, setTimeLeft] = useState<number | null>(null);
  const [isMounted, setIsMounted] = useState(false);
  const { user } = useUser();
  const firestore = useFirestore();
  const [isSubmitting, setIsSubmitting] = useState(false);

  useEffect(() => {
    const sessionData = localStorage.getItem(`test_session_${testId}`);
    if (sessionData) {
      const parsedSession: TestSession = JSON.parse(sessionData);
      setSession(parsedSession);
      if (parsedSession.config.timeLimit) {
        const elapsedTime = Math.floor((Date.now() - parsedSession.startTime) / 1000);
        setTimeLeft(parsedSession.config.timeLimit * 60 - elapsedTime);
      }
    }
    setIsMounted(true);
  }, [testId]);

  const updateSessionState = (newSession: TestSession) => {
    setSession(newSession);
    localStorage.setItem(`test_session_${testId}`, JSON.stringify(newSession));
  }

  useEffect(() => {
    if(!session || session.userAnswers[currentQuestionIndex].status !== 'not-visited') {
      return;
    }
    const newAnswers = [...session.userAnswers];
    newAnswers[currentQuestionIndex].status = 'not-answered';
    updateSessionState({ ...session, userAnswers: newAnswers });
  // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [currentQuestionIndex, session]);
  
  const submitTest = useCallback(async () => {
    if (!session || !user || !firestore) {
        if (!user) {
            toast({ title: 'Not Signed In', description: 'Please sign in to save your results.', variant: 'destructive'});
        }
        return;
    };
    
    setIsSubmitting(true);

    try {
        const endTime = Date.now();
        let correctAnswers = 0;
        let incorrectAnswers = 0;
        let skipped = 0;

        const performanceByCategory: { [key: string]: { correct: number; incorrect: number; skipped: number; total: number; type: 'subject' | 'topic' } } = {};

        session.questions.forEach((q, i) => {
            const categoryName = q.topic || q.subject;
            const categoryType = q.topic ? 'topic' as const : 'subject' as const;

            if (!performanceByCategory[categoryName]) {
                performanceByCategory[categoryName] = { correct: 0, incorrect: 0, skipped: 0, total: 0, type: categoryType };
            }
            performanceByCategory[categoryName].total++;

            const userAnswer = session.userAnswers[i];
            if (userAnswer.status === 'not-visited' || userAnswer.status === 'not-answered' || userAnswer.status === 'marked-for-review') {
                skipped++;
                performanceByCategory[categoryName].skipped++;
            } else if (userAnswer.selectedOption === q.correctAnswerIndex) {
                correctAnswers++;
                performanceByCategory[categoryName].correct++;
            } else {
                incorrectAnswers++;
                performanceByCategory[categoryName].incorrect++;
            }
        });
        
        const attempted = correctAnswers + incorrectAnswers;
        const accuracy = attempted > 0 ? (correctAnswers / attempted) * 100 : 0;
        
        let score = correctAnswers;
        if(session.config.negativeMarking && session.config.negativeMarkingRatio) {
            score -= incorrectAnswers * session.config.negativeMarkingRatio;
        }

        const performanceBreakdown: import('@/lib/types').PerformanceBreakdown[] = Object.entries(performanceByCategory).map(([category, stats]) => {
            const attemptedInCategory = stats.correct + stats.incorrect;
            return {
                category: category,
                categoryType: stats.type,
                accuracy: attemptedInCategory > 0 ? (stats.correct / attemptedInCategory) * 100 : 0,
                correctCount: stats.correct,
                incorrectCount: stats.incorrect,
                skippedCount: stats.skipped,
                totalCount: stats.total,
            };
        });

        const result: TestResult = {
            id: session.id,
            userId: user.uid,
            config: session.config,
            score,
            accuracy,
            totalQuestions: session.questions.length,
            correctAnswers,
            incorrectAnswers,
            skipped,
            takenAt: new Date(session.startTime).toISOString(),
            timeTaken: Math.floor((endTime - session.startTime) / 1000),
            performanceBreakdown,
        };
        
        const finalSession: TestSession = { ...session, userId: user.uid, endTime: endTime, userAnswers: session.userAnswers.map(ua => ({...ua, selectedOption: ua.selectedOption === undefined ? null : ua.selectedOption})) };
        
        const batch = writeBatch(firestore);

        const resultRef = doc(firestore, 'users', user.uid, 'testResults', session.id);
        batch.set(resultRef, result);

        const sessionRef = doc(firestore, 'users', user.uid, 'userTests', session.id);
        batch.set(sessionRef, finalSession);

        await batch.commit();
        
        localStorage.removeItem(`test_session_${session.id}`);

        toast({ title: 'Test Submitted! 🎉', description: 'Your results have been saved. Redirecting...' });
        router.replace(`/results/${session.id}`);

    } catch (error) {
        console.error("Failed to submit test:", error);
        toast({
            title: "Submission Failed",
            description: "Could not save your test results. Please check your connection and try again.",
            variant: "destructive",
        });
        setIsSubmitting(false);
    }
  }, [session, router, user, firestore, toast]);


  useEffect(() => {
    if (timeLeft === null) return;
    if (timeLeft <= 0) {
      submitTest();
      return;
    }
    const timer = setInterval(() => setTimeLeft(t => (t ?? 0) - 1), 1000);
    return () => clearInterval(timer);
  }, [timeLeft, submitTest]);

  const updateAnswer = (questionIndex: number, selectedOption?: number, statusUpdate?: Partial<UserAnswer>) => {
    if (!session) return;
    const newAnswers = [...session.userAnswers];
    const currentAnswer = newAnswers[questionIndex];
    let newSession: TestSession;

    if (selectedOption !== undefined) {
        const newStatus = currentAnswer.status === 'marked-for-review' || currentAnswer.status === 'answered-and-marked-for-review' ? 'answered-and-marked-for-review' : 'answered';
        newAnswers[questionIndex] = { ...currentAnswer, selectedOption, status: newStatus };
        newSession = { ...session, userAnswers: newAnswers };
    } else if (statusUpdate) {
        newAnswers[questionIndex] = { ...currentAnswer, ...statusUpdate };
        newSession = { ...session, userAnswers: newAnswers };
    } else { // Clear response
        const newStatus = currentAnswer.status === 'answered-and-marked-for-review' ? 'marked-for-review' : 'not-answered';
        newAnswers[questionIndex] = { ...currentAnswer, selectedOption: null, status: newStatus };
        newSession = { ...session, userAnswers: newAnswers };
    }

    updateSessionState(newSession);
  };
  
  const handleOptionChange = (value: string) => {
    updateAnswer(currentQuestionIndex, parseInt(value));
  };

  const handleMarkForReview = () => {
    const currentStatus = session?.userAnswers[currentQuestionIndex].status;
    let newStatus: UserAnswer['status'];
    if (currentStatus === 'answered') {
      newStatus = 'answered-and-marked-for-review';
    } else if (currentStatus === 'answered-and-marked-for-review') {
      newStatus = 'answered';
    } else if (currentStatus === 'marked-for-review') {
      newStatus = 'not-answered';
    } else { // not-answered or not-visited
      newStatus = 'marked-for-review';
    }
    updateAnswer(currentQuestionIndex, undefined, { status: newStatus });
  };

  const clearResponse = () => {
      updateAnswer(currentQuestionIndex, undefined, { selectedOption: null });
  };
  
  const handleQuestionSelect = (index: number) => {
    setCurrentQuestionIndex(index);
  };

  const currentQuestion: Question | undefined = session?.questions[currentQuestionIndex];
  const currentUserAnswer: UserAnswer | undefined = session?.userAnswers[currentQuestionIndex];

  if (!isMounted) return <div className="fixed inset-0 flex items-center justify-center bg-background"><Loader2 className="h-8 w-8 animate-spin" /></div>;
  if (!session || !currentQuestion) {
      useEffect(() => notFound(), []);
      return null;
  }
  
  const formatTime = (seconds: number) => {
    const mins = Math.floor(seconds / 60);
    const secs = seconds % 60;
    return `${mins.toString().padStart(2, '0')}:${secs.toString().padStart(2, '0')}`;
  };

  return (
    <div className="bg-gradient-to-br from-[#0c132c] to-[#121a3a] text-white min-h-[calc(100vh-3.5rem)]">
      <main className="max-w-6xl mx-auto p-4 md:p-8">
            <Card className="shadow-2xl bg-white/5 border-white/10 backdrop-blur-lg">
              <CardHeader className="flex flex-col sm:flex-row items-start sm:items-center justify-between gap-4">
                  <div className='flex-1'>
                    <p className='text-sm text-slate-400'>Question {currentQuestionIndex + 1} of {session.questions.length}</p>
                    <CardTitle className="text-xl md:text-2xl mt-1 font-headline">
                      {session.config.exam}
                    </CardTitle>
                  </div>
                {timeLeft !== null && (
                    <div className="text-right shrink-0">
                        <span className="text-sm text-slate-400 block">Time Left</span>
                        <div className="text-2xl font-mono font-bold tabular-nums text-white">{formatTime(timeLeft)}</div>
                    </div>
                )}
              </CardHeader>

              <CardContent className="space-y-6">
                <div className="prose prose-lg prose-invert max-w-none font-semibold text-slate-200">
                    <ReactMarkdown remarkPlugins={[remarkMath]} rehypePlugins={[rehypeKatex]}>
                        {currentQuestion.questionText}
                    </ReactMarkdown>
                </div>
                <RadioGroup key={currentQuestionIndex} value={currentUserAnswer?.selectedOption?.toString()} onValueChange={handleOptionChange} className="space-y-3">
                    {currentQuestion.options.map((option, index) => (
                        <Label key={index} htmlFor={`option-${index}`} className={cn(
                            "flex items-start rounded-lg border p-4 transition-all duration-300 cursor-pointer text-base bg-black/20 border-white/10",
                            "hover:border-primary/70 hover:bg-primary/10",
                            currentUserAnswer?.selectedOption === index && "border-primary bg-primary/20 ring-2 ring-primary"
                         )}>
                            <RadioGroupItem value={index.toString()} id={`option-${index}`} className="h-6 w-6 mt-0.5 border-slate-500 text-primary focus:ring-primary focus:ring-offset-background"/>
                            <span className="ml-4 font-medium flex-1 text-slate-300">
                                <ReactMarkdown components={{ p: React.Fragment }} remarkPlugins={[remarkMath]} rehypePlugins={[rehypeKatex]}>
                                    {option}
                                </ReactMarkdown>
                            </span>
                        </Label>
                    ))}
                </RadioGroup>
              </CardContent>
              
              <div className="px-6 py-4 border-y border-white/10">
                <Palette 
                    session={session} 
                    currentQuestionIndex={currentQuestionIndex}
                    onQuestionSelect={handleQuestionSelect}
                />
              </div>

              <CardFooter className="flex flex-wrap items-center justify-between gap-4 pt-6">
                <div className='flex gap-2 flex-wrap'>
                  <Button variant="outline" className="bg-transparent border-slate-600 hover:bg-slate-700 hover:text-slate-100" onClick={handleMarkForReview}><Bookmark className="mr-2 h-4 w-4"/>Mark for Review</Button>
                  <Button variant="ghost" className="hover:bg-slate-700 hover:text-slate-100" onClick={clearResponse}>Clear Response</Button>
                </div>
                <div className="flex gap-2">
                  <Button variant="secondary" onClick={() => setCurrentQuestionIndex(p => Math.max(0, p - 1))} disabled={currentQuestionIndex === 0}><ChevronLeft className="mr-2 h-4 w-4"/>Previous</Button>
                  {currentQuestionIndex < session.questions.length - 1 ? (
                    <Button onClick={() => setCurrentQuestionIndex(p => Math.min(session.questions.length - 1, p + 1))}>Save & Next <ChevronRight className="ml-2 h-4 w-4"/></Button>
                  ) : (
                     <AlertDialog>
                        <AlertDialogTrigger asChild><Button variant="default" className="bg-green-600 hover:bg-green-700 text-white">Submit Test</Button></AlertDialogTrigger>
                        <AlertDialogContent>
                            <AlertDialogHeader>
                                <AlertDialogTitle>Ready to finish?</AlertDialogTitle>
                                <AlertDialogDescription>
                                    This will end the test and calculate your score. You cannot go back after submitting.
                                </AlertDialogDescription>
                            </AlertDialogHeader>
                            <AlertDialogFooter>
                                <AlertDialogCancel>Cancel</AlertDialogCancel>
                                <AlertDialogAction onClick={submitTest} disabled={isSubmitting}>
                                  {isSubmitting ? (
                                    <>
                                      <Loader2 className="mr-2 h-4 w-4 animate-spin" />
                                      Submitting...
                                    </>
                                  ) : (
                                    'Submit'
                                  )}
                                </AlertDialogAction>
                            </AlertDialogFooter>
                        </AlertDialogContent>
                    </AlertDialog>
                  )}
                </div>
              </CardFooter>
            </Card>
      </main>
    </div>
  );
}