import {
  convertToModelMessages,
  createUIMessageStream,
  createUIMessageStreamResponse,
  smoothStream,
  stepCountIs,
  streamText,
  tool,
} from "ai";

import { eq } from "@turbostarter/db";
import { chat, message, part } from "@turbostarter/db/schema/chat";
import { db } from "@turbostarter/db/server";

import { modelStrategies } from "../chat/strategies";
import { Model, Role } from "../chat/types";

import {
  graphPatchSchema,
  validateEdgeReferences,
  validateEdgeTypes,
  validateNodeTypes,
  validateUniqueIds,
} from "./mutation-schema";
import { buildCopilotSystemPrompt, buildSelectedContext } from "./system-prompt";

import type { CopilotMessagePayload } from "./schema";
import type { DiagramType } from "./types";
import type {
  InsertChat,
  InsertMessage,
  InsertPart,
} from "@turbostarter/db/schema/chat";

const DEFAULT_MODEL = Model.CLAUDE_4_SONNET;

const generateDiagramTool = tool({
  description:
    "Generate or modify a diagram with the given graph data. Call this when the user asks to create, restructure, or modify a diagram.",
  inputSchema: graphPatchSchema,
  execute: async (patch) => {
    const errors = [
      ...validateUniqueIds(patch),
      ...validateNodeTypes(patch),
      ...validateEdgeTypes(patch),
      ...validateEdgeReferences(patch),
    ];
    if (errors.length > 0) {
      return { success: false as const, errors };
    }
    return { success: true as const, data: patch };
  },
});

const createCopilotChat = async (data: InsertChat) =>
  db
    .insert(chat)
    .values(data)
    .onConflictDoUpdate({
      target: chat.id,
      set: data,
    })
    .returning();

const createMessage = async (data: InsertMessage) =>
  db.insert(message).values(data).onConflictDoUpdate({
    target: message.id,
    set: data,
  });

const createParts = async (data: InsertPart[]) =>
  db.insert(part).values(data).onConflictDoNothing();

const getChatMessages = async (chatId: string) =>
  db.query["chat.message"].findMany({
    where: eq(message.chatId, chatId),
    orderBy: (message, { asc }) => [asc(message.createdAt)],
    with: {
      part: {
        orderBy: (part, { asc }) => [asc(part.order)],
      },
    },
  });

const toChatMessage = (msg: Awaited<ReturnType<typeof getChatMessages>>[number]) => ({
  id: msg.id,
  role: msg.role as "user" | "assistant",
  parts: msg.part.map((p) => {
    const details = p.details as Record<string, unknown>;

    // Preserve tool invocation parts so the AI retains generation context
    if (p.type.startsWith("tool-")) {
      return {
        type: p.type as "text",
        ...(details as Record<string, unknown>),
      };
    }

    return {
      type: "text" as const,
      text: (details.text as string) ?? "",
    };
  }),
});

export const getCopilotHistory = async (chatId: string) => {
  const messages = await getChatMessages(chatId);
  return messages.map(toChatMessage);
};

export const streamCopilot = async ({
  chatId,
  diagramId,
  diagramType,
  graphContext,
  selectedElements,
  userId,
  signal,
  ...msg
}: CopilotMessagePayload & { signal: AbortSignal; userId: string }) => {
  // Upsert chat record (linked to diagram)
  await createCopilotChat({
    id: chatId,
    userId,
    diagramId,
    name: `Copilot — ${diagramType.toUpperCase()}`,
  });

  // Load existing messages for context
  const existingMessages = await getChatMessages(chatId);

  // Persist the user message
  await createMessage({ id: msg.id, chatId, role: Role.USER });
  await createParts(
    msg.parts.map(({ type, ...details }, order) => ({
      type,
      order,
      details,
      messageId: msg.id,
    })),
  );

  const systemPrompt = buildCopilotSystemPrompt(diagramType as DiagramType, {
    graphContext,
    selectedElements,
    selectedContext: selectedElements?.length
      ? buildSelectedContext(selectedElements, graphContext)
      : undefined,
  });

  const stream = createUIMessageStream({
    execute: ({ writer }) => {
      const result = streamText({
        model: modelStrategies.languageModel(DEFAULT_MODEL),
        messages: convertToModelMessages([
          ...existingMessages.map(toChatMessage),
          {
            id: msg.id,
            role: msg.role ?? "user",
            parts: msg.parts,
          },
        ]),
        system: systemPrompt,
        tools: { generateDiagram: generateDiagramTool },
        stopWhen: stepCountIs(2),
        abortSignal: signal,
        experimental_transform: smoothStream({
          chunking: "word",
          delayInMs: 15,
        }),
        onError: (error) => {
          console.error("[copilot]", error);
        },
      });

      void result.consumeStream();

      writer.merge(result.toUIMessageStream());
    },
    onFinish: async ({ responseMessage }) => {
      await createMessage({
        id: responseMessage.id,
        chatId,
        role: Role.ASSISTANT,
      });

      await createParts(
        responseMessage.parts.map(({ type, ...details }, order) => ({
          type,
          details,
          messageId: responseMessage.id,
          order,
        })),
      );
    },
  });

  return createUIMessageStreamResponse({
    stream,
    headers: {
      "Content-Type": "application/octet-stream",
      "Content-Encoding": "none",
    },
  });
};
