import { router, publicProcedure, createCallerFactory } from "./server.js";
import type { DraftMessage } from "../../types.js";
import { MODEL_NAME, openrouter } from "../provider.js";
import { generateObject, generateText, jsonSchema } from "ai";

const factsFromNewMessagesSystemPrompt = ({
  previousRunningSummary,
  messagesSincePreviousRunningSummary,
}: {
  previousRunningSummary: string;
  messagesSincePreviousRunningSummary: Array<DraftMessage>;
}) => `You are an expert at extracting facts from conversations.

An AI assistant is in the middle of a conversation whose data is given below. The data consists of a summary of a conversation, and optionally some messages exchanged since that summary was produced. The user will provide you with *new* messages.

Your task is to extract *new* facts that can be gleaned from the *new* messages that the user sends.

* You should not extract any facts that are already in the summary.
* The user should be referred to as "the user" in the fact text.
* The user's pronouns should be either he or she, NOT "they" or "them", because this summary will be read by an AI assistant to give it context; and excessive use of "they" or "them" will make what they refer to unclear or ambiguous.
* The assistant should be referred to as "I" or "me", because these facts will be read by an AI assistant to give it context.

<running_summary>
  ${previousRunningSummary}
</running_summary>

${messagesSincePreviousRunningSummary.map(
  (message) =>
    `<${message.role}_message>${message.parts
      .filter((p) => p.type === "text")
      .map((p) => p.text)
      .join("\n")}</${message.role}_message>`
)}
`;

const factsFromNewMessagesUserPrompt = ({
  newMessages,
}: {
  newMessages: Array<DraftMessage>;
}) =>
  `${newMessages.map(
    (message) =>
      `<${message.role}_message>${message.parts
        .filter((p) => p.type === "text")
        .map((p) => p.text)
        .join("\n")}</${message.role}_message>`
  )}

Extract new facts from these messages.`;

export const facts = router({
  fetchByConversationId: publicProcedure
    .input((x) => x as { conversationId: string })
    .query(async ({ input: { conversationId }, ctx: { dbClient } }) => {
      const rows = await dbClient
        .selectFrom("facts")
        .innerJoin("messages", "messages.id", "facts.sourceMessageId")
        .selectAll("facts")
        .where("conversationId", "=", conversationId)
        .execute();
      return rows;
    }),
  deleteOne: publicProcedure
    .input(
      (x) =>
        x as {
          factId: string;
        }
    )
    .mutation(async ({ input: { factId }, ctx: { dbClient } }) => {
      await dbClient.deleteFrom("facts").where("id", "=", factId).execute();
      return { ok: true };
    }),
  update: publicProcedure
    .input(
      (x) =>
        x as {
          factId: string;
          content: string;
        }
    )
    .mutation(async ({ input: { factId, content }, ctx: { dbClient } }) => {
      await dbClient
        .updateTable("facts")
        .set({ content })
        .where("id", "=", factId)
        .execute();
      return { ok: true };
    }),
  extractFromNewMessages: publicProcedure
    .input(
      (x) =>
        x as {
          previousRunningSummary: string;
          /** will *not* have facts extracted */
          messagesSincePreviousRunningSummary: Array<DraftMessage>;
          /** *will* have facts extracted */
          newMessages: Array<DraftMessage>;
        }
    )
    .query(
      async ({
        input: {
          previousRunningSummary,
          messagesSincePreviousRunningSummary,
          newMessages,
        },
        ctx: { openrouter },
      }) => {
        const factsFromUserMessageResponse = await generateObject({
          model: openrouter(MODEL_NAME),
          messages: [
            {
              role: "system" as const,
              content: factsFromNewMessagesSystemPrompt({
                previousRunningSummary,
                messagesSincePreviousRunningSummary,
              }),
            },
            {
              role: "user" as const,
              content: factsFromNewMessagesUserPrompt({
                newMessages,
              }),
            },
          ],
          schema: jsonSchema<{
            facts: Array<string>;
          }>({
            type: "object",
            properties: {
              facts: {
                type: "array",
                items: {
                  type: "string",
                },
              },
            },
            required: ["facts"],
          }),
          temperature: 0.4,
          maxRetries: 0,
        }).catch((err) => {
          console.error(err);
          return {
            object: {
              facts: [] as Array<string>,
            },
          };
        });
        return factsFromUserMessageResponse;
      }
    ),
});

export const createCaller = createCallerFactory(facts);