src/pages/LugolExamPage.tsx

import { useState, useEffect } from 'react';
import { Camera, Video, ArrowLeft, ArrowRight, CheckCircle2, Info, Pause, X, Edit2, RotateCcw, Save, ChevronRight, Sparkles } from 'lucide-react';
import { ImageAnnotator } from '../components/ImageAnnotator';
import { ImagingObservations } from '../components/ImagingObservations';

type CapturedItem = {
  id: string;
  type: 'image' | 'video';
  url: string;
  timestamp: Date;
  annotations?: any[];
};

type Props = {
  goBack: () => void;
  onNext: () => void;
};

export function LugolExamPage({ goBack, onNext }: Props) {
  const [capturedItems, setCapturedItems] = useState<CapturedItem[]>([]);
  const [isRecording, setIsRecording] = useState(false);
  const [selectedImage, setSelectedImage] = useState<string | null>(null);
  const [annotations, setAnnotations] = useState<any[]>([]);
  const [observations, setObservations] = useState({});
  const [showExitWarning, setShowExitWarning] = useState(false);
  
  // Timer states
  const [timerStarted, setTimerStarted] = useState(false);
  const [seconds, setSeconds] = useState(0);
  const [lugolApplied, setLugolApplied] = useState(false);
  const [showFlash, setShowFlash] = useState(false);
  const audibleAlert = true;
  const [timerPaused, setTimerPaused] = useState(false);
  const [isLiveAILoading, setIsLiveAILoading] = useState(false);
  const [isAIAssistEnabled, setIsAIAssistEnabled] = useState(false);
  const [liveAIResults, setLiveAIResults] = useState<{ cervixDetected: boolean; quality: string; confidence: number } | null>(null);
  const [liveAIError, setLiveAIError] = useState<string | null>(null);

  const cervixImageUrl = "/C87Aceto_(1).jpg";

  // Timer effect
  useEffect(() => {
    if (!timerStarted || !lugolApplied || timerPaused) return;

    const interval = setInterval(() => {
      setSeconds(prev => prev + 1);
    }, 1000);

    return () => clearInterval(interval);
  }, [timerStarted, lugolApplied, timerPaused]);

  // Check for 1 minute and 3 minute marks
  useEffect(() => {
    if (seconds === 60) {
      // 1 minute mark
      setShowFlash(true);
      if (audibleAlert) {
        console.log('BEEP - 1 minute mark');
      }
      setTimeout(() => setShowFlash(false), 3000);
    } else if (seconds === 180) {
      // 3 minute mark
      setShowFlash(true);
      if (audibleAlert) {
        console.log('BEEP - 3 minute mark');
      }
      setTimeout(() => setShowFlash(false), 3000);
    }
  }, [seconds, audibleAlert]);

  const formatTime = (totalSeconds: number) => {
    const mins = Math.floor(totalSeconds / 60);
    const secs = totalSeconds % 60;
    return `${mins.toString().padStart(2, '0')}:${secs.toString().padStart(2, '0')}`;
  };

  const handleLugolApplied = () => {
    setLugolApplied(true);
    setTimerStarted(true);
    setSeconds(0);
  };

  const handleRestartTimer = () => {
    setSeconds(0);
    setTimerStarted(false);
    setLugolApplied(false);
    setShowFlash(false);
    setTimerPaused(false);
  };

  const handleCaptureImage = () => {
    const newCapture: CapturedItem = {
      id: Date.now().toString(),
      type: 'image',
      url: cervixImageUrl,
      timestamp: new Date()
    };
    setCapturedItems(prev => [...prev, newCapture]);
  };

  const handleToggleRecording = () => {
    if (!isRecording) {
      setIsRecording(true);
    } else {
      setIsRecording(false);
      const newCapture: CapturedItem = {
        id: Date.now().toString(),
        type: 'video',
        url: cervixImageUrl,
        timestamp: new Date()
      };
      setCapturedItems(prev => [...prev, newCapture]);
    }
  };

  const handleSaveAnnotations = () => {
    if (!selectedImage) return;
    
    setCapturedItems(prev => prev.map(item =>
      item.id === selectedImage
        ? { ...item, annotations }
        : item
    ));
    setSelectedImage(null);
    setAnnotations([]);
  };

  const handleDeleteCapture = (id: string) => {
    setCapturedItems(prev => prev.filter(item => item.id !== id));
    if (selectedImage === id) {
      setSelectedImage(null);
    }
  };

  const mapQualityLabel = (score: number) => {
    if (score >= 0.8) return 'Excellent';
    if (score >= 0.6) return 'Good';
    return 'Bad';
  };

  const handleAIAssistToggle = async () => {
    if (isLiveAILoading) return;
    if (isAIAssistEnabled) {
      setIsAIAssistEnabled(false);
      setLiveAIResults(null);
      setLiveAIError(null);
      return;
    }
    setIsAIAssistEnabled(true);
    await handleLugolMainAIAssist();
  };

  const handleLugolMainAIAssist = async () => {
    setLiveAIError(null);
    setLiveAIResults(null);
    const imageItems = capturedItems.filter(item => item.type === 'image');
    const targetItem = imageItems[0];

    setIsLiveAILoading(true);

    try {
      const response = await fetch(targetItem ? targetItem.url : cervixImageUrl);
      const blob = await response.blob();

      const formData = new FormData();
      formData.append('file', blob, 'image.jpg');
      const backendResponse = await fetch('/infer/image', {
        method: 'POST',
        body: formData,
      });

      if (!backendResponse.ok) {
        throw new Error(`Backend error: ${backendResponse.statusText}`);
      }

      const result = await backendResponse.json();

      const qualityScore = typeof result.quality_score === 'number'
        ? result.quality_score
        : (typeof result.quality_percent === 'number' ? result.quality_percent / 100 : 0);

      setLiveAIResults({
        cervixDetected: Boolean(result.detected),
        quality: mapQualityLabel(qualityScore),
        confidence: qualityScore
      });
      setIsLiveAILoading(false);
    } catch (error) {
      console.error('Live AI assist error:', error);
      setLiveAIError(error instanceof Error ? error.message : 'Failed to check image quality');
      setIsLiveAILoading(false);
    }
  };

  const selectedItem = selectedImage
    ? capturedItems.find(item => item.id === selectedImage)
    : null;

  const totalCaptures = capturedItems.length;
  const imageCaptures = capturedItems.filter(item => item.type === 'image');
  const videoCaptures = capturedItems.filter(item => item.type === 'video');
  const hasRequiredCapture = imageCaptures.length > 0;

  return (
    <div className="w-full bg-white/95 relative">
      <div className="relative z-10 py-4 md:py-6 lg:py-8">
        <div className="w-full max-w-7xl mx-auto px-4 md:px-6">
          
          {/* Page Header */}
          <div className="mb-4 md:mb-6">
            <div className="flex items-center justify-between mb-4">
              <div className="flex items-center gap-3">
                <button onClick={() => setShowExitWarning(true)} className="p-2 hover:bg-gray-100 rounded-lg transition-colors text-gray-600">
                  <ArrowLeft className="w-5 h-5" />
                </button>
                <h1 className="text-xl md:text-2xl lg:text-3xl font-bold text-[#0A2540]">Lugol Examination</h1>
              </div>
              <div className="flex items-center gap-3">
              </div>
            </div>

            {/* Progress Bar - Capture / Annotation / Comparison View / Report */}
            <div className="mb-4 flex gap-1 md:gap-2 items-center">
              <div className="flex gap-1 md:gap-2 flex-1">
                {['Capture', 'Annotation', 'Comparison View', 'Report'].map((stage, idx) => (
                  <div key={stage} className="flex items-center flex-1">
                    <button
                      className={`flex-1 py-2 px-2 md:px-3 rounded-lg font-medium text-sm md:text-base transition-all border-2 border-[#0A2540] ${
                        (stage === 'Capture' && !selectedImage) ||
                        (stage === 'Annotation' && selectedImage)
                          ? 'bg-[#05998c] text-white shadow-md'
                          : 'bg-gray-100 text-gray-600'
                      }`}
                    >
                      {stage}
                    </button>
                    {idx < 3 && <div className="w-1.5 h-1.5 rounded-full bg-gray-300 mx-1" />}
                  </div>
                ))}
              </div>
              <button onClick={onNext} className="ml-4 px-6 md:px-8 py-2 md:py-3 rounded-xl bg-gray-600 text-white font-bold shadow-lg shadow-gray-500/20 hover:bg-slate-700 hover:shadow-gray-500/30 transition-all flex items-center justify-center gap-2 text-sm md:text-base">
                <Save className="w-4 h-4 md:w-5 md:h-5" />
                <span className="hidden lg:inline">Next</span>
                <span className="inline lg:hidden">Next</span>
                <ChevronRight className="w-4 h-4 md:w-5 md:h-5" />
              </button>
            </div>
          </div>

          {!selectedImage ? (
            // Live Feed View
            <div className="grid grid-cols-1 lg:grid-cols-3 gap-6">
              {/* Main Live Feed */}
              <div className="lg:col-span-2 space-y-4">
                <div className="bg-white rounded-xl shadow-sm border border-gray-100 p-6">
                  <div className="mb-4">
                    <h2 className="text-2xl md:text-3xl font-bold text-[#0A2540] mb-2">
                      Lugol Iodine
                    </h2>
                    <p className="text-gray-600">
                      Lugol iodine application and observation
                    </p>
                  </div>

                  {/* Live Video Feed */}
                  <div className="relative bg-gray-900 rounded-xl overflow-hidden shadow-2xl border-2 border-gray-700 mb-4">
                    <div className="aspect-video flex items-center justify-center">
                      <img src={cervixImageUrl} alt="Live Feed" className="w-full h-full object-cover" />
                      
                      {/* Live indicator */}
                      <div className="absolute top-4 left-4 flex items-center gap-2 bg-red-500 text-white px-3 py-1 rounded-full text-sm font-semibold">
                        <div className={`w-2 h-2 rounded-full ${isRecording ? 'bg-white animate-pulse' : 'bg-white/70'}`} />
                        {isRecording ? 'Recording' : 'Live'}
                      </div>

                      {/* Timer overlay */}
                      {lugolApplied && (
                        <div className="absolute top-4 right-4 bg-black/70 text-white px-4 py-2 rounded-lg">
                          <p className="text-2xl font-mono font-bold">{formatTime(seconds)}</p>
                        </div>
                      )}

                      {/* Flash overlay at 1 and 3 minutes */}
                      {showFlash && (
                        <div className="absolute inset-0 bg-[#05998c]/30 animate-pulse flex items-center justify-center">
                          <div className="bg-white/90 px-6 py-4 rounded-lg">
                            <p className="text-2xl font-bold text-[#0A2540]">
                              {seconds >= 180 ? '3 Minutes!' : '1 Minute!'}
                            </p>
                          </div>
                        </div>
                      )}
                    </div>
                  </div>

                  <div className="mt-4 flex items-center gap-2 text-sm text-gray-500">
                    <Camera className="w-4 h-4" />
                    <span>Captures: {totalCaptures} / 1 required (image)</span>
                    {hasRequiredCapture && <CheckCircle2 className="w-4 h-4 text-green-500 ml-2" />}
                  </div>

                  {!hasRequiredCapture && timerStarted && (
                    <div className="mt-4 flex items-center gap-2 text-sm text-amber-600 bg-amber-50 px-4 py-2 rounded-lg">
                      <Info className="w-4 h-4" />
                      <span>At least one image capture required</span>
                    </div>
                  )}

                  {/* Captured Images Selection for Annotation */}
                  {imageCaptures.length > 0 && (
                    <div className="mt-6 pt-6 border-t border-gray-200">
                      <h4 className="text-sm font-semibold text-gray-700 mb-3">Select Image to Annotate</h4>
                      <div className="grid grid-cols-3 gap-3">
                        {imageCaptures.map(item => (
                          <div 
                            key={item.id} 
                            onClick={() => setSelectedImage(item.id)}
                            className="relative group cursor-pointer"
                          >
                            <div className="aspect-square bg-gray-100 rounded-lg overflow-hidden border-2 border-transparent hover:border-[#05998c] transition-all">
                              <img src={item.url} alt="Capture" className="w-full h-full object-cover" />
                              {item.annotations && item.annotations.length > 0 && (
                                <div className="absolute top-1 right-1 bg-green-500 text-white p-1 rounded">
                                  <CheckCircle2 className="w-3 h-3" />
                                </div>
                              )}
                            </div>
                            <div className="absolute inset-0 bg-black/50 opacity-0 group-hover:opacity-100 transition-opacity rounded-lg flex items-center justify-center">
                              <span className="text-white text-xs font-semibold">Annotate</span>
                            </div>
                          </div>
                        ))}
                      </div>
                    </div>
                  )}
                </div>
              </div>

              {/* Captures Sidebar */}
              <div className="lg:col-span-1">
                <div className="bg-white rounded-xl shadow-sm border border-gray-100 p-6">
                  {/* Lugol Timer Section */}
                  {!timerStarted && (
                    <div className="mb-6">
                      {/* Apply Lugol Message Box */}
                      <div className="bg-gradient-to-r from-yellow-500 to-yellow-600 rounded-lg p-4 mb-4 shadow-md">
                        <div className="flex items-center gap-3 mb-2">
                          <div className="w-10 h-10 bg-white rounded-full flex items-center justify-center">
                            <Info className="w-6 h-6 text-yellow-600" />
                          </div>
                          <div>
                            <p className="text-white font-bold text-lg">Apply Lugol Iodine Now</p>
                            <p className="text-yellow-100 text-sm">Apply Lugol iodine solution to the cervix</p>
                          </div>
                        </div>
                      </div>

                      <button
                        onClick={handleLugolApplied}
                        className="w-full px-6 py-3 bg-[#05998c] text-white rounded-lg font-semibold hover:bg-[#047569] transition-all shadow-md hover:shadow-lg"
                      >
                        Lugol applied — Start timer
                      </button>
                    </div>
                  )}

                  {timerStarted && (
                    <div className={`mb-6 rounded-lg p-4 border-2 transition-all ${
                      seconds < 50 
                        ? 'bg-gradient-to-r from-[#05998c]/10 to-[#0A2540]/10 border-[#05998c]'
                        : seconds < 55
                        ? 'bg-red-50 border-red-200'
                        : seconds < 60
                        ? 'bg-red-100 border-red-300'
                        : seconds >= 60 && seconds <= 60
                        ? 'bg-red-200 border-red-400'
                        : seconds > 60 && seconds < 170
                        ? 'bg-gradient-to-r from-green-50 to-green-50 border-green-300'
                        : seconds < 175
                        ? 'bg-red-50 border-red-200'
                        : seconds < 180
                        ? 'bg-red-100 border-red-300'
                        : 'bg-red-200 border-red-400'
                    }`}>
                      <div className="flex flex-col gap-3">
                        <div>
                          <p className="text-sm text-gray-600 mb-1">Timer</p>
                          <p className={`text-4xl font-bold font-mono ${
                            seconds >= 180 ? 'text-red-600' : 
                            seconds >= 60 ? 'text-amber-500' : 
                            seconds >= 50 ? 'text-amber-400' : 
                            'text-[#0A2540]'
                          }`}>{formatTime(seconds)}</p>
                          {seconds >= 60 && seconds < 180 && (
                            <p className="text-sm text-amber-600 mt-1">Approaching 3-minute mark...</p>
                          )}
                          {seconds >= 180 && (
                            <p className="text-sm text-green-600 mt-1">3-minute observation period complete</p>
                          )}
                        </div>
                        <div className="flex gap-2">
                          <button
                            onClick={() => setTimerPaused(!timerPaused)}
                            className="flex-1 flex items-center justify-center gap-2 px-3 py-2 bg-white border-2 border-gray-300 rounded-lg hover:bg-gray-50 transition-colors"
                          >
                            {timerPaused ? (
                              <>
                                <Video className="w-4 h-4" />
                                <span className="text-sm font-medium">Play</span>
                              </>
                            ) : (
                              <>
                                <Pause className="w-4 h-4" />
                                <span className="text-sm font-medium">Pause</span>
                              </>
                            )}
                          </button>
                          <button
                            onClick={handleRestartTimer}
                            className="flex-1 flex items-center justify-center gap-2 px-3 py-2 bg-white border-2 border-gray-300 rounded-lg hover:bg-gray-50 transition-colors"
                          >
                            <RotateCcw className="w-4 h-4" />
                            <span className="text-sm font-medium">Restart</span>
                          </button>
                        </div>
                      </div>
                      {seconds === 60 && (
                        <div className="mt-3 bg-amber-100 border border-amber-300 rounded-lg p-3">
                          <p className="text-amber-800 font-semibold text-sm">⏰ 1-minute mark - Capture recommended!</p>
                        </div>
                      )}
                      {seconds === 180 && (
                        <div className="mt-3 bg-green-100 border border-green-300 rounded-lg p-3">
                          <p className="text-green-800 font-semibold text-sm">⏰ 3-minute mark - Final capture recommended!</p>
                        </div>
                      )}
                    </div>
                  )}

                  {/* Capture Controls */}
                  <div className="space-y-3 mb-6">
                    <div className="flex gap-2">
                      <button
                        onClick={handleCaptureImage}
                        disabled={!timerStarted}
                        className={`flex-1 flex items-center justify-center gap-2 px-4 py-3 rounded-lg font-semibold transition-colors text-sm ${
                          showFlash && (seconds === 60 || seconds === 180)
                            ? 'bg-[#05998c] text-white animate-pulse'
                            : 'bg-[#05998c] text-white hover:bg-[#047569]'
                        } disabled:opacity-50 disabled:cursor-not-allowed`}
                      >
                        <Camera className="w-4 h-4" />
                        Capture
                      </button>
                      <button
                        onClick={handleToggleRecording}
                        disabled={!timerStarted}
                        className={`flex-1 flex items-center justify-center gap-2 px-4 py-3 rounded-lg font-semibold transition-colors text-sm ${
                          isRecording
                            ? 'bg-red-500 text-white hover:bg-red-600'
                            : 'bg-[#05998c] text-white hover:bg-[#047569]'
                        } disabled:opacity-50 disabled:cursor-not-allowed`}
                      >
                        {isRecording ? <Pause className="w-4 h-4" /> : <Video className="w-4 h-4" />}
                        {isRecording ? 'Stop' : 'Record'}
                      </button>
                    </div>

                    {/* Centered AI Assist Button */}
                    <button
                      onClick={handleAIAssistToggle}
                      disabled={isLiveAILoading}
                      className={`w-full flex items-center justify-center gap-2 px-6 py-4 rounded-lg text-white font-bold transition-all shadow-lg hover:shadow-xl disabled:opacity-50 disabled:cursor-not-allowed text-base ${
                        isAIAssistEnabled
                          ? 'bg-gradient-to-r from-green-500 to-green-600 hover:from-green-600 hover:to-green-700'
                          : 'bg-gradient-to-r from-blue-600 to-blue-700 hover:from-blue-700 hover:to-blue-800'
                      }`}
                      title={isAIAssistEnabled ? 'AI Assist is ON' : 'Run AI model to check quality'}
                    >
                      <div className={`relative inline-flex h-6 w-11 items-center rounded-full transition-colors ${isAIAssistEnabled ? 'bg-white' : 'bg-gray-300'}`}>
                        <span className={`inline-block h-4 w-4 transform rounded-full bg-gradient-to-br transition-transform ${isAIAssistEnabled ? 'translate-x-6 from-green-400 to-green-600' : 'translate-x-1 from-blue-400 to-blue-600'}`} />
                      </div>
                      <Sparkles className="w-6 h-6" />
                      {isLiveAILoading ? 'Checking...' : (isAIAssistEnabled ? 'AI Assist On' : 'AI Assist')}
                    </button>
                    
                    {/* Live AI Results Panel */}
                    {liveAIResults && (
                      <div className="p-4 bg-green-50 border border-green-300 rounded-lg">
                        <div className="flex items-center gap-3 mb-3">
                          <div className="w-3 h-3 bg-green-500 rounded-full"></div>
                          <h4 className="font-bold text-green-800">Quality Check Results</h4>
                        </div>
                        <div className="space-y-2 text-sm">
                          <p className="text-gray-700">
                            <span className="font-semibold">Cervix Detected:</span> {liveAIResults.cervixDetected ? 'Yes' : 'No'} ({((liveAIResults.cervixDetected ? liveAIResults.confidence : 0) * 100).toFixed(1)}%)
                          </p>
                          <p className="text-gray-700">
                            <span className="font-semibold">Quality:</span> {liveAIResults.quality} ({(liveAIResults.confidence * 100).toFixed(1)}%)
                          </p>
                        </div>
                      </div>
                    )}
                    
                    {liveAIError && (
                      <div className="text-sm text-red-600 bg-red-50 border border-red-200 rounded-lg p-3 text-center">
                        {liveAIError}
                      </div>
                    )}
                    <button className="w-full flex items-center justify-center gap-2 px-4 py-3 bg-gray-600 text-white rounded-lg font-semibold hover:bg-slate-700 transition-colors">
                      Next
                      <ArrowRight className="w-4 h-4" />
                    </button>
                  </div>
                  <h3 className="font-bold text-[#0A2540] mb-4">Captured Media</h3>
                  
                  {totalCaptures === 0 ? (
                    <div className="flex flex-col items-center justify-center py-12 text-center">
                      <Camera className="w-16 h-16 text-gray-300 mb-3" />
                      <p className="text-gray-500 font-medium">No captures yet</p>
                      <p className="text-sm text-gray-400 mt-1">Apply Lugol iodine and start capturing</p>
                    </div>
                  ) : (
                    <div className="space-y-3">
                      {/* Image Thumbnails */}
                      {imageCaptures.length > 0 && (
                        <div>
                          <h4 className="text-xs font-semibold text-gray-500 uppercase mb-2">Images ({imageCaptures.length})</h4>
                          <div className="grid grid-cols-2 gap-2">
                            {imageCaptures.map(item => (
                              <div key={item.id} className="relative group">
                                <div 
                                  onClick={() => setSelectedImage(item.id)}
                                  className="aspect-square bg-gray-100 rounded-lg overflow-hidden cursor-pointer border-2 border-transparent hover:border-[#05998c] transition-all"
                                >
                                  <img src={item.url} alt="Capture" className="w-full h-full object-cover" />
                                  {item.annotations && item.annotations.length > 0 && (
                                    <div className="absolute top-1 right-1 bg-green-500 text-white p-1 rounded">
                                      <CheckCircle2 className="w-3 h-3" />
                                    </div>
                                  )}
                                </div>
                                <button
                                  onClick={(e) => {
                                    e.stopPropagation();
                                    handleDeleteCapture(item.id);
                                  }}
                                  className="absolute top-1 left-1 bg-red-500 text-white p-1 rounded opacity-0 group-hover:opacity-100 transition-opacity"
                                >
                                  <X className="w-3 h-3" />
                                </button>
                                <button
                                  onClick={() => setSelectedImage(item.id)}
                                  className="absolute bottom-1 right-1 bg-[#0A2540] text-white p-1 rounded opacity-0 group-hover:opacity-100 transition-opacity"
                                >
                                  <Edit2 className="w-3 h-3" />
                                </button>
                              </div>
                            ))}
                          </div>
                        </div>
                      )}

                      {/* Video Items */}
                      {videoCaptures.length > 0 && (
                        <div className="mt-4">
                          <h4 className="text-xs font-semibold text-gray-500 uppercase mb-2">Videos ({videoCaptures.length})</h4>
                          <div className="space-y-2">
                            {videoCaptures.map(item => (
                              <div key={item.id} className="relative group bg-gray-50 rounded-lg p-3 flex items-center gap-3">
                                <div className="w-12 h-12 bg-gray-200 rounded flex items-center justify-center">
                                  <Video className="w-6 h-6 text-gray-500" />
                                </div>
                                <div className="flex-1">
                                  <p className="text-sm font-medium text-gray-700">Video Recording</p>
                                  <p className="text-xs text-gray-500">{item.timestamp.toLocaleTimeString()}</p>
                                </div>
                                <button
                                  onClick={() => handleDeleteCapture(item.id)}
                                  className="p-1 hover:bg-red-50 rounded text-red-500"
                                >
                                  <X className="w-4 h-4" />
                                </button>
                              </div>
                            ))}
                          </div>
                        </div>
                      )}
                    </div>
                  )}
                </div>
              </div>
            </div>
          ) : (
            // Annotation View
            <div>
              <div className="mb-4 flex items-center justify-between">
                <button
                  onClick={() => {
                    setSelectedImage(null);
                    setAnnotations([]);
                  }}
                  className="flex items-center gap-2 px-4 py-2 text-gray-600 hover:bg-gray-100 rounded-lg transition-colors"
                >
                  <ArrowLeft className="w-4 h-4" />
                  Back to Live Feed
                </button>
                <div className="flex items-center gap-3">
                  <button
                    onClick={handleSaveAnnotations}
                    className="px-6 py-2 bg-[#05998c] text-white rounded-lg font-semibold hover:bg-[#047569] transition-colors"
                  >
                    Save Annotations
                  </button>
                  <button className="px-6 py-2 bg-gray-600 text-white rounded-lg font-semibold hover:bg-slate-700 transition-colors flex items-center gap-2">
                    Next
                    <ArrowRight className="w-4 h-4" />
                  </button>
                </div>
              </div>

              <div className="grid grid-cols-1 lg:grid-cols-3 gap-6">
                <div className="lg:col-span-2">
                  <ImageAnnotator 
                    imageUrl={selectedItem?.url || cervixImageUrl} 
                    onAnnotationsChange={setAnnotations} 
                  />
                </div>
                <div className="lg:col-span-1">
                  <ImagingObservations 
                    onObservationsChange={setObservations}
                    stepId="lugol"
                  />
                </div>
              </div>
            </div>
          )}

          {/* Exit Warning Dialog */}
          {showExitWarning && (
            <div className="fixed inset-0 bg-black/50 flex items-center justify-center z-50">
              <div className="bg-white rounded-xl shadow-2xl p-6 max-w-md mx-4">
                <h3 className="text-xl font-bold text-[#0A2540] mb-3">Leave Examination?</h3>
                <p className="text-gray-600 mb-6">
                  If you go back now, all captures, timer data, and annotations will be lost. Are you sure you want to continue?
                </p>
                <div className="flex gap-3 justify-end">
                  <button
                    onClick={() => setShowExitWarning(false)}
                    className="px-6 py-2 bg-gray-100 text-gray-700 rounded-lg font-semibold hover:bg-gray-200 transition-colors"
                  >
                    Cancel
                  </button>
                  <button
                    onClick={() => {
                      setShowExitWarning(false);
                      goBack();
                    }}
                    className="px-6 py-2 bg-red-500 text-white rounded-lg font-semibold hover:bg-red-600 transition-colors"
                  >
                    Leave Anyway
                  </button>
                </div>
              </div>
            </div>
          )}
        </div>
      </div>
    </div>
  );
}