import { and, eq } from "drizzle-orm";
import {
  fetchAtlasCloudModelPresets,
  resolveAtlasCloudDisplayName,
  shouldInsertMissingAtlasCloudPreset,
} from "../lib/ai/atlascloud-models";
import { db } from "../lib/db";
import { aiModel, aiProvider } from "../lib/db/schema";

async function main() {
  const providers = await db
    .select()
    .from(aiProvider)
    .where(and(eq(aiProvider.type, "atlascloud"), eq(aiProvider.isEnabled, true)));
  if (providers.length === 0) {
    throw new Error("没有找到已启用的 AtlasCloud 服务商");
  }

  const catalog = await fetchAtlasCloudModelPresets();
  if (catalog.errors.length > 0) {
    throw new Error(
      `AtlasCloud 目录有 ${catalog.errors.length} 个 schema 加载失败：${catalog.errors
        .map((item) => item.modelId)
        .join(", ")}`
    );
  }

  const summaries = [];
  for (const provider of providers) {
    let added = 0;
    let updated = 0;
    let preservedDeleted = 0;
    const providerModels = await db
      .select()
      .from(aiModel)
      .where(eq(aiModel.providerId, provider.id));
    const existingByModelId = new Map(
      providerModels.map((model) => [model.modelId, model])
    );
    const existingMediaKinds = new Set(
      providerModels.map((model) => {
        const capabilities = model.capabilities as Record<string, unknown> | null;
        return String(capabilities?.mediaType || "image");
      })
    );
    for (const preset of catalog.presets) {
      const existing = existingByModelId.get(preset.modelId);
      const presetCapabilities = preset.capabilities as Record<string, unknown>;
      const mediaKind = String(presetCapabilities.mediaType || "image");
      if (
        !existing &&
        !shouldInsertMissingAtlasCloudPreset(existingMediaKinds, mediaKind)
      ) {
        preservedDeleted += 1;
        continue;
      }
      const values = {
        displayName: resolveAtlasCloudDisplayName(
          existing?.displayName,
          preset.displayName,
          preset.modelId
        ),
        avatarUrl: preset.avatarUrl || null,
        description: preset.description,
        supportsVision: preset.supportsVision,
        supportsImageGen: preset.supportsImageGen,
        sortOrder: existing?.sortOrder ?? preset.sortOrder,
        capabilities: preset.capabilities,
        defaultParams: preset.defaultParams,
        imageGenConfig: preset.imageGenConfig,
        updatedAt: new Date(),
      };

      if (existing) {
        await db
          .update(aiModel)
          .set(values)
          .where(eq(aiModel.id, existing.id));
        updated += 1;
      } else {
        await db.insert(aiModel).values({
          providerId: provider.id,
          modelId: preset.modelId,
          ...values,
          isEnabled: true,
          isHidden: false,
          isDefault: false,
          isSubModel: false,
          parentModelId: null,
        });
        added += 1;
      }
    }
    summaries.push({
      providerId: provider.id,
      providerName: provider.name,
      added,
      updated,
      preservedDeleted,
    });
  }

  console.log(
    JSON.stringify(
      {
        catalogModels: catalog.catalogCount,
        availableModels: catalog.presets.length,
        skippedUnsupportedVideoModels: catalog.skipped.length,
        providers: summaries,
      },
      null,
      2
    )
  );
}

void main().catch((error) => {
  console.error(error instanceof Error ? error.message : String(error));
  process.exitCode = 1;
});