ai-chatbot/app/api/worldlabs/key-pool/import/route.ts

import { z } from "zod";
import {
  createAiProvider,
  getAiProvidersByType,
} from "@/lib/db/queries";
import { ChatSDKError } from "@/lib/errors";
import { requireAdmin } from "@/lib/middleware/admin";
import {
  sanitizeWorldLabsApiKey,
  sanitizeWorldLabsBaseUrl,
} from "@/lib/worldlabs/server";
import { getWorldLabsKeyPoolConfig } from "@/lib/worldlabs/key-pool-config";

export const runtime = "nodejs";
export const maxDuration = 60;

const DEFAULT_WORLDLABS_BASE_URL = "https://api.worldlabs.ai";

const importPayloadSchema = z.object({
  keys: z
    .array(
      z.union([
        z.string().min(1),
        z.object({
          apiKey: z.string().min(1),
          name: z.string().min(1).max(128).optional(),
          baseUrl: z.string().min(1).optional(),
          isEnabled: z.boolean().optional(),
        }),
      ])
    )
    .max(5000)
    .optional(),
  apiKeysText: z.string().optional(),
  namePrefix: z.string().min(1).max(96).optional(),
  baseUrl: z.string().min(1).optional(),
  isEnabled: z.boolean().optional(),
});

function readImportToken(request: Request) {
  const authHeader = request.headers.get("authorization") || "";
  const bearerToken = authHeader.match(/^Bearer\s+(.+)$/i)?.[1]?.trim();
  const customToken =
    request.headers.get("x-worldlabs-import-token")?.trim() || "";
  return bearerToken || customToken || "";
}

function getConfiguredImportToken() {
  return String(
    process.env.WORLDLABS_KEY_POOL_IMPORT_TOKEN ||
      process.env.WORLDLABS_IMPORT_TOKEN ||
      ""
  ).trim();
}

function normalizeImportItems(payload: z.infer<typeof importPayloadSchema>) {
  const items: Array<{
    apiKey: string;
    name?: string;
    baseUrl?: string;
    isEnabled?: boolean;
  }> = [];

  for (const item of payload.keys || []) {
    if (typeof item === "string") {
      items.push({ apiKey: item });
    } else {
      items.push(item);
    }
  }

  const lines = String(payload.apiKeysText || "")
    .split(/\r?\n/g)
    .map((line) => line.trim())
    .filter(Boolean);
  for (const line of lines) {
    items.push({ apiKey: line });
  }

  return items;
}

function buildProviderName(params: {
  namePrefix: string;
  index: number;
  overrideName?: string;
}) {
  const customName = String(params.overrideName || "").trim();
  if (customName) return customName.slice(0, 128);

  const datePart = new Date()
    .toISOString()
    .replace(/[-:TZ.]/g, "")
    .slice(0, 14);
  return `${params.namePrefix}-${datePart}-${String(params.index + 1).padStart(4, "0")}`.slice(
    0,
    128
  );
}

function maskKey(value: string) {
  if (value.length <= 10) {
    return `${value.slice(0, 2)}****`;
  }
  return `${value.slice(0, 4)}****${value.slice(-4)}`;
}

export async function POST(request: Request) {
  try {
    const dbConfig = await getWorldLabsKeyPoolConfig();
    const configuredToken = dbConfig.importToken || getConfiguredImportToken();
    const requestToken = readImportToken(request);

    let authorizedByAdmin = false;
    try {
      await requireAdmin();
      authorizedByAdmin = true;
    } catch (_error) {
      authorizedByAdmin = false;
    }

    if (!authorizedByAdmin) {
      if (!configuredToken) {
        return Response.json(
          {
            success: false,
            error:
              "服务端未配置 World Labs key 池导入令牌，无法使用专用导入接口",
          },
          { status: 503 }
        );
      }

      if (!requestToken || requestToken !== configuredToken) {
        return Response.json(
          { success: false, error: "导入令牌无效" },
          { status: 401 }
        );
      }
    }

    const rawBody = await request.json().catch(() => ({}));
    const payload = importPayloadSchema.parse(rawBody);
    const inputItems = normalizeImportItems(payload);

    if (inputItems.length === 0) {
      return Response.json(
        { success: false, error: "请至少提供一个 World Labs API Key" },
        { status: 400 }
      );
    }

    const existingProviders = await getAiProvidersByType({
      type: "worldlabs-3d",
      enabledOnly: false,
    });
    const existingKeySet = new Set(
      existingProviders
        .map((provider) => sanitizeWorldLabsApiKey(String(provider.apiKey || "")))
        .filter(Boolean)
    );

    const defaultBaseUrl =
      sanitizeWorldLabsBaseUrl(payload.baseUrl || "") ||
      DEFAULT_WORLDLABS_BASE_URL;
    const defaultIsEnabled = payload.isEnabled !== false;
    const namePrefix = String(payload.namePrefix || "WorldLabsKey").trim();

    const dedupedImports: Array<{
      apiKey: string;
      name?: string;
      baseUrl: string;
      isEnabled: boolean;
    }> = [];
    const requestKeySet = new Set<string>();
    const skipped: Array<{ apiKey: string; reason: string }> = [];

    for (const item of inputItems) {
      const apiKey = sanitizeWorldLabsApiKey(item.apiKey);
      if (!apiKey) {
        skipped.push({ apiKey: "", reason: "empty_key" });
        continue;
      }
      if (existingKeySet.has(apiKey)) {
        skipped.push({ apiKey: maskKey(apiKey), reason: "duplicate_existing" });
        continue;
      }
      if (requestKeySet.has(apiKey)) {
        skipped.push({ apiKey: maskKey(apiKey), reason: "duplicate_in_request" });
        continue;
      }

      requestKeySet.add(apiKey);
      dedupedImports.push({
        apiKey,
        name: item.name,
        baseUrl:
          sanitizeWorldLabsBaseUrl(item.baseUrl || "") || defaultBaseUrl,
        isEnabled: item.isEnabled ?? defaultIsEnabled,
      });
    }

    const imported = [];
    for (const [index, item] of dedupedImports.entries()) {
      const created = await createAiProvider({
        name: buildProviderName({
          namePrefix,
          index,
          overrideName: item.name,
        }),
        type: "worldlabs-3d",
        baseUrl: item.baseUrl,
        apiKey: item.apiKey,
        isEnabled: item.isEnabled,
      });
      imported.push({
        id: created[0]?.id,
        name: created[0]?.name,
        isEnabled: created[0]?.isEnabled === true,
      });
    }

    return Response.json({
      success: true,
      data: {
        importedCount: imported.length,
        skippedCount: skipped.length,
        imported,
        skipped,
      },
    });
  } catch (error) {
    if (error instanceof ChatSDKError) {
      return error.toResponse();
    }
    if (error instanceof z.ZodError) {
      return Response.json(
        { success: false, error: error.flatten() },
        { status: 400 }
      );
    }
    return Response.json(
      {
        success: false,
        error:
          error instanceof Error ? error.message : "导入 World Labs key 失败",
      },
      { status: 500 }
    );
  }
}