SteelXR/src/components/AnnotationEditor.tsx

import { useRef, useState, useEffect, useCallback } from 'react';
import { Button } from '@/components/ui/button';
import { Input } from '@/components/ui/input';
import {
  X, ArrowRight, Circle, Type, Undo2, Download, Check,
} from 'lucide-react';

type Tool = 'arrow' | 'circle' | 'text';

interface Annotation {
  type: Tool;
  x1: number; y1: number;
  x2: number; y2: number;
  text?: string;
}

interface AnnotationEditorProps {
  imageSrc: string;
  onSave: (annotatedDataUrl: string) => void;
  onClose: () => void;
}

export function AnnotationEditor({ imageSrc, onSave, onClose }: AnnotationEditorProps) {
  const canvasRef = useRef<HTMLCanvasElement>(null);
  const imgRef = useRef<HTMLImageElement | null>(null);
  const [tool, setTool] = useState<Tool>('arrow');
  const [annotations, setAnnotations] = useState<Annotation[]>([]);
  const [drawing, setDrawing] = useState(false);
  const [start, setStart] = useState({ x: 0, y: 0 });
  const [current, setCurrent] = useState({ x: 0, y: 0 });
  const [pendingText, setPendingText] = useState<{ x: number; y: number } | null>(null);
  const [textValue, setTextValue] = useState('');
  const [canvasSize, setCanvasSize] = useState({ w: 800, h: 450 });

  // Load image
  useEffect(() => {
    const img = new Image();
    img.onload = () => {
      imgRef.current = img;
      // Fit within container while preserving aspect ratio
      const maxW = Math.min(900, window.innerWidth - 80);
      const maxH = Math.min(600, window.innerHeight - 200);
      const scale = Math.min(maxW / img.width, maxH / img.height, 1);
      setCanvasSize({ w: Math.round(img.width * scale), h: Math.round(img.height * scale) });
    };
    img.src = imageSrc;
  }, [imageSrc]);

  const getPos = useCallback((e: React.MouseEvent<HTMLCanvasElement>) => {
    const rect = canvasRef.current!.getBoundingClientRect();
    return {
      x: ((e.clientX - rect.left) / rect.width) * canvasSize.w,
      y: ((e.clientY - rect.top) / rect.height) * canvasSize.h,
    };
  }, [canvasSize]);

  // Draw everything
  const redraw = useCallback((extra?: Annotation) => {
    const ctx = canvasRef.current?.getContext('2d');
    const img = imgRef.current;
    if (!ctx || !img) return;

    ctx.clearRect(0, 0, canvasSize.w, canvasSize.h);
    ctx.drawImage(img, 0, 0, canvasSize.w, canvasSize.h);

    const allAnnotations = extra ? [...annotations, extra] : annotations;
    allAnnotations.forEach((a) => drawAnnotation(ctx, a));
  }, [annotations, canvasSize]);

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

  const drawAnnotation = (ctx: CanvasRenderingContext2D, a: Annotation) => {
    ctx.strokeStyle = '#ef4444';
    ctx.fillStyle = '#ef4444';
    ctx.lineWidth = 2.5;
    ctx.font = 'bold 16px monospace';

    if (a.type === 'arrow') {
      const dx = a.x2 - a.x1;
      const dy = a.y2 - a.y1;
      const angle = Math.atan2(dy, dx);
      const len = Math.sqrt(dx * dx + dy * dy);
      if (len < 3) return;

      ctx.beginPath();
      ctx.moveTo(a.x1, a.y1);
      ctx.lineTo(a.x2, a.y2);
      ctx.stroke();

      // Arrowhead
      const headLen = Math.min(14, len * 0.3);
      ctx.beginPath();
      ctx.moveTo(a.x2, a.y2);
      ctx.lineTo(a.x2 - headLen * Math.cos(angle - 0.4), a.y2 - headLen * Math.sin(angle - 0.4));
      ctx.lineTo(a.x2 - headLen * Math.cos(angle + 0.4), a.y2 - headLen * Math.sin(angle + 0.4));
      ctx.closePath();
      ctx.fill();
    } else if (a.type === 'circle') {
      const rx = Math.abs(a.x2 - a.x1) / 2;
      const ry = Math.abs(a.y2 - a.y1) / 2;
      const cx = (a.x1 + a.x2) / 2;
      const cy = (a.y1 + a.y2) / 2;
      ctx.beginPath();
      ctx.ellipse(cx, cy, Math.max(rx, 2), Math.max(ry, 2), 0, 0, Math.PI * 2);
      ctx.stroke();
    } else if (a.type === 'text' && a.text) {
      // Background
      const metrics = ctx.measureText(a.text);
      const pad = 4;
      ctx.fillStyle = 'rgba(0,0,0,0.6)';
      ctx.fillRect(a.x1 - pad, a.y1 - 16 - pad, metrics.width + pad * 2, 20 + pad * 2);
      ctx.fillStyle = '#ffffff';
      ctx.fillText(a.text, a.x1, a.y1);
    }
  };

  const handleMouseDown = (e: React.MouseEvent<HTMLCanvasElement>) => {
    if (pendingText) return;
    const pos = getPos(e);
    if (tool === 'text') {
      setPendingText(pos);
      setTextValue('');
      return;
    }
    setDrawing(true);
    setStart(pos);
    setCurrent(pos);
  };

  const handleMouseMove = (e: React.MouseEvent<HTMLCanvasElement>) => {
    if (!drawing) return;
    const pos = getPos(e);
    setCurrent(pos);
    redraw({ type: tool, x1: start.x, y1: start.y, x2: pos.x, y2: pos.y });
  };

  const handleMouseUp = () => {
    if (!drawing) return;
    setDrawing(false);
    setAnnotations((prev) => [
      ...prev,
      { type: tool, x1: start.x, y1: start.y, x2: current.x, y2: current.y },
    ]);
  };

  const confirmText = () => {
    if (!pendingText || !textValue.trim()) {
      setPendingText(null);
      return;
    }
    setAnnotations((prev) => [
      ...prev,
      { type: 'text', x1: pendingText.x, y1: pendingText.y, x2: pendingText.x, y2: pendingText.y, text: textValue.trim() },
    ]);
    setPendingText(null);
    setTextValue('');
  };

  const undo = () => setAnnotations((prev) => prev.slice(0, -1));

  const handleSave = () => {
    // Render full resolution
    const img = imgRef.current;
    if (!img) return;
    const offscreen = document.createElement('canvas');
    offscreen.width = img.width;
    offscreen.height = img.height;
    const ctx = offscreen.getContext('2d')!;
    const scaleX = img.width / canvasSize.w;
    const scaleY = img.height / canvasSize.h;

    ctx.drawImage(img, 0, 0);
    ctx.save();
    ctx.scale(scaleX, scaleY);
    annotations.forEach((a) => drawAnnotation(ctx, a));
    ctx.restore();

    onSave(offscreen.toDataURL('image/png'));
  };

  const tools: { id: Tool; icon: typeof ArrowRight; label: string }[] = [
    { id: 'arrow', icon: ArrowRight, label: 'Seta' },
    { id: 'circle', icon: Circle, label: 'Círculo' },
    { id: 'text', icon: Type, label: 'Texto' },
  ];

  return (
    <div className="fixed inset-0 z-50 flex items-center justify-center bg-black/80 backdrop-blur-sm">
      <div className="flex flex-col gap-3 rounded-xl border bg-card p-4 shadow-2xl max-w-[95vw] max-h-[95vh]">
        {/* Toolbar */}
        <div className="flex items-center justify-between gap-2">
          <div className="flex items-center gap-1">
            {tools.map((t) => (
              <Button
                key={t.id}
                variant={tool === t.id ? 'default' : 'outline'}
                size="sm"
                className="gap-1.5 h-8 text-xs"
                onClick={() => setTool(t.id)}
              >
                <t.icon className="h-3.5 w-3.5" />
                {t.label}
              </Button>
            ))}
            <div className="w-px h-6 bg-border mx-1" />
            <Button variant="ghost" size="sm" className="h-8 text-xs" onClick={undo} disabled={annotations.length === 0}>
              <Undo2 className="h-3.5 w-3.5 mr-1" /> Desfazer
            </Button>
          </div>
          <div className="flex items-center gap-1">
            <Button variant="default" size="sm" className="h-8 text-xs gap-1" onClick={handleSave}>
              <Check className="h-3.5 w-3.5" /> Salvar
            </Button>
            <Button variant="ghost" size="icon" className="h-8 w-8" onClick={onClose}>
              <X className="h-4 w-4" />
            </Button>
          </div>
        </div>

        {/* Canvas area */}
        <div className="relative">
          <canvas
            ref={canvasRef}
            width={canvasSize.w}
            height={canvasSize.h}
            className="rounded-lg border cursor-crosshair"
            style={{ width: canvasSize.w, height: canvasSize.h }}
            onMouseDown={handleMouseDown}
            onMouseMove={handleMouseMove}
            onMouseUp={handleMouseUp}
            onMouseLeave={handleMouseUp}
          />

          {/* Text input overlay */}
          {pendingText && (
            <div
              className="absolute flex items-center gap-1"
              style={{ left: (pendingText.x / canvasSize.w) * 100 + '%', top: (pendingText.y / canvasSize.h) * 100 + '%' }}
            >
              <Input
                autoFocus
                value={textValue}
                onChange={(e) => setTextValue(e.target.value)}
                onKeyDown={(e) => e.key === 'Enter' && confirmText()}
                className="h-7 w-40 text-xs"
                placeholder="Digite o texto…"
              />
              <Button size="sm" className="h-7 w-7 p-0" onClick={confirmText}>
                <Check className="h-3 w-3" />
              </Button>
            </div>
          )}
        </div>

        <p className="font-mono text-[10px] text-muted-foreground text-center">
          {tool === 'arrow' && 'Clique e arraste para desenhar uma seta'}
          {tool === 'circle' && 'Clique e arraste para desenhar um círculo'}
          {tool === 'text' && 'Clique no ponto desejado para adicionar texto'}
        </p>
      </div>
    </div>
  );
}