import { performance } from "node:perf_hooks";
import { prisma } from "../db.js";
import { buildToolLogMessageData } from "./chat-tools.js";
import { getProviderChatAdapter } from "./provider-adapters.js";
import { toPrismaProvider } from "./provider-ids.js";
import type { MultiplexRequest, MultiplexResponse, Provider } from "./types.js";

function asProviderEnum(p: Provider) {
  return toPrismaProvider(p);
}

export async function runMultiplex(req: MultiplexRequest): Promise<MultiplexResponse> {
  const t0 = performance.now();
  const chatId = req.chatId ?? (await prisma.chat.create({ data: {}, select: { id: true } })).id;

  // Persist call record early so we can attach errors.
  const call = await prisma.llmCall.create({
    data: {
      chatId,
      provider: asProviderEnum(req.provider) as any,
      model: req.model,
      request: req as any,
    },
    select: { id: true, chatId: true },
  });

  await prisma.$transaction([
    prisma.chat.update({
      where: { id: chatId },
      data: {
        lastUsedProvider: asProviderEnum(req.provider) as any,
        lastUsedModel: req.model,
      },
    }),
    prisma.chat.updateMany({
      where: { id: chatId, initiatedProvider: null },
      data: {
        initiatedProvider: asProviderEnum(req.provider) as any,
        initiatedModel: req.model,
      },
    }),
  ]);

  try {
    let outText = "";
    let usage: MultiplexResponse["usage"] | undefined;
    let raw: unknown;
    let toolMessages: ReturnType<typeof buildToolLogMessageData>[] = [];
    const adapter = getProviderChatAdapter(req.provider);
    const r = await adapter.complete({
      model: req.model,
      messages: req.messages,
      enabledTools: req.enabledTools,
      userLocation: req.userLocation,
      temperature: req.temperature,
      maxTokens: req.maxTokens,
      logContext: {
        provider: req.provider,
        model: req.model,
        chatId,
      },
    });
    raw = r.raw;
    outText = r.text;
    usage = r.usage;
    toolMessages = r.toolEvents.map((event) => buildToolLogMessageData(call.chatId, event));

    const latencyMs = Math.round(performance.now() - t0);

    // Store tool activity (if any), assistant message, and call record.
    await prisma.$transaction(async (tx) => {
      if (toolMessages.length) {
        await tx.message.createMany({
          data: toolMessages.map((message) => ({
            chatId: message.chatId,
            role: message.role as any,
            content: message.content,
            name: message.name,
            metadata: message.metadata as any,
          })),
        });
      }
      await tx.message.create({
        data: {
          chatId: call.chatId,
          role: "assistant" as any,
          content: outText,
        },
      });
      await tx.llmCall.update({
        where: { id: call.id },
        data: {
          response: raw as any,
          latencyMs,
          inputTokens: usage?.inputTokens,
          outputTokens: usage?.outputTokens,
          totalTokens: usage?.totalTokens,
        },
      });
    });

    return {
      provider: req.provider,
      model: req.model,
      message: { role: "assistant", content: outText },
      usage,
      raw,
    };
  } catch (e: any) {
    const latencyMs = Math.round(performance.now() - t0);
    await prisma.llmCall.update({
      where: { id: call.id },
      data: {
        error: e?.message ?? String(e),
        latencyMs,
      },
    });
    throw e;
  }
}