omniai-web/src/features/canvas/canvasToolPanels.tsx

import { useCallback, useRef, useState } from "react";
import { aiGenerationClient } from "../../api/aiGenerationClient";
import { waitForTask } from "../../api/taskSubscription";
import { toast } from "../../components/toast/toastStore";
import type { CanvasImageNode } from "./canvasTypes";

interface CanvasToolPanelProps {
  imageUrl: string;
  imageNode: CanvasImageNode;
  onComplete: (resultUrl: string) => void;
}

export function CanvasMultiGridPanel({ imageUrl, onComplete }: CanvasToolPanelProps) {
  const [gridMode, setGridMode] = useState<"grid-4" | "grid-9">("grid-4");
  const [prompt, setPrompt] = useState("");
  const [loading, setLoading] = useState(false);
  const cancelRef = useRef(false);

  const handleGenerate = useCallback(async () => {
    if (!imageUrl) return;
    setLoading(true);
    cancelRef.current = false;
    try {
      const { taskId } = await aiGenerationClient.createImageTask({
        model: "gpt-image-2",
        prompt: prompt || "基于参考图生成多宫格变体",
        referenceUrls: [imageUrl],
        gridMode,
      });
      const resultUrl = await waitForTask(taskId, { kind: "image", abortRef: cancelRef });
      if (resultUrl) {
        onComplete(resultUrl);
        toast.success("多宫格生成完成");
      }
    } catch (err: unknown) {
      if (!cancelRef.current) toast.error(err instanceof Error ? err.message : "多宫格生成失败");
    } finally {
      setLoading(false);
    }
  }, [imageUrl, prompt, gridMode, onComplete]);

  return (
    <div className="studio-canvas-tool-panel">
      <div className="studio-canvas-tool-panel__preview"><img src={imageUrl} alt="" /></div>
      <div className="studio-canvas-tool-panel__controls">
        <label className="studio-canvas-tool-panel__label">{"宫格模式"}</label>
        <div className="studio-canvas-tool-panel__options">
          {([["grid-4", "2×2"], ["grid-9", "3×3"]] as const).map(([value, label]) => (
            <button key={value} type="button" className={gridMode === value ? "is-active" : ""} onClick={() => setGridMode(value)}>{label}</button>
          ))}
        </div>
        <label className="studio-canvas-tool-panel__label">{"提示词（可选）"}</label>
        <textarea className="studio-canvas-tool-panel__textarea" value={prompt} onChange={(e) => setPrompt(e.target.value)} placeholder="描述多宫格内容变化" />
        <button type="button" className="studio-canvas-tool-panel__submit" disabled={loading} onClick={handleGenerate}>
          {loading ? "生成中..." : "生成多宫格"}
        </button>
      </div>
    </div>
  );
}

export function CanvasUpscalePanel({ imageUrl, onComplete }: CanvasToolPanelProps) {
  const [scale, setScale] = useState<"2x" | "4x">("2x");
  const [loading, setLoading] = useState(false);
  const cancelRef = useRef(false);

  const handleUpscale = useCallback(async () => {
    if (!imageUrl) return;
    setLoading(true);
    cancelRef.current = false;
    try {
      const { taskId } = await aiGenerationClient.createImageSuperResolveTask({
        imageUrl,
        scale,
      });
      const resultUrl = await waitForTask(taskId, { kind: "image", abortRef: cancelRef });
      if (resultUrl) {
        onComplete(resultUrl);
        toast.success("超分完成");
      }
    } catch (err: unknown) {
      if (!cancelRef.current) toast.error(err instanceof Error ? err.message : "超分失败");
    } finally {
      setLoading(false);
    }
  }, [imageUrl, scale, onComplete]);

  return (
    <div className="studio-canvas-tool-panel">
      <div className="studio-canvas-tool-panel__preview"><img src={imageUrl} alt="" /></div>
      <div className="studio-canvas-tool-panel__controls">
        <label className="studio-canvas-tool-panel__label">{"放大倍数"}</label>
        <div className="studio-canvas-tool-panel__options">
          {(["2x", "4x"] as const).map((s) => (
            <button key={s} type="button" className={scale === s ? "is-active" : ""} onClick={() => setScale(s)}>{s}</button>
          ))}
        </div>
        <button type="button" className="studio-canvas-tool-panel__submit" disabled={loading} onClick={handleUpscale}>
          {loading ? "处理中..." : "开始超分"}
        </button>
      </div>
    </div>
  );
}

export function CanvasInpaintPanel({ imageUrl, onComplete }: CanvasToolPanelProps) {
  const [prompt, setPrompt] = useState("");
  const [brushSize, setBrushSize] = useState(30);
  const [loading, setLoading] = useState(false);
  const canvasRef = useRef<HTMLCanvasElement>(null);
  const isDrawingRef = useRef(false);
  const cancelRef = useRef(false);

  const initCanvas = useCallback(() => {
    const canvas = canvasRef.current;
    if (!canvas) return;
    const img = new Image();
    img.crossOrigin = "anonymous";
    img.onload = () => {
      canvas.width = img.naturalWidth;
      canvas.height = img.naturalHeight;
      const ctx = canvas.getContext("2d");
      if (!ctx) return;
      ctx.clearRect(0, 0, canvas.width, canvas.height);
    };
    img.src = imageUrl;
  }, [imageUrl]);

  const getPos = (e: React.MouseEvent<HTMLCanvasElement>) => {
    const canvas = canvasRef.current!;
    const rect = canvas.getBoundingClientRect();
    const scaleX = canvas.width / rect.width;
    const scaleY = canvas.height / rect.height;
    return { x: (e.clientX - rect.left) * scaleX, y: (e.clientY - rect.top) * scaleY };
  };

  const draw = (e: React.MouseEvent<HTMLCanvasElement>) => {
    if (!isDrawingRef.current) return;
    const ctx = canvasRef.current?.getContext("2d");
    if (!ctx) return;
    const { x, y } = getPos(e);
    ctx.globalCompositeOperation = "source-over";
    ctx.fillStyle = "rgba(255, 0, 0, 0.4)";
    ctx.beginPath();
    ctx.arc(x, y, brushSize, 0, Math.PI * 2);
    ctx.fill();
  };

  const getMaskDataUrl = (): string => {
    const canvas = canvasRef.current!;
    const maskCanvas = document.createElement("canvas");
    maskCanvas.width = canvas.width;
    maskCanvas.height = canvas.height;
    const srcCtx = canvas.getContext("2d")!;
    const maskCtx = maskCanvas.getContext("2d")!;
    const imgData = srcCtx.getImageData(0, 0, canvas.width, canvas.height);
    const maskData = maskCtx.createImageData(canvas.width, canvas.height);
    for (let i = 0; i < imgData.data.length; i += 4) {
      const hasColor = imgData.data[i + 3] > 10;
      maskData.data[i] = hasColor ? 255 : 0;
      maskData.data[i + 1] = hasColor ? 255 : 0;
      maskData.data[i + 2] = hasColor ? 255 : 0;
      maskData.data[i + 3] = 255;
    }
    maskCtx.putImageData(maskData, 0, 0);
    return maskCanvas.toDataURL("image/png");
  };

  const handleInpaint = useCallback(async () => {
    if (!imageUrl || !prompt) {
      toast.error("请输入重绘提示词");
      return;
    }
    setLoading(true);
    cancelRef.current = false;
    try {
      const maskDataUrl = getMaskDataUrl();
      const { taskId } = await aiGenerationClient.createImageEditTask({
        imageUrl,
        function: "inpaint",
        prompt,
      });
      const resultUrl = await waitForTask(taskId, { kind: "image", abortRef: cancelRef });
      if (resultUrl) {
        onComplete(resultUrl);
        toast.success("局部重绘完成");
      }
    } catch (err: unknown) {
      if (!cancelRef.current) toast.error(err instanceof Error ? err.message : "局部重绘失败");
    } finally {
      setLoading(false);
    }
  }, [imageUrl, prompt, onComplete]);
  return (
    <div className="studio-canvas-tool-panel studio-canvas-tool-panel--inpaint">
      <div className="studio-canvas-tool-panel__canvas-wrap">
        <img src={imageUrl} alt="" className="studio-canvas-tool-panel__canvas-bg" />
        <canvas
          ref={canvasRef}
          className="studio-canvas-tool-panel__canvas"
          onMouseDown={(e) => { isDrawingRef.current = true; draw(e); }}
          onMouseMove={draw}
          onMouseUp={() => { isDrawingRef.current = false; }}
          onMouseLeave={() => { isDrawingRef.current = false; }}
        />
      </div>
      <div className="studio-canvas-tool-panel__controls">
        <label className="studio-canvas-tool-panel__label">{"画笔大小"}</label>
        <input type="range" min={5} max={80} value={brushSize} onChange={(e) => setBrushSize(Number(e.target.value))} />
        <label className="studio-canvas-tool-panel__label">{"重绘提示词"}</label>
        <textarea className="studio-canvas-tool-panel__textarea" value={prompt} onChange={(e) => setPrompt(e.target.value)} placeholder="描述需要重绘区域的内容" />
        <div className="studio-canvas-tool-panel__actions">
          <button type="button" className="studio-canvas-tool-panel__reset" onClick={initCanvas}>{"清除蒙版"}</button>
          <button type="button" className="studio-canvas-tool-panel__submit" disabled={loading} onClick={handleInpaint}>
            {loading ? "处理中..." : "开始重绘"}
          </button>
        </div>
      </div>
    </div>
  );
}