vfxreview/components/annotations/AnnotationCanvas.tsx

"use client";

import {
  useRef,
  useEffect,
  useState,
  useCallback,
  MouseEvent,
  type RefObject,
} 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 { letterboxRect } from "@/lib/aspect-overlays";
import type { AnnotationShape, AnnotationDrawingData, AnnotationPoint } from "@/types";
import { v4 as uuidv4 } from "uuid";

interface SaveAnnotationData {
  versionId: string;
  frameNumber: number;
  drawingData: AnnotationDrawingData;
  color: string;
}

interface AddCommentData {
  versionId: string;
  frameNumber: number;
  timestamp: number;
  text: string;
}

interface AnnotationCanvasProps {
  versionId: string;
  /** The video under the canvas — the canvas is sized to its displayed picture. */
  videoRef: RefObject<HTMLVideoElement | null>;
  frameNumber: number;
  fps: number;
  isAnnotating: boolean;
  showAnnotations: boolean;
  existingAnnotations?: unknown[];
  onAnnotationSaved?: (frameNumber: number) => void;
  /** Override the default server action used to persist annotations (e.g. for client-token routes) */
  saveAnnotationFn?: (data: SaveAnnotationData) => Promise<void>;
  /** Override the default server action used to create the companion comment */
  addCommentFn?: (data: AddCommentData) => Promise<void>;
}

/**
 * A shape as held in memory. Legacy (v1.0) shapes keep their points relative to
 * the player box they were drawn in; `legacyBox` is that box, used to map them
 * onto the picture once the video's intrinsic size is known.
 */
type LoadedShape = AnnotationShape & { legacyBox?: { w: number; h: number } };

function shapesOf(drawingData: AnnotationDrawingData | undefined): LoadedShape[] {
  const shapes = drawingData?.shapes ?? [];
  if (!drawingData || drawingData.version === "2.0") return shapes;
  const legacyBox = { w: drawingData.canvasWidth, h: drawingData.canvasHeight };
  return shapes.map((s) => ({ ...s, legacyBox }));
}

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

export function AnnotationCanvas({
  versionId,
  videoRef,
  frameNumber,
  fps,
  isAnnotating,
  showAnnotations,
  existingAnnotations = [],
  onAnnotationSaved,
  saveAnnotationFn,
  addCommentFn,
}: 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, LoadedShape[]>>(() => {
    const map: Record<number, LoadedShape[]> = {};
    for (const a of (existingAnnotations as { frameNumber: number; drawingData: AnnotationDrawingData }[])) {
      const shapes = shapesOf(a.drawingData);
      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, LoadedShape[]> = {};
      const serverShapeIds = new Set<string>();
      for (const a of (existingAnnotations as { frameNumber: number; drawingData: AnnotationDrawingData }[])) {
        const shapes = shapesOf(a.drawingData);
        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]);


  // ── Picture rect ─────────────────────────────────────────────────────────
  // The canvas is laid exactly over the displayed picture (not the letterboxed
  // player box), so normalized points are fractions of the image itself and
  // stay locked to it at any player size. Same measurement as AspectOverlay;
  // clientWidth/Height ignore the zoom transform, which the canvas shares.
  const [pictureRect, setPictureRect] = useState<{
    left: number;
    top: number;
    width: number;
    height: number;
  } | null>(null);

  useEffect(() => {
    const video = videoRef.current;
    if (!video) return;
    const measure = () => {
      const rect = letterboxRect(
        video.clientWidth,
        video.clientHeight,
        video.videoWidth,
        video.videoHeight
      );
      setPictureRect((prev) =>
        prev &&
        rect &&
        prev.left === rect.left &&
        prev.top === rect.top &&
        prev.width === rect.width &&
        prev.height === rect.height
          ? prev
          : rect
      );
    };
    measure();
    const observer = new ResizeObserver(measure);
    observer.observe(video);
    video.addEventListener("loadedmetadata", measure);
    return () => {
      observer.disconnect();
      video.removeEventListener("loadedmetadata", measure);
    };
  }, [videoRef]);

  /** A shape's points as fractions of the picture (maps legacy box-relative ones). */
  const pictureSpacePoints = useCallback(
    (shape: LoadedShape): AnnotationPoint[] => {
      const box = shape.legacyBox;
      const video = videoRef.current;
      if (!box || !video) return shape.points;
      const r = letterboxRect(box.w, box.h, video.videoWidth, video.videoHeight);
      if (!r) return shape.points; // no metadata yet — best effort
      return shape.points.map((p) => ({
        x: (p.x * box.w - r.left) / r.width,
        y: (p.y * box.h - r.top) / r.height,
      }));
    },
    [videoRef]
  );

  // ── 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: LoadedShape,
      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 = pictureSpacePoints(shape).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;
    },
    [pictureSpacePoints]
  );

  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: "2.0",
      };

      const doSaveAnnotation = saveAnnotationFn
        ? saveAnnotationFn
        : (data: SaveAnnotationData) => saveAnnotation(data).then(() => undefined);

      const doAddComment = addCommentFn
        ? addCommentFn
        : (data: AddCommentData) => addComment(data).then(() => undefined);

      await doSaveAnnotation({
        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 doAddComment({
          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, saveAnnotationFn, addCommentFn]);

  // ── Touch events (mirrors mouse events for mobile drawing) ──────────────
  const handleTouchStart = useCallback(
    (e: React.TouchEvent<HTMLCanvasElement>) => {
      if (!isAnnotating) return;
      e.preventDefault();
      const canvas = canvasRef.current;
      if (!canvas) return;
      const rect = canvas.getBoundingClientRect();
      const touch = e.touches[0];
      if (!touch) return;
      const point: AnnotationPoint = {
        x: (touch.clientX - rect.left) / rect.width,
        y: (touch.clientY - rect.top) / rect.height,
      };
      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 handleTouchMove = useCallback(
    (e: React.TouchEvent<HTMLCanvasElement>) => {
      if (!drawingStateRef.current.isDrawing || !drawingStateRef.current.currentShape) return;
      e.preventDefault();
      const canvas = canvasRef.current;
      if (!canvas) return;
      const rect = canvas.getBoundingClientRect();
      const touch = e.touches[0];
      if (!touch) return;
      const point: AnnotationPoint = {
        x: (touch.clientX - rect.left) / rect.width,
        y: (touch.clientY - rect.top) / rect.height,
      };
      setDrawingState((prev) => ({
        ...prev,
        currentShape: prev.currentShape
          ? { ...prev.currentShape, points: [...prev.currentShape.points, point] }
          : null,
      }));
      redraw();
    },
    [redraw]
  );

  return (
    <canvas
      ref={canvasRef}
      className={`annotation-canvas-overlay ${isAnnotating ? "is-annotating" : ""}`}
      style={{
        ...(pictureRect && {
          left: pictureRect.left,
          top: pictureRect.top,
          width: pictureRect.width,
          height: pictureRect.height,
        }),
        cursor: isAnnotating ? "crosshair" : "default",
        touchAction: isAnnotating ? "none" : "auto",
      }}
      onMouseDown={handleMouseDown}
      onMouseMove={handleMouseMove}
      onMouseUp={handleMouseUp}
      onMouseLeave={handleMouseUp}
      onTouchStart={handleTouchStart}
      onTouchMove={handleTouchMove}
      onTouchEnd={() => void handleMouseUp()}
      onTouchCancel={() => void handleMouseUp()}
    />
  );
}