src/components/ImageAnnotator.tsx

import React, { useEffect, useState, useRef, useImperativeHandle, forwardRef } from 'react';
import { Wrench, Trash2, Circle as CircleIcon, Hexagon, Square as SquareIcon, ChevronDown, ChevronUp, Sparkles, Loader } from 'lucide-react';

type ShapeType = 'rect' | 'circle' | 'polygon';

interface Point {
  x: number;
  y: number;
}

interface Annotation {
  id: string;
  type: ShapeType;
  x: number;
  y: number;
  width: number;
  height: number;
  color: string;
  label?: string;
  points?: Point[];
  source?: 'manual' | 'ai';
  identified?: boolean;
  accepted?: boolean;
}

interface ImageAnnotatorProps {
  imageUrl?: string;
  imageUrls?: string[];
  onAnnotationsChange?: (annotations: Annotation[]) => void;
  onAIAssist?: () => Promise<void>;
  isAILoading?: boolean;
}

export interface ImageAnnotatorHandle {
  addAIAnnotations: (aiAnnotations: Annotation[]) => void;
  setImageIndex: (index: number) => void;
  clearAIAnnotations: () => void;
  clearAllAnnotations: () => void;
  resetViewport: () => void;
  waitForImageReady: () => Promise<void>;
}

const ImageAnnotatorComponent = forwardRef<ImageAnnotatorHandle, ImageAnnotatorProps>(({ imageUrl, imageUrls, onAnnotationsChange, onAIAssist, isAILoading: externalAILoading = false }, ref) => {
  const canvasRef = useRef<HTMLCanvasElement | null>(null);
  const containerRef = useRef<HTMLDivElement | null>(null);
  const imageReadyResolveRef = useRef<(() => void) | null>(null);
  const cachedImagesRef = useRef<Record<string, HTMLImageElement>>({});
  const pendingDrawAnimationRef = useRef<number | null>(null);
  const lastDrawStateRef = useRef<string>('');
  const [annotations, setAnnotations] = useState<Annotation[]>([]);
  const [tool, setTool] = useState<ShapeType>('rect');
  const [color, setColor] = useState('#05998c');
  const [labelInput, setLabelInput] = useState('');
  const [isDrawing, setIsDrawing] = useState(false);
  const [startPoint, setStartPoint] = useState<Point | null>(null);
  const [currentAnnotation, setCurrentAnnotation] = useState<Annotation | null>(null);
  const [imageLoaded, setImageLoaded] = useState(false);
  const [imageDimensions, setImageDimensions] = useState({ width: 0, height: 0 });
  const [polygonPoints, setPolygonPoints] = useState<Point[]>([]);
  const [selectedImageIndex, setSelectedImageIndex] = useState(0);
  const [isAnnotationsOpen, setIsAnnotationsOpen] = useState(true);
  const [isLabelDropdownOpen, setIsLabelDropdownOpen] = useState(false);
  const [isAIAssistEnabled, setIsAIAssistEnabled] = useState(false);
  const [internalAILoading, setInternalAILoading] = useState(false);
  
  // Use internal loading state or external prop
  const isAILoading = externalAILoading || internalAILoading;

  // Edit state for annotation list
  const [editingId, setEditingId] = useState<string | null>(null);
  const [editLabel, setEditLabel] = useState('');
  const [editColor, setEditColor] = useState('#05998c');
  
  // Annotation metadata state
  const [_annotationIdentified, setAnnotationIdentified] = useState<Record<string, boolean>>({});
  const [annotationAccepted, setAnnotationAccepted] = useState<Record<string, boolean>>({});
  // Predefined label options
  const labelOptions = [
    'Cervix',
    'SCJ',
    'OS',
    'TZ',
    'Blood Discharge',
    'Scarring',
    'Growth',
    'Ulcer',
    'Inflammation',
    'Polypoidal Growth',
    'Cauliflower-like Growth',
    'Fungating Mass'
  ];

  // Filter labels based on input
  const filteredLabels = labelOptions.filter(label =>
    label.toLowerCase().includes(labelInput.toLowerCase())
  );

  // Expose annotator methods via ref
  useImperativeHandle(ref, () => ({
    addAIAnnotations: (aiAnnotations: Annotation[]) => {
      setAnnotations(aiAnnotations);
    },
    setImageIndex: (index: number) => {
      setSelectedImageIndex(index);
    },
    clearAIAnnotations: () => {
      setAnnotations(prev => prev.filter(ann => ann.source !== 'ai'));
    },
    clearAllAnnotations: () => {
      setAnnotations([]);
      setPolygonPoints([]);
      setCurrentAnnotation(null);
    },
    resetViewport: () => {
      setSelectedImageIndex(0);
      setAnnotations([]);
      setPolygonPoints([]);
      setCurrentAnnotation(null);
    },
    waitForImageReady: () => {
      return new Promise<void>((resolve) => {
        if (imageLoaded) {
          resolve();
        } else {
          imageReadyResolveRef.current = resolve;
        }
      });
    }
  }));

  const images = imageUrls || (imageUrl ? [imageUrl] : []);
  const currentImageUrl = images[selectedImageIndex];

  // Clear annotations when switching to a different image
  useEffect(() => {
    console.log('🖼️ Image switched, clearing annotations');
    setAnnotations([]);
    setPolygonPoints([]);
    setCurrentAnnotation(null);
    setImageLoaded(false);
  }, [selectedImageIndex]);

  useEffect(() => {
    const img = new Image();
    img.src = currentImageUrl;
    img.onload = () => {
      console.log('✅ Image loaded:', { width: img.width, height: img.height });
      setImageDimensions({ width: img.width, height: img.height });
      setImageLoaded(true);
      // Resolve waitForImageReady promise if pending
      if (imageReadyResolveRef.current) {
        console.log('🔔 Image ready resolver called');
        imageReadyResolveRef.current();
        imageReadyResolveRef.current = null;
      }
      // Force immediate canvas redraw with new dimensions
      setTimeout(() => drawCanvas(), 0);
    };
  }, [currentImageUrl]);

  useEffect(() => {
    if (imageLoaded) drawCanvas();
  }, [annotations, currentAnnotation, polygonPoints, imageLoaded]);

  useEffect(() => {
    if (onAnnotationsChange) onAnnotationsChange(annotations);
  }, [annotations, onAnnotationsChange]);

  // Cleanup pending draw on unmount
  useEffect(() => {
    return () => {
      if (pendingDrawAnimationRef.current) {
        cancelAnimationFrame(pendingDrawAnimationRef.current);
      }
    };
  }, []);

  const clearAnnotations = () => setAnnotations([]);
  const deleteLastAnnotation = () => setAnnotations(prev => prev.slice(0, -1));

  const getCanvasCoordinates = (e: React.MouseEvent<HTMLCanvasElement>) => {
    const canvas = canvasRef.current;
    if (!canvas) return { x: 0, y: 0 };
    const rect = canvas.getBoundingClientRect();
    const scaleX = imageDimensions.width / canvas.width || 1;
    const scaleY = imageDimensions.height / canvas.height || 1;
    return { x: (e.clientX - rect.left) * scaleX, y: (e.clientY - rect.top) * scaleY };
  };

  const startDrawing = (point: Point) => {
    setStartPoint(point);
    setIsDrawing(true);
  };

  const finishDrawingRectOrCircle = () => {
    if (currentAnnotation && (currentAnnotation.width > 5 || currentAnnotation.height > 5)) {
      const ann: Annotation = { ...currentAnnotation, id: Date.now().toString(), label: labelInput, source: 'manual', identified: true };
      setAnnotations(prev => [...prev, ann]);
    }
    setIsDrawing(false);
    setStartPoint(null);
    setCurrentAnnotation(null);
  };

  const handleMouseDown = (e: React.MouseEvent<HTMLCanvasElement>) => {
    const p = getCanvasCoordinates(e);
    if (tool === 'polygon') {
      // add point to polygon
      setPolygonPoints(prev => [...prev, p]);
      return;
    }
    startDrawing(p);
  };

  const handleMouseMove = (e: React.MouseEvent<HTMLCanvasElement>) => {
    if (tool === 'polygon') return; // polygon preview handled by polygonPoints
    if (!isDrawing || !startPoint) return;
    const p = getCanvasCoordinates(e);
    const w = p.x - startPoint.x;
    const h = p.y - startPoint.y;
    const ann: Annotation = {
      id: 'temp',
      type: tool,
      x: w > 0 ? startPoint.x : p.x,
      y: h > 0 ? startPoint.y : p.y,
      width: Math.abs(w),
      height: Math.abs(h),
      color,
      label: labelInput
    };
    setCurrentAnnotation(ann);
  };

  const handleMouseUp = () => {
    if (tool === 'polygon') return;
    if (isDrawing) finishDrawingRectOrCircle();
  };

  const finishPolygon = () => {
    if (polygonPoints.length < 3) return;
    const bounds = getBoundsFromPoints(polygonPoints);
    const ann: Annotation = {
      id: Date.now().toString(),
      type: 'polygon',
      x: bounds.x,
      y: bounds.y,
      width: bounds.width,
      height: bounds.height,
      color,
      label: labelInput,
      points: polygonPoints,
      source: 'manual',
      identified: true
    };
    setAnnotations(prev => [...prev, ann]);
    setPolygonPoints([]);
    setLabelInput('');
  };

  const cancelPolygon = () => setPolygonPoints([]);

  const getBoundsFromPoints = (pts: Point[]) => {
    const xs = pts.map(p => p.x);
    const ys = pts.map(p => p.y);
    const minX = Math.min(...xs);
    const minY = Math.min(...ys);
    const maxX = Math.max(...xs);
    const maxY = Math.max(...ys);
    return { x: minX, y: minY, width: maxX - minX, height: maxY - minY };
  };

  const drawCanvas = () => {
    const canvas = canvasRef.current;
    const container = containerRef.current;
    if (!canvas || !container) return;
    
    // Cancel previous animation frame if pending
    if (pendingDrawAnimationRef.current) {
      cancelAnimationFrame(pendingDrawAnimationRef.current);
    }
    
    // Use requestAnimationFrame for smooth, optimized drawing
    pendingDrawAnimationRef.current = requestAnimationFrame(() => {
      performDraw();
    });
  };
  
  const performDraw = () => {
    const canvas = canvasRef.current;
    const container = containerRef.current;
    if (!canvas || !container) return;
    
    const ctx = canvas.getContext('2d');
    if (!ctx) return;
    
    // Create state key to prevent redundant draws
    const stateKey = `${selectedImageIndex}-${annotations.length}-${polygonPoints.length}-${currentAnnotation?.id}`;
    if (lastDrawStateRef.current === stateKey && canvas.width > 0) {
      return; // Skip if nothing changed
    }
    lastDrawStateRef.current = stateKey;
    
    // Get or create cached image
    let img = cachedImagesRef.current[currentImageUrl];
    if (!img) {
      img = new Image();
      img.src = currentImageUrl;
      cachedImagesRef.current[currentImageUrl] = img;
    }
    
    // If image not loaded yet, wait for it
    if (!img.complete) {
      img.onload = () => performDraw();
      return;
    }
    
    // Resize canvas if needed
    const containerWidth = container.clientWidth;
    const aspectRatio = img.height / img.width || 1;
    const canvasHeight = containerWidth * aspectRatio;
    
    if (canvas.width !== containerWidth || canvas.height !== canvasHeight) {
      canvas.width = containerWidth;
      canvas.height = canvasHeight;
    }
    
    // Draw image
    ctx.clearRect(0, 0, canvas.width, canvas.height);
    ctx.drawImage(img, 0, 0, canvas.width, canvas.height);

    // Draw saved annotations
    annotations.forEach(a => drawAnnotation(ctx, a, canvas.width, canvas.height));

    // Draw polygon being created
    if (polygonPoints.length > 0) {
      drawPreviewPolygon(ctx, polygonPoints, canvas.width, canvas.height, color);
    }

    // Draw current temp annotation
    if (currentAnnotation) drawAnnotation(ctx, currentAnnotation, canvas.width, canvas.height);
    
    pendingDrawAnimationRef.current = null;
  };

  const drawPreviewPolygon = (ctx: CanvasRenderingContext2D, pts: Point[], canvasWidth: number, canvasHeight: number, col: string) => {
    const scaleX = canvasWidth / imageDimensions.width || 1;
    const scaleY = canvasHeight / imageDimensions.height || 1;
    ctx.strokeStyle = col;
    ctx.lineWidth = 2;
    ctx.setLineDash([6, 4]);
    ctx.beginPath();
    pts.forEach((p, i) => {
      const x = p.x * scaleX;
      const y = p.y * scaleY;
      if (i === 0) ctx.moveTo(x, y); else ctx.lineTo(x, y);
    });
    ctx.stroke();
    ctx.setLineDash([]);
  };

  const drawAnnotation = (ctx: CanvasRenderingContext2D, annotation: Annotation, canvasWidth: number, canvasHeight: number) => {
    const scaleX = canvasWidth / imageDimensions.width || 1;
    const scaleY = canvasHeight / imageDimensions.height || 1;
    ctx.strokeStyle = annotation.color;
    ctx.lineWidth = 3;
    ctx.setLineDash([]);
    if (annotation.type === 'rect') {
      ctx.strokeRect(annotation.x * scaleX, annotation.y * scaleY, annotation.width * scaleX, annotation.height * scaleY);
      // Only fill for manual annotations, not for AI-detected bounding boxes
      if (annotation.source !== 'ai') {
        ctx.fillStyle = annotation.color + '20';
        ctx.fillRect(annotation.x * scaleX, annotation.y * scaleY, annotation.width * scaleX, annotation.height * scaleY);
      }
    } else if (annotation.type === 'circle') {
      const cx = (annotation.x + annotation.width / 2) * scaleX;
      const cy = (annotation.y + annotation.height / 2) * scaleY;
      const r = Math.max(annotation.width * scaleX, annotation.height * scaleY) / 2;
      ctx.beginPath();
      ctx.arc(cx, cy, r, 0, Math.PI * 2);
      ctx.stroke();
      // Only fill for manual annotations, not for AI-detected shapes
      if (annotation.source !== 'ai') {
        ctx.fillStyle = annotation.color + '20';
        ctx.fill();
      }
    } else if (annotation.type === 'polygon' && annotation.points && Array.isArray(annotation.points) && annotation.points.length > 0) {
      ctx.beginPath();
      annotation.points.forEach((p, i) => {
        if (p && typeof p.x === 'number' && typeof p.y === 'number') {
          const x = p.x * scaleX;
          const y = p.y * scaleY;
          if (i === 0) ctx.moveTo(x, y); else ctx.lineTo(x, y);
        }
      });
      ctx.closePath();
      ctx.stroke();
      // Only fill for manual annotations, not for AI-detected shapes
      if (annotation.source !== 'ai') {
        ctx.fillStyle = annotation.color + '20';
        ctx.fill();
      }
    }
  };

  // list editing
  const startEdit = (ann: Annotation) => {
    setEditingId(ann.id);
    setEditLabel(ann.label || '');
    setEditColor(ann.color || '#05998c');
  };
  const saveEdit = () => {
    setAnnotations(prev => prev.map(a => a.id === editingId ? { ...a, label: editLabel, color: editColor } : a));
    setEditingId(null);
  };
  const deleteAnnotation = (id: string) => setAnnotations(prev => prev.filter(a => a.id !== id));

  const handleAIAssistToggle = async () => {
    if (isAILoading) return;
    if (isAIAssistEnabled) {
      setIsAIAssistEnabled(false);
      return;
    }
    setIsAIAssistEnabled(true);
    if (onAIAssist) {
      setInternalAILoading(true);
      try {
        await onAIAssist();
      } catch (error) {
        console.error('AI Assist error:', error);
      } finally {
        setInternalAILoading(false);
      }
    }
  };

  const getShapeTypeName = (type: ShapeType): string => {
    const typeMap: Record<ShapeType, string> = {
      'rect': 'Rectangle',
      'circle': 'Circle',
      'polygon': 'Polygon'
    };
    return typeMap[type] || type;
  };

  return (
    <div className="space-y-3 md:space-y-4">
      {/* Image Selection - Show if multiple images */}
      {images.length > 1 && (
        <div>
          <label className="block text-xs font-semibold text-gray-500 uppercase mb-2 md:mb-3">
            Select Image
          </label>
          <div className="flex gap-2 overflow-x-auto pb-2">
            {images.map((imgUrl, idx) => (
              <button
                key={idx}
                onClick={() => setSelectedImageIndex(idx)}
                className={`relative flex-shrink-0 w-16 h-16 md:w-20 md:h-20 rounded-lg overflow-hidden border-2 transition-all ${
                  selectedImageIndex === idx
                    ? 'border-[#05998c] ring-2 ring-[#05998c]/50'
                    : 'border-gray-300'
                }`}
              >
                <img
                  src={imgUrl}
                  alt={`Image ${idx + 1}`}
                  className="w-full h-full object-cover"
                />
                {/* Grey overlay for selected image */}
                {selectedImageIndex === idx && (
                  <div className="absolute inset-0 bg-black/30 flex items-center justify-center">
                    <div className="w-5 h-5 rounded-full bg-[#05998c] flex items-center justify-center">
                      <div className="w-2 h-2 bg-white rounded-full" />
                    </div>
                  </div>
                )}
              </button>
            ))}
          </div>
        </div>
      )}

      <div className="flex flex-col md:flex-row md:items-center md:justify-between gap-2 md:gap-3">
        <div className="flex items-center gap-2">
          <div className="flex items-center gap-2 text-xs md:text-sm text-gray-600 bg-blue-50 px-3 py-1.5 rounded-lg border border-blue-100 whitespace-nowrap">
            <Wrench className="w-4 h-4 text-[#05998c]" />
            <span className="font-medium">Tools</span>
          </div>
          <div className="flex items-center gap-1 md:gap-2">
            <button onClick={() => setTool('rect')} className={`px-2 md:px-3 py-1 rounded text-xs md:text-sm ${tool === 'rect' ? 'bg-gray-800 text-white' : 'bg-white text-gray-700'} border`}><SquareIcon className="inline w-4 h-4" /></button>
            <button onClick={() => setTool('circle')} className={`px-2 md:px-3 py-1 rounded text-xs md:text-sm ${tool === 'circle' ? 'bg-gray-800 text-white' : 'bg-white text-gray-700'} border`}><CircleIcon className="inline w-4 h-4" /></button>
            <button onClick={() => setTool('polygon')} className={`px-2 md:px-3 py-1 rounded text-xs md:text-sm ${tool === 'polygon' ? 'bg-gray-800 text-white' : 'bg-white text-gray-700'} border`}><Hexagon className="inline w-4 h-4" /></button>
          </div>
        </div>

        <div className="flex flex-col gap-2">
          {/* First row: AI Assist on right side */}
          <div className="flex justify-end">
            {onAIAssist && (
              <button 
                onClick={handleAIAssistToggle} 
                disabled={isAILoading} 
                className={`px-4 md:px-6 py-2 md:py-3 text-sm md:text-base font-bold text-white rounded-lg disabled:opacity-50 disabled:cursor-not-allowed transition-all shadow-lg hover:shadow-xl flex items-center justify-center gap-2 ${
                  isAIAssistEnabled
                    ? 'bg-gradient-to-r from-green-500 to-green-600 border border-green-600 hover:from-green-600 hover:to-green-700'
                    : 'bg-gradient-to-r from-blue-600 to-blue-700 border border-blue-700 hover:from-blue-700 hover:to-blue-800'
                }`}
                title={isAIAssistEnabled ? 'AI Assist is ON' : 'Run AI model to automatically detect annotations'}
              >
                {isAILoading ? (
                  <>
                    <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>
                    <Loader className="w-5 h-5 animate-spin" />
                    <span>Analyzing...</span>
                  </>
                ) : (
                  <>
                    <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="h-5 w-5" />
                    <span>{isAIAssistEnabled ? 'AI Assist On' : 'AI Assist'}</span>
                  </>
                )}
              </button>
            )}
          </div>

          {/* Second row: Label, Color, Undo on left, Clear All on right */}
          <div className="flex flex-col md:flex-row md:items-center gap-2">
            <div className="relative">
              <input 
                aria-label="Annotation label" 
                placeholder="Search or select label" 
                value={labelInput} 
                onChange={e => setLabelInput(e.target.value)}
                onFocus={() => setIsLabelDropdownOpen(true)}
                onBlur={() => setTimeout(() => setIsLabelDropdownOpen(false), 200)}
                className="px-3 py-1 border rounded text-xs md:text-sm w-48" 
              />
              {isLabelDropdownOpen && filteredLabels.length > 0 && (
                <div className="absolute top-full left-0 mt-1 w-48 bg-white border border-gray-300 rounded-lg shadow-lg max-h-60 overflow-y-auto z-50">
                  {filteredLabels.map((label, idx) => (
                    <button
                      key={idx}
                      type="button"
                      onMouseDown={(e) => {
                        e.preventDefault();
                        setLabelInput(label);
                        setIsLabelDropdownOpen(false);
                      }}
                      className="w-full text-left px-3 py-2 text-xs md:text-sm hover:bg-gray-100 transition-colors"
                    >
                      {label}
                    </button>
                  ))}
                </div>
              )}
            </div>
            <input aria-label="Annotation color" type="color" value={color} onChange={e => setColor(e.target.value)} className="w-10 h-8 p-0 border rounded" />
            <button onClick={deleteLastAnnotation} disabled={annotations.length === 0} className="px-3 md:px-4 py-1 md:py-2 text-xs md:text-sm font-medium text-gray-600 bg-white border border-gray-200 rounded-lg hover:bg-gray-50 disabled:opacity-50 disabled:cursor-not-allowed transition-colors flex items-center justify-center gap-1 md:gap-2">
              <Trash2 className="w-4 h-4" />
              <span className="hidden md:inline">Undo</span>
              <span className="inline md:hidden">Undo</span>
            </button>
            <div className="flex-1"></div>
            <button onClick={clearAnnotations} disabled={annotations.length === 0} className="px-3 md:px-4 py-1 md:py-2 text-xs md:text-sm font-medium text-red-600 bg-white border border-red-200 rounded-lg hover:bg-red-50 disabled:opacity-50 disabled:cursor-not-allowed transition-colors">Clear All</button>
          </div>
        </div>
      </div>

      <div ref={containerRef} className="relative bg-gray-900 rounded-xl overflow-hidden shadow-2xl border-2 border-gray-700">
        <canvas ref={canvasRef} onMouseDown={handleMouseDown} onMouseMove={handleMouseMove} onMouseUp={handleMouseUp} onMouseLeave={handleMouseUp} className="w-full cursor-crosshair" />
        {/* show guide */}
        {annotations.length === 0 && polygonPoints.length === 0 && !isDrawing && (
          <div className="absolute inset-0 flex items-center justify-center pointer-events-none text-white/70 text-xs md:text-sm px-4">
        
          </div>
        )}
      </div>

      {tool === 'polygon' && (
        <div className="flex flex-col md:flex-row md:items-center gap-2">
          <span className="text-xs md:text-sm text-gray-600">Polygon points: {polygonPoints.length}</span>
          <button onClick={finishPolygon} disabled={polygonPoints.length < 3} className="px-3 py-1 text-xs md:text-sm bg-green-600 text-white rounded disabled:opacity-50">Finish Polygon</button>
          <button onClick={cancelPolygon} disabled={polygonPoints.length === 0} className="px-3 py-1 text-xs md:text-sm bg-gray-200 rounded">Cancel</button>
        </div>
      )}

      {/* Annotations Table */}
      <div className="border border-gray-200 rounded-lg overflow-hidden">
        <button
          onClick={() => setIsAnnotationsOpen(!isAnnotationsOpen)}
          className="w-full flex items-center justify-between p-3 md:p-4 bg-gray-50 hover:bg-gray-100 transition-colors"
        >
          <div className="flex items-center gap-2">
            <span className="text-xs md:text-sm font-semibold text-gray-700">
              Annotations ({annotations.length})
            </span>
          </div>
          {isAnnotationsOpen ? (
            <ChevronUp className="w-4 h-4 text-gray-600" />
          ) : (
            <ChevronDown className="w-4 h-4 text-gray-600" />
          )}
        </button>

        {isAnnotationsOpen && (
          <div className="overflow-x-auto border-t border-gray-200 max-h-96 overflow-y-auto">
            {annotations.length === 0 ? (
              <div className="p-3 md:p-4 bg-white text-center">
                <p className="text-xs text-gray-500">No annotations yet. Draw on the image to create annotations.</p>
              </div>
            ) : (
              <table className="w-full">
                <thead className="bg-gray-50 sticky top-0 border-b border-gray-200">
                  <tr>
                    <th className="px-3 md:px-4 py-2 text-left text-xs font-semibold text-gray-700">Annotation</th>
                    <th className="px-3 md:px-4 py-2 text-center text-xs font-semibold text-gray-700">Identified</th>
                    <th className="px-3 md:px-4 py-2 text-center text-xs font-semibold text-gray-700">Source</th>
                    <th className="px-3 md:px-4 py-2 text-center text-xs font-semibold text-gray-700">Action</th>
                  </tr>
                </thead>
                <tbody>
                  {annotations.map((a) => (
                    <tr key={a.id} className="border-b border-gray-100 hover:bg-gray-50 transition-colors">
                      {/* Annotation Column - shown in color it was drawn */}
                      <td className="px-3 md:px-4 py-3">
                        {editingId === a.id ? (
                          <div className="flex items-center gap-2">
                            <input 
                              type="color" 
                              value={editColor} 
                              onChange={e => setEditColor(e.target.value)} 
                              className="w-6 h-6 rounded p-0 border cursor-pointer" 
                            />
                            <input 
                              value={editLabel} 
                              onChange={e => setEditLabel(e.target.value)} 
                              className="px-2 py-1 border rounded text-xs flex-1" 
                              placeholder="Label"
                            />
                          </div>
                        ) : (
                          <div className="flex items-center gap-3">
                            <div className="w-6 h-6 rounded-sm flex-shrink-0" style={{ backgroundColor: a.color }} />
                            <div>
                              <div className="text-sm font-medium text-gray-900">{a.label || '(no label)'}</div>
                              <div className="text-xs text-gray-500">{getShapeTypeName(a.type)}</div>
                            </div>
                          </div>
                        )}
                      </td>

                      {/* Identified Checkbox Column */}
                      <td className="px-3 md:px-4 py-3 text-center">
                        <input
                          type="checkbox"
                          checked={a.identified || false}
                          onChange={(e) => {
                            setAnnotations(prev => prev.map(ann => 
                              ann.id === a.id ? { ...ann, identified: e.target.checked } : ann
                            ));
                          }}
                          className="w-4 h-4 rounded border-gray-300 cursor-pointer"
                        />
                      </td>

                      {/* Source Column (Manual/AI) */}
                      <td className="px-3 md:px-4 py-3 text-center">
                        <span className={`inline-block px-2 py-1 rounded text-xs font-medium ${
                          a.source === 'ai' 
                            ? 'bg-blue-100 text-blue-800' 
                            : 'bg-gray-100 text-gray-800'
                        }`}>
                          {a.source === 'ai' ? 'AI' : 'Manual'}
                        </span>
                      </td>

                      {/* Action Column - Accept/Reject for AI, Edit/Delete for Manual */}
                      <td className="px-3 md:px-4 py-3">
                        <div className="flex items-center justify-center gap-2">
                          {editingId === a.id ? (
                            <>
                              <button
                                onClick={saveEdit}
                                className="px-2 py-1 text-xs font-medium bg-blue-600 text-white rounded hover:bg-blue-700 transition-colors"
                              >
                                Save
                              </button>
                              <button
                                onClick={() => setEditingId(null)}
                                className="px-2 py-1 text-xs font-medium bg-gray-300 text-gray-800 rounded hover:bg-gray-400 transition-colors"
                              >
                                Cancel
                              </button>
                            </>
                          ) : a.source === 'ai' ? (
                            <>
                              {annotationAccepted[a.id] === true ? (
                                <div className="flex items-center gap-2">
                                  <span className="px-3 py-1 text-xs font-medium bg-green-100 text-green-700 rounded-full">
                                    ✓ Accepted
                                  </span>
                                  <button
                                    onClick={() => {
                                      setAnnotationAccepted(prev => {
                                        const next = { ...prev };
                                        delete next[a.id];
                                        return next;
                                      });
                                      setAnnotationIdentified(prev => {
                                        const next = { ...prev };
                                        delete next[a.id];
                                        return next;
                                      });
                                    }}
                                    className="px-2 py-1 text-[11px] font-medium bg-white border border-gray-200 text-gray-700 rounded hover:bg-gray-50 transition-colors"
                                  >
                                    Undo
                                  </button>
                                </div>
                              ) : annotationAccepted[a.id] === false ? (
                                <span className="px-3 py-1 text-xs font-medium bg-red-100 text-red-700 rounded-full">
                                  ✕ Rejected
                                </span>
                              ) : (
                                <>
                                  <button
                                    onClick={() => {
                                      setAnnotations(prev => prev.map(ann =>
                                        ann.id === a.id ? { ...ann, identified: true } : ann
                                      ));
                                      setAnnotationAccepted(prev => ({
                                        ...prev,
                                        [a.id]: true
                                      }));
                                      setAnnotationIdentified(prev => ({
                                        ...prev,
                                        [a.id]: true
                                      }));
                                    }}
                                    className="px-2 py-1 text-xs font-medium bg-green-50 text-green-700 border border-green-200 rounded hover:bg-green-100 transition-colors"
                                  >
                                    ✓ Accept
                                  </button>
                                  <button
                                    onClick={() => {
                                      setAnnotations(prev => prev.filter(item => item.id !== a.id));
                                      setAnnotationAccepted(prev => {
                                        const next = { ...prev };
                                        delete next[a.id];
                                        return next;
                                      });
                                      setAnnotationIdentified(prev => {
                                        const next = { ...prev };
                                        delete next[a.id];
                                        return next;
                                      });
                                    }}
                                    className="px-2 py-1 text-xs font-medium bg-red-50 text-red-700 border border-red-200 rounded hover:bg-red-100 transition-colors"
                                  >
                                    ✕ Reject
                                  </button>
                                </>
                              )}
                            </>
                          ) : (
                            <>
                              <button
                                onClick={() => startEdit(a)}
                                className="px-2 py-1 text-xs font-medium bg-white border border-gray-300 rounded hover:bg-gray-50 transition-colors"
                              >
                                Edit
                              </button>
                              <button
                                onClick={() => deleteAnnotation(a.id)}
                                className="px-2 py-1 text-xs font-medium bg-red-50 text-red-700 border border-red-200 rounded hover:bg-red-100 transition-colors"
                              >
                                Delete
                              </button>
                            </>
                          )}
                        </div>
                      </td>
                    </tr>
                  ))}
                </tbody>
              </table>
            )}
          </div>
        )}
      </div>
    </div>
  );
});

ImageAnnotatorComponent.displayName = 'ImageAnnotator';

export const ImageAnnotator = ImageAnnotatorComponent;