import { defineWrappedResponseHandler } from "#server/utils/handler";
import { R } from "#server/utils/response";
import { requireAdmin } from "#server/utils/admin-guard";
import { listAgents } from "#server/service/agent/agent";
import { dbGlobal } from "drizzle-pkg/lib/db";
import { agentToolAssociations } from "drizzle-pkg/lib/schema/agent";
import { llmModels, llmProviders } from "drizzle-pkg/lib/schema/llm";
import { eq, sql, inArray } from "drizzle-orm";

export default defineWrappedResponseHandler(async (event) => {
  await requireAdmin(event);

  const agents = await listAgents(false);

  if (agents.length === 0) {
    return R.success({ list: [] });
  }

  const agentIds = agents.map((a) => a.id);
  const countRows = await dbGlobal
    .select({
      agentId: agentToolAssociations.agentId,
      count: sql<number>`count(*)`.as("count"),
    })
    .from(agentToolAssociations)
    .where(inArray(agentToolAssociations.agentId, agentIds))
    .groupBy(agentToolAssociations.agentId);

  const countMap = new Map<number, number>();
  for (const row of countRows) {
    countMap.set(row.agentId, Number(row.count));
  }

  const modelIds = agents
    .flatMap((a) => [a.defaultModelId, a.titleModelId])
    .filter((id): id is number => id != null);

  const modelMap = new Map<number, {
    modelName: string;
    modelId: string;
    providerName: string;
    providerStatus: string;
    apiKeySet: boolean;
  }>();

  if (modelIds.length > 0) {
    const modelRows = await dbGlobal
      .select({
        modelId: llmModels.id,
        modelName: llmModels.name,
        modelIdentifier: llmModels.modelId,
        providerName: llmProviders.name,
        providerStatus: llmProviders.status,
        apiKey: llmProviders.apiKey,
      })
      .from(llmModels)
      .innerJoin(llmProviders, eq(llmModels.providerId, llmProviders.id))
      .where(inArray(llmModels.id, modelIds));

    for (const row of modelRows) {
      modelMap.set(row.modelId, {
        modelName: row.modelName,
        modelId: row.modelIdentifier,
        providerName: row.providerName,
        providerStatus: row.providerStatus,
        apiKeySet: !!row.apiKey,
      });
    }
  }

  const list = agents.map((a) => ({
    ...a,
    toolCount: countMap.get(a.id) ?? 0,
    defaultModel: a.defaultModelId ? modelMap.get(a.defaultModelId) ?? null : null,
    titleModel: a.titleModelId ? modelMap.get(a.titleModelId) ?? null : null,
  }));

  return R.success({ list });
});