import {
  router,
  publicProcedure,
  createCallerFactory,
  Validator,
  // Validator
} from "../../trpc/server";
import { createOpenRouter } from "@openrouter/ai-sdk-provider";
import { generateText } from "ai";
import type { Message as UIMessage } from "ai";
import type { OtherParameters } from "../../types.js";
import { env } from "../../server/env.js";
// import { client } from "../../database/milvus";
// import {
//   ConsistencyLevelEnum,
//   type NumberArrayId,
// } from "@zilliz/milvus2-sdk-node";
import { db } from "../../database/postgres";
import type { ConversationsId } from "../../database/generated/public/Conversations";
import type { UsersId } from "../../database/generated/public/Users";

const openrouter = createOpenRouter({
  apiKey: env.OPENROUTER_API_KEY,
});

export const chat = router({
  listConversations: publicProcedure.query(async () => {
    const rows = await db.selectFrom("conversations").selectAll().execute();
    return rows;
  }),
  fetchConversation: publicProcedure
    .input((x) => x as { id: number })
    .query(async ({ input: { id } }) => {
      const row = await db
        .selectFrom("conversations")
        .selectAll()
        .where("id", "=", id as ConversationsId)
        .executeTakeFirst();
      return row;
    }),
  createConversation: publicProcedure.mutation(async () => {
    const title = "New Conversation";
    const row = await db
      .insertInto("conversations")
      .values({
        title,
        user_id: 1 as UsersId,
      })
      .returningAll()
      .executeTakeFirst();
    return row;
  }),
  deleteConversation: publicProcedure
    .input((x) => x as { id: number })
    .mutation(async ({ input: { id } }) => {
      const result = await db
        .deleteFrom("conversations")
        .where("id", "=", id as ConversationsId)
        .execute();
      return result;
    }),
  updateConversationTitle: publicProcedure
    .input(
      (x) =>
        x as {
          id: number;
          title: string;
        },
    )
    .mutation(async ({ input: { id, title } }) => {
      const result = await db
        .updateTable("conversations")
        .set({ title })
        .where("id", "=", id as ConversationsId)
        .execute();
      return result[0];
    }),
  sendMessage: publicProcedure
    .input(
      (x) =>
        x as {
          messages: Array<UIMessage>;
          systemPrompt: string;
          parameters: OtherParameters;
        },
    )
    .query(async ({ input: { messages, systemPrompt, parameters } }) => {
      const response = await generateText({
        model: openrouter("mistralai/mistral-nemo"),
        messages: [
          { role: "system" as const, content: systemPrompt },
          ...messages,
        ],
        maxSteps: 3,
        tools: undefined,
        ...parameters,
      });
      return response;
    }),
});

export const createCaller = createCallerFactory(chat);