frontend/src/components/viewer/AnnotationCanvas.tsx

import { useCallback, useEffect, useRef, useState } from "react";
import OpenSeadragon from "openseadragon";

export interface Annotation {
  id: string;
  type: "rectangle" | "polygon" | "ellipse" | "brush";
  points: Array<{ x: number; y: number }>;
  color: string;
  label: string;
  completed: boolean;
}

interface AnnotationCanvasProps {
  viewer: OpenSeadragon.Viewer | null;
  tool: "rectangle" | "polygon" | "ellipse" | "brush" | "select" | "none";
  onAnnotationComplete: (annotation: Annotation) => void;
  activeLabel: string;
  onAnnotationSelected?: (annotationId: string | null) => void;
  annotations: Annotation[];
  selectedAnnotationId?: string | null;
  showAnnotations: boolean;
}

export function AnnotationCanvas({
  viewer,
  tool,
  onAnnotationComplete,
  activeLabel,
  onAnnotationSelected,
  annotations,
  selectedAnnotationId,
  showAnnotations,
}: AnnotationCanvasProps) {
  const canvasRef = useRef<HTMLCanvasElement>(null);
  const [isDrawing, setIsDrawing] = useState(false);
  const [currentPoints, setCurrentPoints] = useState<Array<{ x: number; y: number }>>([]);
  const [hoveredAnnotationId, setHoveredAnnotationId] = useState<string | null>(null);
  const [drawingTool, setDrawingTool] = useState<"rectangle" | "polygon" | "ellipse" | "brush" | null>(null);
  const [isPanning, setIsPanning] = useState(false);
  const [lastPanPoint, setLastPanPoint] = useState<{ x: number; y: number } | null>(null);

  const labelColor = (label: string) => {
    switch (label) {
      case "Tumor":
        return "#EF4444";
      case "Benign":
        return "#FACC15";
      case "Stroma":
        return "#EC4899";
      case "Necrosis":
        return "#22C55E";
      case "DCIS":
        return "#3B82F6";
      case "Invasive":
        return "#8B5CF6";
      default:
        return "#9CA3AF";
    }
  };

  const toImagePoint = (point: { x: number; y: number }) => {
    if (!viewer) return point;
    const viewportPoint = viewer.viewport.pointFromPixel(
      new OpenSeadragon.Point(point.x, point.y)
    );
    const imagePoint = viewer.viewport.viewportToImageCoordinates(viewportPoint);
    return { x: imagePoint.x, y: imagePoint.y };
  };

  const toScreenPoint = (point: { x: number; y: number }) => {
    if (!viewer) return point;
    const viewportPoint = viewer.viewport.imageToViewportCoordinates(
      new OpenSeadragon.Point(point.x, point.y)
    );
    const pixelPoint = viewer.viewport.pixelFromPoint(viewportPoint, true);
    return { x: pixelPoint.x, y: pixelPoint.y };
  };

  // Setup canvas
  useEffect(() => {
    const canvas = canvasRef.current;
    if (!canvas || !viewer) return;

    const updateCanvasSize = () => {
      canvas.width = viewer.container.clientWidth;
      canvas.height = viewer.container.clientHeight;
      redrawAnnotations();
    };

    updateCanvasSize();
    
    // Update canvas size when viewer resizes
    const resizeObserver = new ResizeObserver(updateCanvasSize);
    resizeObserver.observe(viewer.container);

    return () => resizeObserver.disconnect();
  }, [viewer]);

  // Helper function to check if point is inside annotation
  const isPointInAnnotation = (point: { x: number; y: number }, annotation: Annotation): boolean => {
    const tolerance = 10; // pixels
    const screenPoints = annotation.points.map(toScreenPoint);
    
    if (annotation.type === "rectangle" && screenPoints.length === 2) {
      const [p1, p2] = screenPoints;
      const minX = Math.min(p1.x, p2.x);
      const maxX = Math.max(p1.x, p2.x);
      const minY = Math.min(p1.y, p2.y);
      const maxY = Math.max(p1.y, p2.y);
      
      return point.x >= minX - tolerance && point.x <= maxX + tolerance &&
             point.y >= minY - tolerance && point.y <= maxY + tolerance;
    } else if (annotation.type === "ellipse" && screenPoints.length === 2) {
      const [p1, p2] = screenPoints;
      const cx = (p1.x + p2.x) / 2;
      const cy = (p1.y + p2.y) / 2;
      const rx = Math.max(1, Math.abs(p2.x - p1.x) / 2);
      const ry = Math.max(1, Math.abs(p2.y - p1.y) / 2);
      const dx = (point.x - cx) / rx;
      const dy = (point.y - cy) / ry;
      const dist = Math.abs(dx * dx + dy * dy - 1);
      const tol = tolerance / Math.max(rx, ry);
      return dist <= tol;
    } else if (annotation.type === "polygon" && screenPoints.length > 2) {
      // Check if point is near any line segment of the polygon
      const points = screenPoints;
      for (let i = 0; i < points.length; i++) {
        const p1 = points[i];
        const p2 = points[(i + 1) % points.length];
        
        const dist = distanceToLineSegment(point, p1, p2);
        if (dist < tolerance) return true;
      }
      // Also check if near any vertex
      return screenPoints.some(p => 
        Math.hypot(p.x - point.x, p.y - point.y) < tolerance
      );
    } else if (annotation.type === "brush" && screenPoints.length > 1) {
      for (let i = 0; i < screenPoints.length - 1; i++) {
        const dist = distanceToLineSegment(point, screenPoints[i], screenPoints[i + 1]);
        if (dist < tolerance) return true;
      }
    }
    return false;
  };

  // Helper function to calculate distance from point to line segment
  const distanceToLineSegment = (
    point: { x: number; y: number },
    p1: { x: number; y: number },
    p2: { x: number; y: number }
  ): number => {
    const dx = p2.x - p1.x;
    const dy = p2.y - p1.y;
    const t = Math.max(0, Math.min(1, ((point.x - p1.x) * dx + (point.y - p1.y) * dy) / (dx * dx + dy * dy)));
    const closestX = p1.x + t * dx;
    const closestY = p1.y + t * dy;
    return Math.hypot(point.x - closestX, point.y - closestY);
  };

  // Handle mouse events for drawing
  useEffect(() => {
    const canvas = canvasRef.current;
    if (!canvas || !viewer) return;

    const handleMouseMove = (e: MouseEvent) => {
      const rect = canvas.getBoundingClientRect();
      const x = e.clientX - rect.left;
      const y = e.clientY - rect.top;
      const point = { x, y };

      // Update hover state based on tool
      if (tool === "select") {
        const hoveredAnnotation = annotations.find(ann => isPointInAnnotation(point, ann));
        setHoveredAnnotationId(hoveredAnnotation?.id || null);
        
        // Change cursor when hovering over annotation
        if (hoveredAnnotation) {
          canvas.style.cursor = "pointer";
        } else {
          canvas.style.cursor = "default";
        }
      } else {
        setHoveredAnnotationId(null);
      }

      // Handle select tool panning
      if (tool === "select" && isPanning && lastPanPoint && viewer) {
        const dx = x - lastPanPoint.x;
        const dy = y - lastPanPoint.y;
        const delta = viewer.viewport.deltaPointsFromPixels(
          new OpenSeadragon.Point(-dx, -dy),
          true
        );
        viewer.viewport.panBy(delta);
        viewer.viewport.applyConstraints();
        setLastPanPoint({ x, y });
        return;
      }

      // Handle drawing preview for rectangle/ellipse tool
      if (isDrawing && (tool === "rectangle" || tool === "ellipse") && currentPoints.length > 0) {
        const imagePoint = toImagePoint(point);
        setCurrentPoints([currentPoints[0], imagePoint]);
      }

      if (isDrawing && tool === "brush" && currentPoints.length > 0) {
        const imagePoint = toImagePoint(point);
        const lastPoint = currentPoints[currentPoints.length - 1];
        if (Math.hypot(imagePoint.x - lastPoint.x, imagePoint.y - lastPoint.y) > 0.5) {
          setCurrentPoints([...currentPoints, imagePoint]);
        }
      }
    };

    const handleMouseDown = (e: MouseEvent) => {
      const rect = canvas.getBoundingClientRect();
      const x = e.clientX - rect.left;
      const y = e.clientY - rect.top;
      const point = { x, y };

      // Allow selection with Ctrl+Click on any annotation (works with any tool)
      if (e.ctrlKey || e.metaKey) {
        const selectedAnnotation = annotations.find(ann => isPointInAnnotation(point, ann));
        if (onAnnotationSelected) {
          onAnnotationSelected(selectedAnnotation?.id || null);
        }
        return;
      }

      // Handle select tool - click to select annotation, drag to pan
      if (tool === "select") {
        const selectedAnnotation = annotations.find(ann => isPointInAnnotation(point, ann));
        if (selectedAnnotation) {
          if (onAnnotationSelected) {
            onAnnotationSelected(selectedAnnotation.id);
          }
          return;
        }

        if (viewer) {
          setIsPanning(true);
          setLastPanPoint({ x, y });
        }
        return;
      }

      // Handle drawing tools
      if (tool === "rectangle") {
        setIsDrawing(true);
        setDrawingTool("rectangle");
        setCurrentPoints([toImagePoint(point)]);
      } else if (tool === "ellipse") {
        setIsDrawing(true);
        setDrawingTool("ellipse");
        setCurrentPoints([toImagePoint(point)]);
      } else if (tool === "polygon") {
        // For polygon, add point on click
        const imagePoint = toImagePoint(point);
        const newPoints = [...currentPoints, imagePoint];
        setCurrentPoints(newPoints);
        if (currentPoints.length === 0) {
          setIsDrawing(true);
          setDrawingTool("polygon");
        }
      } else if (tool === "brush") {
        setIsDrawing(true);
        setDrawingTool("brush");
        setCurrentPoints([toImagePoint(point)]);
      }
    };

    const handleMouseUp = (e: MouseEvent) => {
      if (tool === "select" && isPanning) {
        setIsPanning(false);
        setLastPanPoint(null);
        return;
      }

      if (!isDrawing || (tool !== "rectangle" && tool !== "ellipse" && tool !== "brush")) return;

      const rect = canvas.getBoundingClientRect();
      const x = e.clientX - rect.left;
      const y = e.clientY - rect.top;

      const imagePoint = toImagePoint({ x, y });
      if (tool === "brush") {
        const newPoints = [...currentPoints, imagePoint];
        setCurrentPoints(newPoints);
        completeAnnotation(undefined, newPoints);
        return;
      }

      setCurrentPoints([currentPoints[0], imagePoint]);
      completeAnnotation();
    };

    const handleDblClick = () => {
      if (tool === "polygon" && isDrawing && currentPoints.length >= 3) {
        completeAnnotation();
      }
    };

    const handleWheel = (e: WheelEvent) => {
      if (!viewer) return;
      e.preventDefault();

      const rect = canvas.getBoundingClientRect();
      const x = e.clientX - rect.left;
      const y = e.clientY - rect.top;
      const zoomPoint = viewer.viewport.pointFromPixel(
        new OpenSeadragon.Point(x, y)
      );

      const zoomFactor = e.deltaY < 0 ? 1.1 : 0.9;
      viewer.viewport.zoomBy(zoomFactor, zoomPoint);
      viewer.viewport.applyConstraints();
    };

    const handleKeyDown = (e: KeyboardEvent) => {
      if (e.key === "Escape" && isDrawing) {
        setIsDrawing(false);
        setCurrentPoints([]);
      }
    };

    canvas.addEventListener("mousemove", handleMouseMove);
    canvas.addEventListener("mousedown", handleMouseDown);
    canvas.addEventListener("mouseup", handleMouseUp);
    canvas.addEventListener("dblclick", handleDblClick);
    canvas.addEventListener("wheel", handleWheel, { passive: false });
    document.addEventListener("keydown", handleKeyDown);

    return () => {
      canvas.removeEventListener("mousemove", handleMouseMove);
      canvas.removeEventListener("mousedown", handleMouseDown);
      canvas.removeEventListener("mouseup", handleMouseUp);
      canvas.removeEventListener("dblclick", handleDblClick);
      canvas.removeEventListener("wheel", handleWheel);
      document.removeEventListener("keydown", handleKeyDown);
    };
  }, [isDrawing, currentPoints, tool, viewer, annotations, selectedAnnotationId, onAnnotationSelected, hoveredAnnotationId, isPanning, lastPanPoint]);

  const completeAnnotation = (
    forcedTool?: "rectangle" | "polygon" | "ellipse" | "brush",
    forcedPoints?: Array<{ x: number; y: number }>
  ) => {
    const points = forcedPoints ?? currentPoints;
    if (points.length < 2) {
      setCurrentPoints([]);
      setIsDrawing(false);
      setDrawingTool(null);
      return;
    }

    const finalTool = forcedTool ?? drawingTool ?? (tool as "rectangle" | "polygon" | "ellipse" | "brush");
    const annotation: Annotation = {
      id: `annotation-${Date.now()}`,
      type: finalTool,
      points,
      color: labelColor(activeLabel),
      label: activeLabel,
      completed: true,
    };

    onAnnotationComplete(annotation);
    setCurrentPoints([]);
    setIsDrawing(false);
    setDrawingTool(null);
    redrawAnnotations();
  };

  const redrawAnnotations = useCallback(() => {
    const canvas = canvasRef.current;
    if (!canvas) return;

    const ctx = canvas.getContext("2d");
    if (!ctx) return;

    // Clear canvas
    ctx.clearRect(0, 0, canvas.width, canvas.height);

    if (!showAnnotations) return;

    // Draw completed annotations
    annotations.forEach((annotation) => {
      const isSelected = selectedAnnotationId === annotation.id;
      const isHovered = hoveredAnnotationId === annotation.id;
      const screenPoints = annotation.points.map(toScreenPoint);
      
      if (annotation.type === "rectangle" && screenPoints.length === 2) {
        const [p1, p2] = screenPoints;
        const x = Math.min(p1.x, p2.x);
        const y = Math.min(p1.y, p2.y);
        const width = Math.abs(p2.x - p1.x);
        const height = Math.abs(p2.y - p1.y);

        // Draw only outline, no fill
        if (isSelected) {
          ctx.strokeStyle = "#FCD34D";
          ctx.lineWidth = 4;
        } else if (isHovered) {
          ctx.strokeStyle = "#60A5FA"; // Light blue for hover
          ctx.lineWidth = 3;
        } else {
          ctx.strokeStyle = annotation.color;
          ctx.lineWidth = 2;
        }
        ctx.strokeRect(x, y, width, height);
        
        // Add hover indicator
        if (isHovered) {
          // Add hover indicator
          ctx.fillStyle = "#60A5FA";
          ctx.fillRect(x, y - 22, 65, 18);
          ctx.fillStyle = "#FFF";
          ctx.font = "11px Arial";
          ctx.textAlign = "left";
          ctx.textBaseline = "top";
          ctx.fillText("Click to select", x + 4, y - 18);
        }
      } else if (annotation.type === "ellipse" && screenPoints.length === 2) {
        const [p1, p2] = screenPoints;
        const cx = (p1.x + p2.x) / 2;
        const cy = (p1.y + p2.y) / 2;
        const rx = Math.abs(p2.x - p1.x) / 2;
        const ry = Math.abs(p2.y - p1.y) / 2;

        if (isSelected) {
          ctx.strokeStyle = "#FCD34D";
          ctx.lineWidth = 4;
        } else if (isHovered) {
          ctx.strokeStyle = "#60A5FA";
          ctx.lineWidth = 3;
        } else {
          ctx.strokeStyle = annotation.color;
          ctx.lineWidth = 2;
        }

        ctx.beginPath();
        ctx.ellipse(cx, cy, Math.max(1, rx), Math.max(1, ry), 0, 0, Math.PI * 2);
        ctx.stroke();

        if (isHovered) {
          ctx.fillStyle = "#60A5FA";
          ctx.fillRect(cx - 35, cy - ry - 26, 70, 18);
          ctx.fillStyle = "#FFF";
          ctx.font = "11px Arial";
          ctx.textAlign = "center";
          ctx.textBaseline = "top";
          ctx.fillText("Click to select", cx, cy - ry - 22);
        }
      } else if (annotation.type === "polygon" && screenPoints.length > 1) {
        // Determine colors based on state
        let strokeColor, pointColor;
        let lineWidth = 2;
        
        if (isSelected) {
          strokeColor = "#FCD34D";
          pointColor = "#FCD34D";
          lineWidth = 4;
        } else if (isHovered) {
          strokeColor = "#60A5FA";
          pointColor = "#60A5FA";
          lineWidth = 3;
        } else {
          strokeColor = annotation.color;
          pointColor = annotation.color;
        }
        
        ctx.strokeStyle = strokeColor;
        ctx.lineWidth = lineWidth;

        // Draw main polygon
        ctx.strokeStyle = strokeColor;
        ctx.beginPath();
        ctx.moveTo(screenPoints[0].x, screenPoints[0].y);
        screenPoints.slice(1).forEach((p) => {
          ctx.lineTo(p.x, p.y);
        });
        ctx.closePath();
        ctx.stroke();

        // Draw points as circles
        ctx.fillStyle = pointColor;
        screenPoints.forEach((p, index) => {
          ctx.beginPath();
          const radius = isSelected ? 7 : isHovered ? 6 : 5;
          ctx.arc(p.x, p.y, radius, 0, Math.PI * 2);
          ctx.fill();
          
          // Draw point numbers
          ctx.fillStyle = "#FFFFFF";
          ctx.font = "bold 11px Arial";
          ctx.textAlign = "center";
          ctx.textBaseline = "middle";
          ctx.fillText((index + 1).toString(), p.x, p.y);
          ctx.fillStyle = pointColor;
        });
        
        // Add hover indicator
        if (isHovered && screenPoints.length > 0) {
          const firstPoint = screenPoints[0];
          ctx.fillStyle = "#60A5FA";
          ctx.fillRect(firstPoint.x + 10, firstPoint.y - 22, 65, 18);
          ctx.fillStyle = "#FFF";
          ctx.font = "11px Arial";
          ctx.textAlign = "left";
          ctx.textBaseline = "top";
          ctx.fillText("Click to select", firstPoint.x + 14, firstPoint.y - 18);
        }
      } else if (annotation.type === "brush" && screenPoints.length > 1) {
        if (isSelected) {
          ctx.strokeStyle = "#FCD34D";
          ctx.lineWidth = 4;
        } else if (isHovered) {
          ctx.strokeStyle = "#60A5FA";
          ctx.lineWidth = 3;
        } else {
          ctx.strokeStyle = annotation.color;
          ctx.lineWidth = 2;
        }

        ctx.beginPath();
        ctx.moveTo(screenPoints[0].x, screenPoints[0].y);
        screenPoints.slice(1).forEach((p) => {
          ctx.lineTo(p.x, p.y);
        });
        ctx.stroke();

        if (isHovered) {
          const firstPoint = screenPoints[0];
          ctx.fillStyle = "#60A5FA";
          ctx.fillRect(firstPoint.x + 10, firstPoint.y - 22, 65, 18);
          ctx.fillStyle = "#FFF";
          ctx.font = "11px Arial";
          ctx.textAlign = "left";
          ctx.textBaseline = "top";
          ctx.fillText("Click to select", firstPoint.x + 14, firstPoint.y - 18);
        }
      }
    });

    // Draw current drawing (preview)
    if (isDrawing && currentPoints.length > 0) {
      const color = tool === "rectangle" ? "#EF4444" : "#3B82F6";
      const previewPoints = currentPoints.map(toScreenPoint);

      if ((tool === "rectangle" || tool === "ellipse") && previewPoints.length === 2) {
        const [p1, p2] = previewPoints;
        const x = Math.min(p1.x, p2.x);
        const y = Math.min(p1.y, p2.y);
        const width = Math.abs(p2.x - p1.x);
        const height = Math.abs(p2.y - p1.y);

        // Draw only outline with dashed style, no fill
        ctx.strokeStyle = color;
        ctx.lineWidth = 2;
        ctx.setLineDash([5, 5]);
        if (tool === "rectangle") {
          ctx.strokeRect(x, y, width, height);
        } else {
          const cx = x + width / 2;
          const cy = y + height / 2;
          ctx.beginPath();
          ctx.ellipse(cx, cy, Math.max(1, width / 2), Math.max(1, height / 2), 0, 0, Math.PI * 2);
          ctx.stroke();
        }
        ctx.setLineDash([]);
      } else if (tool === "polygon" && previewPoints.length > 1) {
        ctx.strokeStyle = color;
        ctx.lineWidth = 2;
        ctx.setLineDash([5, 5]);

        ctx.beginPath();
        ctx.moveTo(previewPoints[0].x, previewPoints[0].y);
        previewPoints.slice(1).forEach((p) => {
          ctx.lineTo(p.x, p.y);
        });
        
        // Close the path only if we have 3+ points
        if (previewPoints.length >= 3) {
          ctx.closePath();
        }
        
        ctx.stroke();
        ctx.setLineDash([]);

        // Draw points as solid circles with numbers
        ctx.fillStyle = color;
        previewPoints.forEach((p, index) => {
          ctx.beginPath();
          ctx.arc(p.x, p.y, 5, 0, Math.PI * 2);
          ctx.fill();
          
          // Draw point numbers
          ctx.fillStyle = "#FFFFFF";
          ctx.font = "bold 11px Arial";
          ctx.textAlign = "center";
          ctx.textBaseline = "middle";
          ctx.fillText((index + 1).toString(), p.x, p.y);
          ctx.fillStyle = color;
        });
      } else if (tool === "brush" && previewPoints.length > 1) {
        ctx.strokeStyle = "#EF4444";
        ctx.lineWidth = 2;
        ctx.setLineDash([5, 5]);
        ctx.beginPath();
        ctx.moveTo(previewPoints[0].x, previewPoints[0].y);
        previewPoints.slice(1).forEach((p) => {
          ctx.lineTo(p.x, p.y);
        });
        ctx.stroke();
        ctx.setLineDash([]);
      }
    }
  }, [annotations, currentPoints, hoveredAnnotationId, isDrawing, selectedAnnotationId, showAnnotations, tool, viewer]);

  useEffect(() => {
    redrawAnnotations();
  }, [redrawAnnotations]);

  useEffect(() => {
    if (!isDrawing || !drawingTool) return;
    if (tool === drawingTool) return;

    if (drawingTool === "polygon" && currentPoints.length >= 3) {
      completeAnnotation(drawingTool);
      return;
    }

    if ((drawingTool === "rectangle" || drawingTool === "ellipse") && currentPoints.length === 2) {
      completeAnnotation(drawingTool);
      return;
    }

    if (drawingTool === "brush" && currentPoints.length >= 2) {
      completeAnnotation(drawingTool);
      return;
    }

    setIsDrawing(false);
    setCurrentPoints([]);
    setDrawingTool(null);
  }, [tool, drawingTool, currentPoints, isDrawing]);

  useEffect(() => {
    if (!viewer) return;
    const handleViewportChange = () => redrawAnnotations();
    viewer.addHandler("zoom", handleViewportChange);
    viewer.addHandler("pan", handleViewportChange);
    viewer.addHandler("animation", handleViewportChange);
    viewer.addHandler("open", handleViewportChange);

    return () => {
      viewer.removeHandler("zoom", handleViewportChange);
      viewer.removeHandler("pan", handleViewportChange);
      viewer.removeHandler("animation", handleViewportChange);
      viewer.removeHandler("open", handleViewportChange);
    };
  }, [redrawAnnotations, viewer]);

  return (
    <canvas
      ref={canvasRef}
      className={`absolute inset-0 z-40 ${
        tool === "none" ? "pointer-events-none cursor-grab" : ""
      } ${tool === "select" ? "cursor-default" : ""} ${
        tool === "rectangle" || tool === "polygon" || tool === "ellipse" || tool === "brush"
          ? "cursor-crosshair"
          : ""
      }`}
      style={{ width: "100%", height: "100%" }}
    />
  );
}