omniai-web/src/features/workbench/workbenchDownload.ts

import { aiGenerationClient } from "../../api/aiGenerationClient";
import { buildApiUrl, buildAuthHeaders } from "../../api/serverConnection";

function sanitizeDownloadFilename(value: string, fallback: string) {
  const normalized = value
    .trim()
    .replace(/[\\/:*?"<>|]+/g, "-")
    .replace(/\s+/g, " ")
    .slice(0, 80)
    .trim();
  return normalized || fallback;
}

function getResultDownloadExtension(url: string, contentType: string | null, isVideo: boolean) {
  const mime = (contentType || "").split(";")[0]?.trim().toLowerCase();
  if (mime === "image/jpeg") return "jpg";
  if (mime === "image/png") return "png";
  if (mime === "image/webp") return "webp";
  if (mime === "image/gif") return "gif";
  if (mime === "video/webm") return "webm";
  if (mime === "video/quicktime") return "mov";
  if (mime === "video/mp4") return "mp4";

  try {
    const pathname = new URL(url, window.location.href).pathname;
    const matched = pathname.match(/\.([a-z0-9]{2,5})$/i);
    if (matched?.[1]) return matched[1].toLowerCase();
  } catch {
    // Keep the fallback below for malformed but still browser-loadable URLs.
  }

  return isVideo ? "mp4" : "png";
}

type LocalFilePickerWritable = {
  write: (data: Blob) => Promise<void>;
  close: () => Promise<void>;
};

type LocalFilePickerHandle = {
  createWritable: () => Promise<LocalFilePickerWritable>;
};

type LocalFilePickerWindow = Window & {
  showSaveFilePicker?: (options?: {
    suggestedName?: string;
    types?: Array<{
      description: string;
      accept: Record<string, string[]>;
    }>;
  }) => Promise<LocalFilePickerHandle>;
};

async function saveBlobToLocal(blob: Blob, filename: string, isVideo: boolean): Promise<"saved" | "started"> {
  const startBrowserDownload = () => {
    const objectUrl = URL.createObjectURL(blob);
    const anchor = document.createElement("a");
    anchor.href = objectUrl;
    anchor.download = filename;
    anchor.rel = "noopener";
    document.body.appendChild(anchor);
    anchor.click();
    anchor.remove();

    window.setTimeout(() => URL.revokeObjectURL(objectUrl), 30_000);
    return "started" as const;
  };

  const savePicker = (window as LocalFilePickerWindow).showSaveFilePicker;
  if (savePicker) {
    try {
      const pickerTypes = [
        isVideo
          ? { description: "Video", accept: { [blob.type || "video/mp4"]: [`.${filename.split(".").pop() || "mp4"}`] } }
          : { description: "Image", accept: { [blob.type || "image/png"]: [`.${filename.split(".").pop() || "png"}`] } },
      ];
      const handle = await savePicker({ suggestedName: filename, types: pickerTypes });
      const writable = await handle.createWritable();
      await writable.write(blob);
      await writable.close();
      return "saved";
    } catch (error) {
      if (error instanceof DOMException && error.name === "AbortError") throw error;
      return startBrowserDownload();
    }
  }

  return startBrowserDownload();
}

async function assertDownloadBlobIsUsable(blob: Blob, isVideo: boolean) {
  const mime = String(blob.type || "").toLowerCase();
  if (/^(?:application|text)\/(?:json|xml|html|plain)|\+xml/.test(mime)) {
    throw new Error("下载失败：结果链接已过期或返回了错误内容，请重新生成后再下载。");
  }

  const header = new Uint8Array(await blob.slice(0, 16).arrayBuffer());
  const asText = new TextDecoder("utf-8", { fatal: false }).decode(header);
  const isPng = header[0] === 0x89 && header[1] === 0x50 && header[2] === 0x4e && header[3] === 0x47;
  const isJpeg = header[0] === 0xff && header[1] === 0xd8 && header[2] === 0xff;
  const isGif = asText.startsWith("GIF8");
  const isWebp = asText.startsWith("RIFF") && asText.slice(8, 12) === "WEBP";
  const isMp4 = asText.slice(4, 8) === "ftyp";
  const looksLikeErrorDocument = asText.startsWith("<?xml") || asText.startsWith("<Error") || asText.startsWith("{\"") || asText.startsWith("<!DO");

  if (looksLikeErrorDocument) {
    throw new Error("下载失败：结果链接已过期或返回了错误内容，请重新生成后再下载。");
  }

  const hasExpectedMagic = isVideo ? isMp4 || mime.startsWith("video/") : isPng || isJpeg || isGif || isWebp || mime.startsWith("image/");
  if (!hasExpectedMagic) {
    throw new Error("下载失败：未收到有效的媒体文件。");
  }
}

export async function downloadResultAsset(url: string, filenameBase: string, isVideo: boolean, taskId?: string): Promise<"saved" | "started"> {
  let blob: Blob | null = null;
  let contentType: string | null = null;

  if (taskId) {
    try {
      const result = await aiGenerationClient.downloadTaskResult(taskId);
      blob = result.blob;
      contentType = result.contentType || result.blob.type || null;
    } catch {
      blob = null;
    }
  }

  if (!blob) {
    const needsProxy = /aliyuncs\.com/i.test(url);
    const fetchUrl = needsProxy
      ? buildApiUrl(`ai/proxy-download?url=${encodeURIComponent(url)}`)
      : url;
    const fetchOpts: RequestInit = needsProxy
      ? { headers: buildAuthHeaders() }
      : { credentials: "omit" };
    const response = await fetch(fetchUrl, fetchOpts);
    if (!response.ok) {
      throw new Error(`下载失败：资源请求返回 ${response.status}`);
    }
    blob = await response.blob();
    contentType = blob.type || response.headers.get("content-type");
  }

  if (!blob.size) {
    throw new Error("下载失败：资源内容为空");
  }
  await assertDownloadBlobIsUsable(blob, isVideo);

  const extension = getResultDownloadExtension(url, contentType, isVideo);
  const filename = `${sanitizeDownloadFilename(filenameBase, isVideo ? "generated-video" : "generated-image")}.${extension}`;
  return saveBlobToLocal(blob, filename, isVideo);
}