import { and, eq, inArray } from "drizzle-orm";
import { fetchAtlasCloudImagePresets } 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 fetchAtlasCloudImagePresets();
  if (catalog.errors.length > 0) {
    throw new Error(
      `AtlasCloud 目录有 ${catalog.errors.length} 个 schema 加载失败：${catalog.errors
        .map((item) => item.modelId)
        .join(", ")}`
    );
  }

  const summaries = [];
  const catalogModelIds = new Set(
    catalog.presets.map((preset) => preset.modelId)
  );
  for (const provider of providers) {
    let added = 0;
    let updated = 0;
    for (const preset of catalog.presets) {
      const existing = await db.query.aiModel.findFirst({
        where: and(
          eq(aiModel.providerId, provider.id),
          eq(aiModel.modelId, preset.modelId)
        ),
      });
      const values = {
        displayName: preset.displayName,
        avatarUrl: preset.avatarUrl || null,
        description: preset.description,
        supportsVision: preset.supportsVision,
        supportsImageGen: preset.supportsImageGen,
        isEnabled: true,
        isHidden: false,
        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,
          isDefault: false,
          isSubModel: false,
          parentModelId: null,
        });
        added += 1;
      }
    }
    const providerModels = await db
      .select({ modelId: aiModel.modelId })
      .from(aiModel)
      .where(eq(aiModel.providerId, provider.id));
    const staleModelIds = providerModels
      .map((model) => model.modelId)
      .filter((modelId) => !catalogModelIds.has(modelId));
    if (staleModelIds.length > 0) {
      await db
        .update(aiModel)
        .set({
          isEnabled: false,
          isHidden: true,
          updatedAt: new Date(),
        })
        .where(
          and(
            eq(aiModel.providerId, provider.id),
            inArray(aiModel.modelId, staleModelIds)
          )
        );
    }
    summaries.push({
      providerId: provider.id,
      providerName: provider.name,
      added,
      updated,
      disabledStale: staleModelIds.length,
    });
  }

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

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