src/components/AceticAnnotator.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 AceticAnnotatorProps {
  imageUrl?: string;
  imageUrls?: string[];
  onAnnotationsChange?: (annotations: Annotation[]) => void;
}

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

const AceticAnnotatorComponent = forwardRef<AceticAnnotatorHandle, AceticAnnotatorProps>(({ imageUrl, imageUrls, onAnnotationsChange }, ref) => {
  const canvasRef = useRef<HTMLCanvasElement | null>(null);
  const containerRef = useRef<HTMLDivElement | null>(null);
  const imageReadyResolveRef = useRef<(() => void) | null>(null);
  const emptyAnnotationsRef = useRef<Annotation[]>([]);
  const cachedImagesRef = useRef<Record<string, HTMLImageElement>>({});
  const pendingDrawAnimationRef = useRef<number | null>(null);
  const lastDrawStateRef = useRef<string>('');
  
  // Store annotations per image index
  const [annotationsByImage, setAnnotationsByImage] = useState<Record<number, 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);

  // 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>>({});
  
  // AI state
  const [isAILoading, setIsAILoading] = useState(false);
  const [aiError, setAIError] = useState<string | null>(null);
  const [isAIAssistEnabled, setIsAIAssistEnabled] = useState(false);
  
  // Get annotations for current image
  const annotations = annotationsByImage[selectedImageIndex] ?? emptyAnnotationsRef.current;
  
  // Helper to update annotations for current image
  const setAnnotations = (updater: Annotation[] | ((prev: Annotation[]) => Annotation[])) => {
    setAnnotationsByImage(prev => {
      const currentAnnotations = prev[selectedImageIndex] || [];
      const newAnnotations = typeof updater === 'function' ? updater(currentAnnotations) : updater;
      return {
        ...prev,
        [selectedImageIndex]: newAnnotations
      };
    });
  };
  
  // Predefined label options for Acetic Acid step (ONLY CHANGE from ImageAnnotator)
  const labelOptions = [
    'Acetowhite',
    'Mosaic',
    'Punctation',
    'Blood Vessel',
    'Border'
  ];



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



  // Helper function to get bounding box from polygon points
  const getBoundsFromPoints = (points: any[]) => {
    if (!points || !Array.isArray(points) || points.length === 0) {
      console.warn('⚠️ Invalid points array in getBoundsFromPoints:', points);
      return { x: 0, y: 0, width: 0, height: 0 };
    }
    
    // Handle both array [x, y] and object {x, y} formats
    const getX = (p: any) => Array.isArray(p) ? p[0] : (p?.x ?? 0);
    const getY = (p: any) => Array.isArray(p) ? p[1] : (p?.y ?? 0);
    
    const xs = points.map(getX).filter(x => typeof x === 'number' && !isNaN(x));
    const ys = points.map(getY).filter(y => typeof y === 'number' && !isNaN(y));
    
    if (xs.length === 0 || ys.length === 0) {
      console.warn('⚠️ No valid coordinates in points:', points);
      return { x: 0, y: 0, width: 0, height: 0 };
    }
    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 };
  };

  // AI Assist - Run acetowhite contour detection on current image
  const runAIAssist = async () => {
    const images = imageUrls || (imageUrl ? [imageUrl] : []);
    if (images.length === 0) {
      setAIError('No images available to analyze');
      return;
    }

    const currentImageUrl = images[selectedImageIndex];

    setIsAILoading(true);
    setAIError(null);

    try {
      // Ensure image is loaded first
      if (!imageLoaded || imageDimensions.width === 0 || imageDimensions.height === 0) {
        console.log('⏳ Waiting for image to load before AI processing...');
        await new Promise<void>((resolve) => {
          imageReadyResolveRef.current = resolve;
        });
      }

      const response = await fetch(currentImageUrl);
      const blob = await response.blob();

      const formData = new FormData();
      formData.append('file', blob, 'image.jpg');
      formData.append('conf_threshold', '0.4');

      console.log('🚀 Sending to backend with image dimensions:', imageDimensions);

      const backendResponse = await fetch('/api/infer-aw-contour', {
        method: 'POST',
        body: formData,
      });

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

      const result = await backendResponse.json();

      console.log('📥 Backend response received');

      if (result.status !== 'success') {
        throw new Error(result.error || 'AI inference failed');
      }

      if (!result.contours || !Array.isArray(result.contours)) {
        console.warn('⚠️ Invalid contours format:', result.contours);
        setAIError('Invalid response format from server');
        setIsAILoading(false);
        return;
      }

      const contours = result.contours;
      if (contours.length === 0) {
        console.log('ℹ️ No contours detected');
        setAIError('No contours detected for this image');
        setIsAILoading(false);
        return;
      }

      console.log(`✅ Processing ${contours.length} contour(s)`);


      const aiAnnotations = contours
        .filter((contour: any) => contour !== null && contour !== undefined)
        .map((contour: any, idx: number) => {
          let points: { x: number; y: number }[];
          let area: number | undefined;
          let confidence: number | undefined;

          if (!contour) {
            console.warn(`⚠️ Contour ${idx} is null or undefined`);
            return null;
          }

          // Handle two formats:
          // 1. OLD: contour is an array of [x, y] pairs: [[x,y], [x,y], ...]
          // 2. NEW: contour is an object: {points: [[x,y], ...], area: ..., confidence: ...}
          
          if (Array.isArray(contour)) {
            // OLD FORMAT: contour is directly an array of points
            try {
              // Use points directly as-is from backend
              points = contour as any[];
              
              // Handle OpenCV format: [[[x, y]], [[x, y]], ...] (nested arrays)
              if (Array.isArray(contour[0]) && Array.isArray(contour[0][0])) {
                points = contour.map((p: any) => {
                  if (Array.isArray(p) && Array.isArray(p[0])) {
                    return p[0]; // Extract [x, y] from [[x, y]]
                  }
                  return p;
                });
              }
            } catch (e) {
              console.warn(`⚠️ Error processing contour ${idx}:`, e);
              return null;
            }
            area = undefined;
            confidence = undefined;
          } else if (contour && typeof contour === 'object' && contour.points && Array.isArray(contour.points)) {
            // NEW FORMAT: contour is an object with points property
            try {
              // Use points directly as-is from backend
              points = contour.points as any[];
              
              // Handle OpenCV format: [[[x, y]], [[x, y]], ...] (nested arrays)
              if (Array.isArray(contour.points[0]) && Array.isArray(contour.points[0][0])) {
                points = contour.points.map((p: any) => {
                  if (Array.isArray(p) && Array.isArray(p[0])) {
                    return p[0]; // Extract [x, y] from [[x, y]]
                  }
                  return p;
                });
              }
            } catch (e) {
              console.warn(`⚠️ Error processing contour points ${idx}:`, e);
              return null;
            }
            area = contour.area;
            confidence = contour.confidence;
          } else {
            console.warn(`⚠️ Invalid contour format at index ${idx}:`, contour);
            return null;
          }

          if (!points || points.length < 3) {
            console.warn(`⚠️ Contour ${idx} has insufficient points (${points?.length || 0})`);
            return null;
          }

          // Normalize point scale if backend returns normalized (0-1) or percent (0-100) coordinates
          let scaledPoints = points;
          
          if (imageDimensions.width > 0 && imageDimensions.height > 0) {
            // Helper to extract x from various formats
            const getX = (p: any) => {
              if (Array.isArray(p) && typeof p[0] === 'number') return p[0];
              if (Array.isArray(p) && Array.isArray(p[0])) return p[0][0];
              if (typeof p === 'object' && p !== null) return p.x;
              return undefined;
            };
            
            // Helper to extract y from various formats
            const getY = (p: any) => {
              if (Array.isArray(p) && typeof p[0] === 'number') return p[1];
              if (Array.isArray(p) && Array.isArray(p[0])) return p[0][1];
              if (typeof p === 'object' && p !== null) return p.y;
              return undefined;
            };
            
            const xValues = points.map(getX).filter(x => typeof x === 'number' && !isNaN(x));
            const yValues = points.map(getY).filter(y => typeof y === 'number' && !isNaN(y));
            
            if (xValues.length === 0 || yValues.length === 0) {
              // Fallback: assume [x, y] or [[x, y]] array format
              scaledPoints = points.map((p) => {
                let x, y;
                if (Array.isArray(p) && Array.isArray(p[0])) {
                  x = Number(p[0][0]);
                  y = Number(p[0][1]);
                } else if (Array.isArray(p) && p.length >= 2) {
                  x = Number(p[0]);
                  y = Number(p[1]);
                } else {
                  x = NaN;
                  y = NaN;
                }
                return { x, y };
              });
            } else {
              const maxX = Math.max(...xValues);
              const maxY = Math.max(...yValues);
              
              if (maxX <= 1 && maxY <= 1) {
                scaledPoints = points.map(p => {
                  const x = getX(p);
                  const y = getY(p);
                  return { x: x * imageDimensions.width, y: y * imageDimensions.height };
                });
              } else if (maxX <= 100 && maxY <= 100) {
                scaledPoints = points.map(p => {
                  const x = getX(p);
                  const y = getY(p);
                  return { x: (x / 100) * imageDimensions.width, y: (y / 100) * imageDimensions.height };
                });
              } else {
                scaledPoints = points.map(p => {
                  const x = getX(p);
                  const y = getY(p);
                  return { x: Number(x), y: Number(y) };
                });
              }
            }
          } else {
            // No image dimensions yet, just convert to {x, y} format
            scaledPoints = points.map(p => {
              let x, y;
              if (Array.isArray(p) && Array.isArray(p[0])) {
                x = Number(p[0][0]);
                y = Number(p[0][1]);
              } else if (Array.isArray(p) && p.length >= 2) {
                x = Number(p[0]);
                y = Number(p[1]);
              } else if (typeof p === 'object' && p !== null) {
                x = Number(p.x);
                y = Number(p.y);
              } else {
                x = NaN;
                y = NaN;
              }
              return { x, y };
            });
          }

          const bounds = getBoundsFromPoints(scaledPoints);
          const annotation = {
            id: `ai_${Date.now()}_${idx}`,
            type: 'polygon' as const,
            x: bounds.x,
            y: bounds.y,
            width: bounds.width,
            height: bounds.height,
            color: '#05998c',
            label: 'Acetowhite',
            points: scaledPoints,
            source: 'ai' as const,
            identified: false,
            accepted: false,
            confidence: confidence,
            area: area
          };
          
          return annotation;
        })
        .filter(Boolean) as Annotation[];


      if (aiAnnotations.length > 0) {
        // Replace all annotations with new AI annotations
        // The useEffect will automatically trigger drawCanvas when annotationsByImage updates
        console.log(`🎨 Setting ${aiAnnotations.length} AI annotations and triggering redraw`);
        setAnnotations(aiAnnotations);
      } else {
        setAIError('No valid contours detected');
      }

      setIsAILoading(false);
    } catch (error) {
      console.error('AI assist error:', error);
      setAIError(error instanceof Error ? error.message : 'Failed to run AI analysis');
      setIsAILoading(false);
    }
  };

  // Expose annotator methods via ref
  useImperativeHandle(ref, () => ({
    addAIAnnotations: (aiAnnotations: Annotation[]) => {
      setAnnotations(aiAnnotations);
    },
    setImageIndex: (index: number) => {
      setSelectedImageIndex(index);
    },
    getCurrentImageIndex: () => {
      return selectedImageIndex;
    },
    clearAIAnnotations: () => {
      setAnnotations(prev => prev.filter(ann => ann.source !== 'ai'));
    },
    clearAllAnnotations: () => {
      setAnnotations([]);
      setPolygonPoints([]);
      setCurrentAnnotation(null);
    },
    resetViewport: () => {
      setSelectedImageIndex(0);
      setAnnotationsByImage({});  // Clear all images' annotations
      setPolygonPoints([]);
      setCurrentAnnotation(null);
    },
    runAIAssist: async () => {
      await runAIAssist();
    },
    waitForImageReady: () => {
      return new Promise<void>((resolve) => {
        if (imageLoaded) {
          console.log('✅ Image already ready');
          resolve();
        } else {
          console.log('⏳ Waiting for image to load');
          imageReadyResolveRef.current = resolve;
        }
      });
    }
  }));

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

  // Reset drawing state when switching images (but keep annotations per-image)
  useEffect(() => {
    
    // Clear temporary drawing state only
    setPolygonPoints([]);
    setCurrentAnnotation(null);
    setAnnotationIdentified({});
    setAnnotationAccepted({});
    setImageLoaded(false);
    
    // Immediately clear the canvas (will be redrawn with correct image's annotations)
    const canvas = canvasRef.current;
    if (canvas) {
      const ctx = canvas.getContext('2d');
      if (ctx) {
        ctx.clearRect(0, 0, canvas.width, canvas.height);
      }
    }
  }, [selectedImageIndex]);

  useEffect(() => {
    const img = new Image();
    img.src = currentImageUrl;
    img.onload = () => {
      setImageDimensions({ width: img.width, height: img.height });
      setImageLoaded(true);
      // Resolve waitForImageReady promise if pending
      if (imageReadyResolveRef.current) {
        imageReadyResolveRef.current();
        imageReadyResolveRef.current = null;
      }
      // drawCanvas will be called by useEffect when imageLoaded changes
    };
  }, [currentImageUrl]);

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

  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: false };
      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: false
    };
    setAnnotations(prev => [...prev, ann]);
    setPolygonPoints([]);
    setLabelInput('');
  };

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

  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 currentAnnotations = annotationsByImage[selectedImageIndex] || [];
    const stateKey = `${selectedImageIndex}-${currentAnnotations.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
    currentAnnotations.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 = imageDimensions.width > 0 ? canvasWidth / imageDimensions.width : 1;
    const scaleY = imageDimensions.height > 0 ? 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);
      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();
      ctx.fillStyle = annotation.color + '20';
      ctx.fill();
    } else if (annotation.type === 'polygon' && annotation.points && Array.isArray(annotation.points) && annotation.points.length > 0) {
      console.log(`🎨 Drawing polygon with ${annotation.points.length} points, scale: ${scaleX}x${scaleY}`);
      
      ctx.beginPath();
      let validPointsCount = 0;
      annotation.points.forEach((p, i) => {
        if (p && typeof p.x === 'number' && typeof p.y === 'number' && Number.isFinite(p.x) && Number.isFinite(p.y)) {
          const x = p.x * scaleX;
          const y = p.y * scaleY;
          if (i === 0) ctx.moveTo(x, y);
          else ctx.lineTo(x, y);
          validPointsCount++;
        } else {
          console.warn(`⚠️ Invalid point at index ${i}:`, p);
        }
      });
      
      if (validPointsCount < 3) {
        console.warn(`⚠️ Polygon has only ${validPointsCount} valid points, skipping draw`);
        return;
      }
      
      ctx.closePath();
      ctx.stroke();
      ctx.fillStyle = annotation.color + '20';
      ctx.fill();
      
      console.log(`✅ Polygon drawn with ${validPointsCount} points`);
    }
  };

  // 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);
    await runAIAssist();
  };

  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>
      )}

      {/* Tab Switcher with AI Assist Button */}
      <div className="flex justify-between items-center">
        <div className="flex justify-center flex-1">
          <div className="inline-flex bg-gray-100 rounded-lg p-1 border border-gray-300">
          </div>
        </div>
        <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>

      {/* Annotate Section */}
      <>
          <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 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>
              <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>
          
          {/* AI Error Message */}
          {aiError && (
            <div className="bg-red-50 border border-red-200 text-red-700 px-3 py-2 rounded-lg text-xs md:text-sm">
              {aiError}
            </div>
          )}

          <div className="flex flex-col">
            <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" />
              </div>
            </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={annotationIdentified[a.id] || false}
                              onChange={(e) => {
                                setAnnotationIdentified(prev => ({
                                  ...prev,
                                  [a.id]: e.target.checked
                                }));
                              }}
                              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>
  );
});

AceticAnnotatorComponent.displayName = 'AceticAnnotator';

export const AceticAnnotator = AceticAnnotatorComponent;
export type { AceticAnnotatorProps, Annotation };