omniai-web/src/features/resolution-upscale/ResolutionUpscalePage.tsx

import {
  CameraOutlined,
  ColumnWidthOutlined,
  CustomerServiceOutlined,
  DeleteOutlined,
  DownloadOutlined,
  EditOutlined,
  FileImageOutlined,
  FolderAddOutlined,
  FontSizeOutlined,
  LinkOutlined,
  PictureOutlined,
  ScissorOutlined,
  SwapOutlined,
  ThunderboltOutlined,
  VideoCameraOutlined,
} from "@ant-design/icons";
import { useCallback, useEffect, useRef, useState, type DragEvent } from "react";
import { aiGenerationClient } from "../../api/aiGenerationClient";
import { waitForTask } from "../../api/taskSubscription";
import { saveToolTaskState, loadToolTaskState, clearToolTaskState } from "../workbench/toolKeepalive";
import { translateTaskError } from "../../utils/translateTaskError";
import { getServerBaseUrl, isServerRequestError } from "../../api/serverConnection";
import { summarizeUrl, formatFileSize, fileToDataUrl, wait } from "../../utils/toolPageUtils";
import TaskStatusBar from "../../components/TaskStatusBar";
import BeforeAfterCompare from "../../components/BeforeAfterCompare";
import { addToolResultToAssetLibrary, saveToolResultToLocal } from "../workbench/toolResultActions";
import type { WebImageWorkbenchTool, WebViewKey } from "../../types";

type UpscaleMode = "image" | "video";
type ImageScale = "2x" | "4x";
type VideoMinLen = 540 | 720;

const videoStyles = [
  { value: 0, label: "日式漫画", color: "#f9a8d4" },
  { value: 1, label: "美式漫画", color: "#fbbf24" },
  { value: 2, label: "清新漫画", color: "#86efac" },
  { value: 3, label: "3D卡通", color: "#93c5fd" },
  { value: 4, label: "国风卡通", color: "#fca5a5" },
  { value: 5, label: "纸艺风格", color: "#d6d3d1" },
  { value: 6, label: "简易插画", color: "#a5b4fc" },
  { value: 7, label: "国风水墨", color: "#94a3b8" },
];

function formatCreateTaskError(error: unknown, taskMode: UpscaleMode): string {
  if (isServerRequestError(error) && error.status === 404) {
    const baseUrl = getServerBaseUrl() || window.location.origin;
    const route = taskMode === "image" ? "/api/ai/image/super-resolve" : "/api/ai/video/super-resolve";
    return `当前连接的服务端 ${baseUrl} 还没有部署 ${route} 接口，请更新并重启服务端后再试`;
  }
  return error instanceof Error ? error.message : "超分任务创建失败，请稍后重试";
}

interface ResolutionUpscalePageProps {
  isAuthenticated?: boolean;
  onOpenMore?: () => void;
  onOpenImageTool?: (tool: WebImageWorkbenchTool) => void;
  onSelectView?: (view: WebViewKey) => void;
}

function ResolutionUpscalePage({
  isAuthenticated = false,
  onOpenMore,
  onOpenImageTool,
  onSelectView,
}: ResolutionUpscalePageProps) {
  const fileInputRef = useRef<HTMLInputElement>(null);
  const pollRunRef = useRef(0);
  const cancelRef = useRef(false);
  const [mode, setMode] = useState<UpscaleMode>("image");
  const [imageScale, setImageScale] = useState<ImageScale>("2x");
  const [keepOriginalStyle, setKeepOriginalStyle] = useState(true);
  const [videoStyle, setVideoStyle] = useState(0);
  const [videoMinLen, setVideoMinLen] = useState<VideoMinLen>(540);
  const [videoFps, setVideoFps] = useState(15);
  const [useSR, setUseSR] = useState(true);
  const [sourceName, setSourceName] = useState("");
  const [sourceFile, setSourceFile] = useState<File | null>(null);
  const [sourceUrl, setSourceUrl] = useState("");
  const [sourcePreview, setSourcePreview] = useState("");
  const [resultPreview, setResultPreview] = useState("");
  const [sourceDimensions, setSourceDimensions] = useState<{ width: number; height: number } | null>(null);
  const [videoViewMode, setVideoViewMode] = useState<"source" | "result">("source");
  const [status, setStatus] = useState("上传素材并选择参数后点击开始超分");
  const [activeTaskId, setActiveTaskId] = useState("");
  const [taskProgress, setTaskProgress] = useState(0);
  const [isProcessing, setIsProcessing] = useState(false);
  const [isDownloading, setIsDownloading] = useState(false);
  const [isSavingAsset, setIsSavingAsset] = useState(false);
  const [isDragging, setIsDragging] = useState(false);
  const activeTaskIdRef = useRef(activeTaskId);
  activeTaskIdRef.current = activeTaskId;
  const keepaliveRestoredRef = useRef(false);

  useEffect(() => {
    return () => {
      if (sourcePreview.startsWith("blob:")) URL.revokeObjectURL(sourcePreview);
    };
  }, [sourcePreview]);

  useEffect(() => {
    return () => {
      if (resultPreview.startsWith("blob:")) URL.revokeObjectURL(resultPreview);
    };
  }, [resultPreview]);

  // Keep-alive: restore saved task on mount
  useEffect(() => {
    if (keepaliveRestoredRef.current) return;
    keepaliveRestoredRef.current = true;
    const saved = loadToolTaskState("upscale");
    if (!saved || saved.resultUrl) return;
    setSourceName(saved.sourceName || "");
    setSourceUrl(saved.sourceUrl || "");
    setIsProcessing(true);
    cancelRef.current = false;
    pollRunRef.current += 1;
    void waitForTaskResult(saved.taskId, mode).catch(() => {});
  }, []);

  useEffect(() => {
    return () => {
      // Stop polling but keep server task alive — keep-alive will resume on remount
      pollRunRef.current += 1;
      cancelRef.current = true;
    };
  }, []);

  const clearSource = () => {
    if (sourcePreview.startsWith("blob:")) URL.revokeObjectURL(sourcePreview);
    setSourceName("");
    setSourceFile(null);
    setSourceUrl("");
    setSourcePreview("");
    setResultPreview("");
    setSourceDimensions(null);
    setVideoViewMode("source");
    setActiveTaskId("");
    setTaskProgress(0);
    setStatus("已清空素材");
  };

  const handleModeChange = (nextMode: UpscaleMode) => {
    if (nextMode === mode) return;
    clearSource();
    setMode(nextMode);
    setStatus(nextMode === "image" ? "已切换为图片超分，请上传图片素材" : "已切换为视频超分，请上传视频素材");
  };

  const handleFileChange = (event: React.ChangeEvent<HTMLInputElement>) => {
    const file = event.target.files?.[0];
    if (!file) return;

    if (sourcePreview.startsWith("blob:")) URL.revokeObjectURL(sourcePreview);
    setSourceName(file.name);
    setSourceFile(file);
    setSourceUrl("");
    setSourcePreview(URL.createObjectURL(file));
    setResultPreview("");
    setSourceDimensions(null);
    setVideoViewMode("source");
    setActiveTaskId("");
    setTaskProgress(0);
    setStatus(`已导入 ${file.name}`);
    event.currentTarget.value = "";
  };

  const processDroppedFile = (file: File) => {
    if (sourcePreview.startsWith("blob:")) URL.revokeObjectURL(sourcePreview);
    setSourceName(file.name);
    setSourceFile(file);
    setSourceUrl("");
    setSourcePreview(URL.createObjectURL(file));
    setResultPreview("");
    setSourceDimensions(null);
    setVideoViewMode("source");
    setActiveTaskId("");
    setTaskProgress(0);
    setStatus(`已导入 ${file.name}`);
  };

  const handleDragOver = (e: DragEvent<HTMLDivElement>) => { e.preventDefault(); e.stopPropagation(); if (e.dataTransfer.types.includes("Files")) setIsDragging(true); };
  const handleDragLeave = (e: DragEvent<HTMLDivElement>) => { e.preventDefault(); e.stopPropagation(); if (e.currentTarget === e.target || !e.currentTarget.contains(e.relatedTarget as Node)) setIsDragging(false); };
  const handleDrop = (e: DragEvent<HTMLDivElement>) => { e.preventDefault(); e.stopPropagation(); setIsDragging(false); const file = e.dataTransfer.files[0]; if (file) processDroppedFile(file); };

  const handleImportUrl = () => {
    const normalizedUrl = sourceUrl.trim();
    if (!/^https?:\/\//i.test(normalizedUrl)) {
      setStatus("请输入可访问的 HTTP/HTTPS 素材 URL");
      return;
    }
    if (sourcePreview.startsWith("blob:")) URL.revokeObjectURL(sourcePreview);
    setSourceFile(null);
    setSourceName(summarizeUrl(normalizedUrl));
    setSourcePreview(normalizedUrl);
    setResultPreview("");
    setVideoViewMode("source");
    setActiveTaskId("");
    setTaskProgress(0);
    setStatus(`已导入 URL：${summarizeUrl(normalizedUrl)}`);
  };

  const getSourceGenerationUrl = async () => {
    if (!sourceFile) return sourcePreview;
    setStatus(`正在上传 ${sourceFile.name}，${formatFileSize(sourceFile.size)}`);
    const uploaded = await aiGenerationClient.uploadAsset({
      dataUrl: await fileToDataUrl(sourceFile),
      name: sourceFile.name,
      mimeType: sourceFile.type || (mode === "image" ? "image/png" : "video/mp4"),
    });
    return uploaded.signedUrl || uploaded.url;
  };

  const waitForTaskResult = async (taskId: string, taskMode: UpscaleMode) => {
    const runId = ++pollRunRef.current;
    setActiveTaskId(taskId);
    setTaskProgress(5);
    saveToolTaskState("upscale", { taskId, sourceName, sourceUrl, status: `任务 ${taskId}`, progress: 5 });

    await waitForTask(taskId, {
      abortRef: cancelRef,
      onProgress: (e) => {
        if (pollRunRef.current !== runId) return;
        const progress = Math.max(0, Math.min(100, Math.trunc(e.progress || 0)));
        setTaskProgress(progress);
        setStatus(`任务 ${taskId} ${e.status}，进度 ${progress}%`);
        if (e.status === "completed" && e.resultUrl) {
          setResultPreview(e.resultUrl);
          setVideoViewMode("result");
          setStatus(`${taskMode === "image" ? "图片" : "视频"}超分完成：${summarizeUrl(e.resultUrl)}`);
          setTaskProgress(100);
          clearToolTaskState("upscale");
          saveToolTaskState("upscale", { taskId, resultUrl: e.resultUrl, resultPreview: e.resultUrl, sourceName, sourceUrl, status: "完成", progress: 100 });
        }
      },
    });

    if (cancelRef.current) throw new Error("已取消");
    if (pollRunRef.current !== runId) throw new Error("当前任务已被新的超分任务替换");
  };

  const handleCancel = useCallback(() => {
    cancelRef.current = true;
    if (activeTaskId) {
      aiGenerationClient.cancelTask(activeTaskId).catch(() => {});
    }
    setIsProcessing(false);
    setStatus("已取消");
    clearToolTaskState("upscale");
  }, [activeTaskId]);

  const handleDownload = async () => {
    if (!resultPreview || isDownloading) return;
    setIsDownloading(true);
    try {
      const isVideo = mode === "video";
      const filenameBase = isVideo ? `upscaled-video-${Date.now()}` : `upscaled-image-${Date.now()}`;
      await saveToolResultToLocal({
        url: resultPreview,
        name: filenameBase,
        type: isVideo ? "video" : "image",
        isVideo,
        taskId: activeTaskId || undefined,
        tags: ["工具盒", "超分辨率", isVideo ? "生成视频" : "生成图片"],
      });
      setStatus("下载完成");
    } catch (error) {
      setStatus(error instanceof Error ? error.message : "下载失败");
    } finally {
      setIsDownloading(false);
    }
  };

  const handleSaveToAssets = async () => {
    if (!resultPreview || isSavingAsset) return;
    setIsSavingAsset(true);
    try {
      const isVideo = mode === "video";
      const status = await addToolResultToAssetLibrary({
        url: resultPreview,
        name: `upscaled-${isVideo ? "video" : "image"}-${Date.now()}.${isVideo ? "mp4" : "png"}`,
        description: `工具盒${isVideo ? "视频" : "图片"}超分生成结果`,
        type: isVideo ? "video" : "image",
        isVideo,
        taskId: activeTaskId || undefined,
        tags: ["工具盒", "超分辨率", isVideo ? "生成视频" : "生成图片"],
      });
      setStatus(status === "server" ? "已加入资产库" : "已加入本地资产库");
    } catch (error) {
      setStatus(error instanceof Error ? error.message : "加入资产库失败");
    } finally {
      setIsSavingAsset(false);
    }
  };

  const handleStart = async () => {
    if (!isAuthenticated) {
      setStatus("请先登录后再创建超分任务");
      return;
    }
    if (!sourcePreview) {
      setStatus(mode === "image" ? "请先上传需要提升分辨率的图片" : "请先上传需要提升分辨率的视频");
      return;
    }
    if (isProcessing) return;

    setIsProcessing(true);
    setResultPreview("");
    setTaskProgress(0);
    cancelRef.current = false;
    try {
      const generationUrl = await getSourceGenerationUrl();
      setStatus(`素材已就绪：${summarizeUrl(generationUrl)}，正在提交任务`);

      if (mode === "image") {
        const result = await aiGenerationClient.createImageSuperResolveTask({
          imageUrl: generationUrl,
          scale: imageScale,
        });
        await waitForTaskResult(result.taskId, "image");
      } else if (keepOriginalStyle) {
        const result = await aiGenerationClient.createVideoSuperResolveTask({
          videoUrl: generationUrl,
        });
        await waitForTaskResult(result.taskId, "video");
      } else {
        const result = await aiGenerationClient.createVideoSuperResolveTask({
          videoUrl: generationUrl,
          provider: "dashscope-style-transform",
          style: videoStyle,
          videoFps,
          minLen: videoMinLen,
          useSR,
          animateEmotion: true,
        });
        await waitForTaskResult(result.taskId, "video");
      }
    } catch (error) {
      setStatus(formatCreateTaskError(error, mode));
    } finally {
      setIsProcessing(false);
    }
  };

  const currentVideoStyleLabel = keepOriginalStyle
    ? "纯超分"
    : videoStyles.find((style) => style.value === videoStyle)?.label || "日式漫画";
  const scaleFactor = Number.parseInt(imageScale, 10);
  const sourceSizeText = sourceDimensions ? `原图 ${sourceDimensions.width}*${sourceDimensions.height}` : "原图";
  const resultSizeText = sourceDimensions
    ? `结果图 ${sourceDimensions.width * scaleFactor}*${sourceDimensions.height * scaleFactor}`
    : "结果图";
  const modeLabel = mode === "image" ? `图片超分 ${imageScale}` : keepOriginalStyle ? "纯超分（保持原风格）" : `${currentVideoStyleLabel} ${videoMinLen}P${useSR ? " + SR" : ""}`;

  return (
    <section className="image-workbench-page resolution-upscale-page">
      <header className="image-workbench-topbar">
        <button type="button" className="image-workbench-back-to-more" onClick={onOpenMore}>
          工具盒
        </button>
        <div className="image-workbench-tool-strip" aria-label="功能入口">
          <button type="button" onClick={() => onOpenImageTool?.("workbench")}>
            <EditOutlined />
            图片工作台
          </button>
          <button type="button" onClick={() => onOpenImageTool?.("inpaint")}>
            <ScissorOutlined />
            局部重绘
          </button>
          <button type="button" onClick={() => onOpenImageTool?.("camera")}>
            <CameraOutlined />
            镜头实验室
          </button>
          <button type="button" onClick={() => onSelectView?.("digitalHuman")}>
            <CustomerServiceOutlined />
            数字人
          </button>
          <button type="button" onClick={() => onSelectView?.("characterMix")}>
            <SwapOutlined />
            角色迁移
          </button>
          <button type="button" className="is-active" onClick={() => onSelectView?.("resolutionUpscale")}>
            <ColumnWidthOutlined />
            分辨率提升
          </button>
          <button type="button" onClick={() => onSelectView?.("watermarkRemoval")}>
            <DeleteOutlined />
            去水印
          </button>
          <button type="button" onClick={() => onSelectView?.("subtitleRemoval")}>
            <FontSizeOutlined />
            字幕去除
          </button>
        </div>
      </header>

      <div className="image-workbench-subbar">
        <button type="button" className="image-workbench-icon-btn" aria-label="返回工具盒" onClick={onOpenMore}>
          <ColumnWidthOutlined />
        </button>
        <strong>分辨率提升</strong>
        <div className="image-workbench-mode-tabs" role="tablist" aria-label="素材类型">
          <button type="button" className={mode === "image" ? "is-active" : ""} onClick={() => handleModeChange("image")}>
            <PictureOutlined />
            图片
          </button>
          <button type="button" className={mode === "video" ? "is-active" : ""} onClick={() => handleModeChange("video")}>
            <VideoCameraOutlined />
            视频
          </button>
        </div>
      </div>

      <main className="image-workbench-layout image-workbench-layout--camera resolution-upscale-layout">
        <aside className="image-workbench-panel image-workbench-panel--left">
          <section className="image-workbench-control-card">
            <div className="image-workbench-section-title">
              <h3>素材</h3>
              <span>{sourcePreview ? "已上传" : "待上传"}</span>
            </div>
            <input
              ref={fileInputRef}
              type="file"
              accept={mode === "image" ? "image/png,image/jpeg,image/webp" : "video/mp4,video/quicktime,video/webm,video/*"}
              onChange={handleFileChange}
            />
            <div
              className={`image-workbench-upload-shell${isDragging ? " is-dragging" : ""}`}
              onDragOver={handleDragOver}
              onDragLeave={handleDragLeave}
              onDrop={handleDrop}
            >
              {isDragging ? <div className="image-workbench-upload-drop-overlay"><span>释放文件以上传</span></div> : null}
              <button type="button" className="image-workbench-upload" onClick={() => fileInputRef.current?.click()}>
                {sourcePreview && mode === "image" ? <img src={sourcePreview} alt="" /> : <FileImageOutlined />}
                <strong>{sourceName || (mode === "image" ? "选择图片" : "选择视频")}</strong>
                <span>{mode === "image" ? "PNG / JPG / WebP" : "MP4 / MOV / WebM"}</span>
              </button>
              {sourcePreview ? (
                <button type="button" className="image-workbench-upload-remove" aria-label="删除素材" onClick={clearSource}>
                  ×
                </button>
              ) : null}
            </div>
            <div className="image-workbench-url-row">
              <label>
                <LinkOutlined />
                <input
                  value={sourceUrl}
                  placeholder="粘贴素材 URL"
                  onChange={(event) => setSourceUrl(event.target.value)}
                />
              </label>
              <button type="button" onClick={handleImportUrl}>导入</button>
            </div>
          </section>

          <section className="image-workbench-control-card">
            <h3>参数</h3>
            {mode === "image" ? (
              <label className="image-workbench-select">
                <span>放大倍数</span>
                <select value={imageScale} onChange={(event) => setImageScale(event.target.value as ImageScale)}>
                  <option value="2x">2x</option>
                  <option value="4x">4x</option>
                </select>
              </label>
            ) : (
              <>
                <div className="resolution-upscale-style-chips">
                  <span className="resolution-upscale-style-chips__title">模式</span>
                  <div className="resolution-upscale-style-chips__grid">
                    <button
                      type="button"
                      className={`resolution-upscale-style-chip${keepOriginalStyle ? " is-active" : ""}`}
                      onClick={() => setKeepOriginalStyle(true)}
                    >
                      <span className="resolution-upscale-style-chip__swatch" style={{ background: "#6ee7b7" }} />
                      保持原风格
                    </button>
                    {videoStyles.map((style) => (
                      <button
                        key={style.value}
                        type="button"
                        className={`resolution-upscale-style-chip${!keepOriginalStyle && videoStyle === style.value ? " is-active" : ""}`}
                        onClick={() => { setKeepOriginalStyle(false); setVideoStyle(style.value); }}
                      >
                        <span className="resolution-upscale-style-chip__swatch" style={{ background: style.color }} />
                        {style.label}
                      </button>
                    ))}
                  </div>
                </div>
                {!keepOriginalStyle && (
                  <>
                    <label className="image-workbench-select">
                      <span>输出短边</span>
                      <select value={videoMinLen} onChange={(event) => setVideoMinLen(Number(event.target.value) as VideoMinLen)}>
                        <option value={540}>540P</option>
                        <option value={720}>720P</option>
                      </select>
                    </label>
                    <label className="image-workbench-select">
                      <span>帧率</span>
                      <select value={videoFps} onChange={(event) => setVideoFps(Number(event.target.value))}>
                        <option value={15}>15 FPS</option>
                        <option value={20}>20 FPS</option>
                        <option value={25}>25 FPS</option>
                      </select>
                    </label>
                    <div className="image-workbench-section-title">
                      <h3>超分增强</h3>
                      <button
                        type="button"
                        className={`image-workbench-toggle ${useSR ? "is-active" : ""}`}
                        aria-pressed={useSR}
                        onClick={() => setUseSR((current) => !current)}
                      />
                    </div>
                  </>
                )}
              </>
            )}
          </section>

          <div className="image-workbench-actions">
            <button type="button" className="image-workbench-primary" onClick={() => void handleStart()} disabled={isProcessing}>
              <ColumnWidthOutlined />
              {isProcessing ? "处理中" : "开始超分"}
            </button>
            {isProcessing && (
              <button type="button" className="image-workbench-cancel" onClick={handleCancel} style={{ marginTop: 6 }}>
                取消
              </button>
            )}
          </div>
        </aside>

        <section className="image-workbench-canvas image-workbench-canvas--camera resolution-upscale-canvas" aria-label="分辨率提升画布">
          {sourcePreview ? (
            mode === "image" ? (
              <div className="resolution-upscale-image-stage">
                <BeforeAfterCompare
                  sourceSrc={sourcePreview}
                  resultSrc={resultPreview || sourcePreview}
                  sourceLabel={sourceSizeText}
                  resultLabel={resultPreview ? resultSizeText : "等待结果"}
                  sourceAlt="原图预览"
                  resultAlt="超分结果预览"
                  onSourceLoad={(width, height) => setSourceDimensions({ width, height })}
                />
                {resultPreview && (
                  <div className="resolution-upscale-result-actions">
                    <button type="button" onClick={() => void handleSaveToAssets()} disabled={isSavingAsset}>
                      <FolderAddOutlined />
                      {isSavingAsset ? "加入中" : "加入资产库"}
                    </button>
                    <button type="button" onClick={() => void handleDownload()} disabled={isDownloading}>
                      <DownloadOutlined />
                      {isDownloading ? "保存中" : "保存本地"}
                    </button>
                  </div>
                )}
              </div>
            ) : (
              <div className="resolution-upscale-video-stage">
                {resultPreview && (
                  <div className="resolution-upscale-video-header">
                    <div className="resolution-upscale-video-toggle">
                      <button
                        type="button"
                        className={videoViewMode === "source" ? "is-active" : ""}
                        onClick={() => setVideoViewMode("source")}
                      >
                        原始
                      </button>
                      <button
                        type="button"
                        className={videoViewMode === "result" ? "is-active" : ""}
                        onClick={() => setVideoViewMode("result")}
                      >
                        超分结果
                      </button>
                    </div>
                    <div className="resolution-upscale-result-actions">
                      <button type="button" onClick={() => void handleSaveToAssets()} disabled={isSavingAsset}>
                        <FolderAddOutlined />
                        {isSavingAsset ? "加入中" : "加入资产库"}
                      </button>
                      <button type="button" onClick={() => void handleDownload()} disabled={isDownloading}>
                        <DownloadOutlined />
                        {isDownloading ? "保存中" : "保存本地"}
                      </button>
                    </div>
                  </div>
                )}
                <video src={videoViewMode === "result" && resultPreview ? resultPreview : sourcePreview} controls muted playsInline />
                <span>{resultPreview ? (videoViewMode === "result" ? "超分结果" : "原始视频") : "原始视频"} · {modeLabel}</span>
              </div>
            )
          ) : (
            <div className="studio-canvas-ghost">
              <div className="studio-canvas-ghost__icon">
                <ThunderboltOutlined />
              </div>
              <div className="studio-canvas-ghost__title">{mode === "image" ? "拖拽或选择图片" : "拖拽或选择视频"}</div>
              <div className="studio-canvas-ghost__hint">{mode === "image" ? "支持 PNG / JPG / WebP" : "支持 MP4 / MOV / WebM"}</div>
            </div>
          )}
        </section>
      </main>

      <TaskStatusBar taskId={activeTaskId} status={status} progress={taskProgress} idleLabel={modeLabel} />
    </section>
  );
}

export default ResolutionUpscalePage;