lovart3/app/api/watermark-remove/route.ts

import { NextRequest, NextResponse } from "next/server";
import { auth } from "@/app/(auth)/auth";
import { db } from "@/lib/db";
import { systemSettings } from "@/lib/db/schema";
import { eq } from "drizzle-orm";
import { uploadFile } from "@/lib/storage/service";

// 获取去水印API配置
async function getWatermarkRemoveConfig() {
  try {
    const setting = await db
      .select()
      .from(systemSettings)
      .where(eq(systemSettings.key, "watermark_remove_config"))
      .limit(1);

    if (setting.length > 0 && setting[0].value) {
      return setting[0].value as { apiKey: string };
    }
    return null;
  } catch (error) {
    console.error("[WatermarkRemove] 获取配置失败:", error);
    return null;
  }
}

// 轮询获取处理结果
async function pollForResult(taskId: string, apiKey: string): Promise<{ success: boolean; imageUrl?: string; error?: string }> {
  const maxAttempts = 30; // 最多轮询30次（30秒）
  const pollInterval = 1000; // 每次间隔1秒

  for (let attempt = 0; attempt < maxAttempts; attempt++) {
    try {
      const response = await fetch(
        `https://techsz.aoscdn.com/api/tasks/visual/external/watermark-remove/${taskId}`,
        {
          method: "GET",
          headers: {
            "X-API-KEY": apiKey,
          },
        }
      );

      if (!response.ok) {
        const errorText = await response.text();
        console.error(`[WatermarkRemove] 轮询失败 (attempt ${attempt + 1}):`, errorText);
        // 继续轮询
        await new Promise(resolve => setTimeout(resolve, pollInterval));
        continue;
      }

      const result = await response.json();
      console.log(`[WatermarkRemove] 轮询结果 (attempt ${attempt + 1}):`, JSON.stringify(result).substring(0, 200));

      // 检查任务状态
      if (result.data?.state === 1) {
        // 任务完成
        const imageUrl = result.data?.image;
        if (imageUrl) {
          return { success: true, imageUrl };
        }
        return { success: false, error: "处理完成但未返回图片URL" };
      } else if (result.data?.state === -1) {
        // 任务失败
        return { success: false, error: result.data?.error || "去水印处理失败" };
      }
      // state === 0 表示处理中，继续轮询

      await new Promise(resolve => setTimeout(resolve, pollInterval));
    } catch (error) {
      console.error(`[WatermarkRemove] 轮询错误 (attempt ${attempt + 1}):`, error);
      await new Promise(resolve => setTimeout(resolve, pollInterval));
    }
  }

  return { success: false, error: "处理超时，请稍后重试" };
}

export async function POST(request: NextRequest) {
  try {
    // 验证用户登录
    const session = await auth();
    if (!session?.user?.id) {
      return NextResponse.json({ error: "请先登录" }, { status: 401 });
    }

    const startTime = Date.now();

    // 获取请求参数
    const body = await request.json();
    const { image_url, image_base64 } = body;

    if (!image_url && !image_base64) {
      return NextResponse.json(
        { error: "请提供图片URL或Base64数据" },
        { status: 400 }
      );
    }

    // 获取API配置
    const config = await getWatermarkRemoveConfig();
    if (!config?.apiKey) {
      return NextResponse.json(
        { error: "去水印服务未配置，请联系管理员在系统设置中配置API密钥" },
        { status: 500 }
      );
    }

    // 如果传入的是base64，需要先上传到S3获取URL
    let imageUrl = image_url;
    if (image_base64 && !image_url) {
      try {
        const buffer = Buffer.from(image_base64, "base64");
        const arrayBuffer = buffer.buffer.slice(buffer.byteOffset, buffer.byteOffset + buffer.byteLength);
        const filename = `watermark-temp-${Date.now()}.jpeg`;
        const uploadResult = await uploadFile(filename, arrayBuffer);
        imageUrl = uploadResult.url;
        console.log("[WatermarkRemove] Base64图片已上传到S3:", imageUrl);
      } catch (uploadError) {
        console.error("[WatermarkRemove] 上传图片失败:", uploadError);
        return NextResponse.json(
          { error: "上传图片失败" },
          { status: 500 }
        );
      }
    }

    console.log("[WatermarkRemove] 开始处理图片:", imageUrl);

    // 创建去水印任务
    const formData = new FormData();
    formData.append("sync", "0");
    formData.append("url", imageUrl);

    const createResponse = await fetch(
      "https://techsz.aoscdn.com/api/tasks/visual/external/watermark-remove",
      {
        method: "POST",
        headers: {
          "X-API-KEY": config.apiKey,
        },
        body: formData,
      }
    );

    if (!createResponse.ok) {
      const errorText = await createResponse.text();
      console.error("[WatermarkRemove] 创建任务失败:", errorText);
      return NextResponse.json(
        { error: "创建去水印任务失败" },
        { status: 500 }
      );
    }

    const createResult = await createResponse.json();
    console.log("[WatermarkRemove] 创建任务结果:", JSON.stringify(createResult).substring(0, 200));

    const taskId = createResult.data?.task_id;
    if (!taskId) {
      return NextResponse.json(
        { error: "未获取到任务ID" },
        { status: 500 }
      );
    }

    // 轮询获取结果
    const pollResult = await pollForResult(taskId, config.apiKey);

    if (!pollResult.success) {
      return NextResponse.json(
        { error: pollResult.error || "去水印处理失败" },
        { status: 500 }
      );
    }

    // 下载处理后的图片并上传到S3
    const processedImageUrl = pollResult.imageUrl!;
    console.log("[WatermarkRemove] 处理完成，结果URL:", processedImageUrl);

    try {
      // 下载处理后的图片
      const imageResponse = await fetch(processedImageUrl);
      if (!imageResponse.ok) {
        throw new Error("下载处理后的图片失败");
      }
      const imageBuffer = await imageResponse.arrayBuffer();

      // 上传到S3
      const filename = `watermark-removed-${Date.now()}.png`;
      const uploadResult = await uploadFile(filename, imageBuffer);

      const processingTime = (Date.now() - startTime) / 1000;
      console.log(`[WatermarkRemove] 处理完成，耗时: ${processingTime}s`);

      return NextResponse.json({
        success: true,
        original_url: uploadResult.url,
        thumbnail_url: uploadResult.url,
        processing_time: processingTime,
      });
    } catch (uploadError) {
      console.error("[WatermarkRemove] 上传结果失败:", uploadError);
      // 如果上传失败，直接返回原始URL
      const processingTime = (Date.now() - startTime) / 1000;
      return NextResponse.json({
        success: true,
        original_url: processedImageUrl,
        thumbnail_url: processedImageUrl,
        processing_time: processingTime,
      });
    }
  } catch (error) {
    console.error("[WatermarkRemove] 处理错误:", error);
    return NextResponse.json(
      { error: error instanceof Error ? error.message : "去水印处理失败" },
      { status: 500 }
    );
  }
}