vfxreview/components/annotations/AnnotationCanvas.tsx

"use client";

import {
  useRef,
  useEffect,
  useState,
  useCallback,
  MouseEvent,
} from "react";
import { useReviewStore } from "@/hooks/use-review-player";
import { saveAnnotation } from "@/actions/annotations";
import { addComment } from "@/actions/comments";
import { useToast } from "@/components/ui/use-toast";
import type { AnnotationShape, AnnotationDrawingData, AnnotationPoint } from "@/types";
import { v4 as uuidv4 } from "uuid";

interface AnnotationCanvasProps {
  versionId: string;
  frameNumber: number;
  fps: number;
  isAnnotating: boolean;
  showAnnotations: boolean;
  existingAnnotations?: unknown[];
  onAnnotationSaved?: (frameNumber: number) => void;
}

type DrawingState = {
  isDrawing: boolean;
  currentShape: AnnotationShape | null;
  shapes: AnnotationShape[];
};

export function AnnotationCanvas({
  versionId,
  frameNumber,
  fps,
  isAnnotating,
  showAnnotations,
  existingAnnotations = [],
  onAnnotationSaved,
}: AnnotationCanvasProps) {
  const canvasRef = useRef<HTMLCanvasElement>(null);
  const { selectedTool, selectedColor, strokeWidth } = useReviewStore();
  const { toast } = useToast();

  const [drawingState, setDrawingState] = useState<DrawingState>({
    isDrawing: false,
    currentShape: null,
    shapes: [],
  });
  const drawingStateRef = useRef(drawingState);
  drawingStateRef.current = drawingState;

  // Per-frame saved shapes — persists across frame navigation so annotations
  // remain visible without depending on a parent API refetch.
  const [savedShapesByFrame, setSavedShapesByFrame] = useState<Record<number, AnnotationShape[]>>(() => {
    const map: Record<number, AnnotationShape[]> = {};
    for (const a of (existingAnnotations as { frameNumber: number; drawingData: AnnotationDrawingData }[])) {
      const shapes = a.drawingData?.shapes ?? [];
      map[a.frameNumber] = [...(map[a.frameNumber] ?? []), ...shapes];
    }
    return map;
  });
  // Ref so redraw closure can read latest without being recreated
  const savedShapesByFrameRef = useRef(savedShapesByFrame);
  savedShapesByFrameRef.current = savedShapesByFrame;

  // Track frames that already have an annotation comment this session,
  // initialized from existing annotations so we don't double-comment on refresh.
  const [annotationCommentedFrames, setAnnotationCommentedFrames] = useState<Set<number>>(() => {
    const frames = new Set<number>();
    for (const a of (existingAnnotations as { frameNumber: number }[])) {
      frames.add(a.frameNumber);
    }
    return frames;
  });

  // Sync savedShapesByFrame when the existingAnnotations prop changes (e.g. after parent API refresh).
  // MERGE server shapes with locally-drawn shapes (deduplicate by shape id) so locally-drawn
  // annotations are never wiped by a new array reference from the parent.
  const prevExistingLengthRef = useRef(existingAnnotations.length);
  const prevExistingRef = useRef(existingAnnotations);
  useEffect(() => {
    if (
      prevExistingRef.current === existingAnnotations ||
      prevExistingLengthRef.current === existingAnnotations.length
    ) {
      prevExistingRef.current = existingAnnotations;
      prevExistingLengthRef.current = existingAnnotations.length;
      return;
    }
    prevExistingRef.current = existingAnnotations;
    prevExistingLengthRef.current = existingAnnotations.length;

    setSavedShapesByFrame((prev) => {
      // Build map from server data
      const serverMap: Record<number, AnnotationShape[]> = {};
      const serverShapeIds = new Set<string>();
      for (const a of (existingAnnotations as { frameNumber: number; drawingData: AnnotationDrawingData }[])) {
        const shapes = a.drawingData?.shapes ?? [];
        serverMap[a.frameNumber] = [...(serverMap[a.frameNumber] ?? []), ...shapes];
        for (const s of shapes) serverShapeIds.add(s.id);
      }
      // Merge: keep local shapes not yet in server data (pending save race)
      const merged = { ...serverMap };
      for (const [frame, shapes] of Object.entries(prev)) {
        for (const s of shapes) {
          if (!serverShapeIds.has(s.id)) {
            const f = Number(frame);
            merged[f] = [...(merged[f] ?? []), s];
          }
        }
      }
      return merged;
    });

    // Also update frames that now have server-side annotation comments
    setAnnotationCommentedFrames((prev) => {
      const next = new Set(prev);
      for (const a of (existingAnnotations as { frameNumber: number }[])) {
        next.add(a.frameNumber);
      }
      return next;
    });
  }, [existingAnnotations]);


  // ── Coordinate normalization ─────────────────────────────────────────────
  const toNormalized = (e: MouseEvent): AnnotationPoint => {
    const canvas = canvasRef.current!;
    const rect = canvas.getBoundingClientRect();
    return {
      x: (e.clientX - rect.left) / rect.width,
      y: (e.clientY - rect.top) / rect.height,
    };
  };

  const toPixel = (p: AnnotationPoint, w: number, h: number) => ({
    x: p.x * w,
    y: p.y * h,
  });

  // ── Rendering ─────────────────────────────────────────────────────────────
  const renderShape = useCallback(
    (
      ctx: CanvasRenderingContext2D,
      shape: AnnotationShape,
      w: number,
      h: number
    ) => {
      if (shape.points.length === 0) return;

      ctx.strokeStyle = shape.color;
      ctx.lineWidth = shape.strokeWidth;
      ctx.lineCap = "round";
      ctx.lineJoin = "round";
      ctx.globalAlpha = 0.85;

      const pts = shape.points.map((p) => toPixel(p, w, h));

      switch (shape.tool) {
        case "freehand": {
          if (pts.length < 2) return;
          ctx.beginPath();
          ctx.moveTo(pts[0].x, pts[0].y);
          for (let i = 1; i < pts.length; i++) {
            ctx.lineTo(pts[i].x, pts[i].y);
          }
          ctx.stroke();
          break;
        }
        case "rectangle": {
          if (pts.length < 2) return;
          const x = Math.min(pts[0].x, pts[pts.length - 1].x);
          const y = Math.min(pts[0].y, pts[pts.length - 1].y);
          const width = Math.abs(pts[pts.length - 1].x - pts[0].x);
          const height = Math.abs(pts[pts.length - 1].y - pts[0].y);
          ctx.beginPath();
          ctx.rect(x, y, width, height);
          ctx.stroke();
          break;
        }
        case "circle": {
          if (pts.length < 2) return;
          const cx = pts[0].x;
          const cy = pts[0].y;
          const dx = pts[pts.length - 1].x - cx;
          const dy = pts[pts.length - 1].y - cy;
          const radius = Math.sqrt(dx * dx + dy * dy);
          ctx.beginPath();
          ctx.arc(cx, cy, radius, 0, Math.PI * 2);
          ctx.stroke();
          break;
        }
        case "arrow": {
          if (pts.length < 2) return;
          const start = pts[0];
          const end = pts[pts.length - 1];
          const angle = Math.atan2(end.y - start.y, end.x - start.x);
          const arrowLen = Math.max(10, shape.strokeWidth * 5);

          ctx.beginPath();
          ctx.moveTo(start.x, start.y);
          ctx.lineTo(end.x, end.y);
          ctx.stroke();

          // Arrowhead
          ctx.fillStyle = shape.color;
          ctx.globalAlpha = 0.85;
          ctx.beginPath();
          ctx.moveTo(end.x, end.y);
          ctx.lineTo(
            end.x - arrowLen * Math.cos(angle - Math.PI / 7),
            end.y - arrowLen * Math.sin(angle - Math.PI / 7)
          );
          ctx.lineTo(
            end.x - arrowLen * Math.cos(angle + Math.PI / 7),
            end.y - arrowLen * Math.sin(angle + Math.PI / 7)
          );
          ctx.closePath();
          ctx.fill();
          break;
        }
      }
      ctx.globalAlpha = 1;
    },
    []
  );

  const redraw = useCallback(() => {
    const canvas = canvasRef.current;
    if (!canvas) return;
    const ctx = canvas.getContext("2d");
    if (!ctx) return;
    const w = canvas.width;
    const h = canvas.height;

    ctx.clearRect(0, 0, w, h);

    if (!showAnnotations) return;

    // Draw saved shapes for this frame (seeded from DB + newly saved this session)
    (savedShapesByFrameRef.current[frameNumber] ?? []).forEach((shape) => {
      renderShape(ctx, shape, w, h);
    });

    // Draw current (in-progress) shape
    const current = drawingStateRef.current.currentShape;
    if (current && drawingStateRef.current.isDrawing) {
      renderShape(ctx, current, w, h);
    }
  }, [frameNumber, showAnnotations, renderShape]);

  // Resize observer — stable, only recreated when redraw changes (which is rare now)
  useEffect(() => {
    const canvas = canvasRef.current;
    if (!canvas) return;
    const observer = new ResizeObserver(() => {
      canvas.width = canvas.offsetWidth;
      canvas.height = canvas.offsetHeight;
      redraw();
    });
    observer.observe(canvas);
    return () => observer.disconnect();
  }, [redraw]);

  // Redraw whenever frame, saved shapes map, in-progress shape, or visibility changes
  useEffect(() => {
    redraw();
  }, [redraw, savedShapesByFrame, drawingState.currentShape, showAnnotations]);

  // ── Mouse events ─────────────────────────────────────────────────────────
  const handleMouseDown = useCallback(
    (e: MouseEvent<HTMLCanvasElement>) => {
      if (!isAnnotating) return;
      const point = toNormalized(e);
      const newShape: AnnotationShape = {
        id: uuidv4(),
        tool: selectedTool,
        points: [point],
        color: selectedColor,
        strokeWidth,
        frameNumber,
      };
      setDrawingState((prev) => ({
        ...prev,
        isDrawing: true,
        currentShape: newShape,
      }));
    },
    [isAnnotating, selectedTool, selectedColor, strokeWidth, frameNumber]
  );

  const handleMouseMove = useCallback(
    (e: MouseEvent<HTMLCanvasElement>) => {
      if (!drawingStateRef.current.isDrawing || !drawingStateRef.current.currentShape) return;
      const point = toNormalized(e);
      setDrawingState((prev) => ({
        ...prev,
        currentShape: prev.currentShape
          ? { ...prev.currentShape, points: [...prev.currentShape.points, point] }
          : null,
      }));
      redraw();
    },
    [redraw]
  );

  const handleMouseUp = useCallback(async () => {
    const state = drawingStateRef.current;
    if (!state.isDrawing || !state.currentShape) return;

    const finishedShape = state.currentShape;
    // Clear the in-progress drawing state
    setDrawingState((prev) => ({
      ...prev,
      isDrawing: false,
      currentShape: null,
      shapes: [],
    }));

    // Immediately persist shape into the per-frame cache so it survives frame navigation
    setSavedShapesByFrame((prev) => ({
      ...prev,
      [frameNumber]: [...(prev[frameNumber] ?? []), finishedShape],
    }));

    // Save to DB
    try {
      const canvas = canvasRef.current;
      const drawingData: AnnotationDrawingData = {
        shapes: [finishedShape],
        canvasWidth: canvas?.offsetWidth ?? 1920,
        canvasHeight: canvas?.offsetHeight ?? 1080,
        version: "1.0",
      };
      await saveAnnotation({
        versionId,
        frameNumber,
        drawingData,
        color: selectedColor,
      });

      // Create a companion comment ONCE per frame so the annotation appears in the panel.
      // Additional strokes on the same frame update the existing comment count silently.
      if (!annotationCommentedFrames.has(frameNumber)) {
        await addComment({
          versionId,
          frameNumber,
          timestamp: frameNumber / fps,
          text: `✏️ Annotation at frame ${frameNumber}`,
        });
        setAnnotationCommentedFrames((prev) => new Set([...prev, frameNumber]));
      }

      onAnnotationSaved?.(frameNumber);
    } catch {
      toast({
        title: "Failed to save annotation",
        variant: "destructive",
      });
    }
  }, [versionId, frameNumber, fps, selectedColor, annotationCommentedFrames, onAnnotationSaved, toast]);

  return (
    <canvas
      ref={canvasRef}
      className={`annotation-canvas-overlay ${isAnnotating ? "is-annotating" : ""}`}
      style={{
        cursor: isAnnotating ? "crosshair" : "default",
      }}
      onMouseDown={handleMouseDown}
      onMouseMove={handleMouseMove}
      onMouseUp={handleMouseUp}
      onMouseLeave={handleMouseUp}
    />
  );
}