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

type StreamUsage = {
  inputTokens?: number;
  outputTokens?: number;
  totalTokens?: number;
};

export type StreamEvent =
  | { type: "meta"; chatId: string | null; callId: string | null; provider: Provider; model: string }
  | { type: "tool_call"; event: ToolExecutionEvent }
  | { type: "delta"; text: string }
  | { type: "done"; text: string; usage?: StreamUsage }
  | { type: "error"; message: string };

function getChatIdOrCreate(chatId?: string) {
  if (chatId) return Promise.resolve(chatId);
  return prisma.chat.create({ data: {}, select: { id: true } }).then((c) => c.id);
}

export async function* runMultiplexStream(req: MultiplexRequest): AsyncGenerator<StreamEvent> {
  const t0 = performance.now();
  const shouldPersist = req.persist !== false;
  const chatId = shouldPersist ? await getChatIdOrCreate(req.chatId) : null;

  const call =
    shouldPersist && chatId
      ? await prisma.llmCall.create({
          data: {
            chatId,
            provider: toPrismaProvider(req.provider) as any,
            model: req.model,
            request: req as any,
          },
          select: { id: true },
        })
      : null;

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

  yield { type: "meta", chatId, callId: call?.id ?? null, provider: req.provider, model: req.model };

  let text = "";
  let usage: StreamUsage | undefined;
  let raw: unknown = { streamed: true };

  try {
    const adapter = getProviderChatAdapter(req.provider);
    const streamEvents = adapter.stream({
      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: chatId ?? undefined,
      },
    });

    for await (const ev of streamEvents) {
      if (ev.type === "delta") {
        text += ev.text;
        yield { type: "delta", text: ev.text };
        continue;
      }

      if (ev.type === "tool_call") {
        if (ev.event.status !== "initiated" && shouldPersist && chatId) {
          const toolMessage = buildToolLogMessageData(chatId, ev.event);
          await prisma.message.create({
            data: {
              chatId: toolMessage.chatId,
              role: toolMessage.role as any,
              content: toolMessage.content,
              name: toolMessage.name,
              metadata: toolMessage.metadata as any,
            },
          });
        }
        yield { type: "tool_call", event: ev.event };
        continue;
      }

      raw = ev.result.raw;
      usage = ev.result.usage;
      text = ev.result.text;
    }

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

    if (shouldPersist && chatId && call) {
      await prisma.$transaction(async (tx) => {
        await tx.message.create({
          data: {
            chatId,
            role: "assistant" as any,
            content: text,
            metadata: req.clientRequestId ? ({ clientRequestId: req.clientRequestId } as any) : undefined,
          },
        });
        await tx.llmCall.update({
          where: { id: call.id },
          data: {
            response: raw as any,
            latencyMs,
            inputTokens: usage?.inputTokens,
            outputTokens: usage?.outputTokens,
            totalTokens: usage?.totalTokens,
          },
        });
      });
    }

    yield { type: "done", text, usage };
  } catch (e: any) {
    const latencyMs = Math.round(performance.now() - t0);
    if (shouldPersist && call) {
      await prisma.llmCall.update({
        where: { id: call.id },
        data: {
          error: e?.message ?? String(e),
          latencyMs,
        },
      });
    }
    yield { type: "error", message: e?.message ?? String(e) };
  }
}
