routers.ts

import { COOKIE_NAME } from "@shared/const";
import { getSessionCookieOptions } from "./_core/cookies";
import { systemRouter } from "./_core/systemRouter";
import { publicProcedure, protectedProcedure, router } from "./_core/trpc";
import { z } from "zod";
import { generateResponseWithReasoning, generateImage } from "./llm";
import { searchOnline, formatSearchResults, sanitizeSearchQuery } from "./search";
import { checkRateLimit } from "./rateLimit";

export const appRouter = router({
  // if you need to use socket.io, read and register route in server/_core/index.ts, all api should start with '/api/' so that the gateway can route correctly
  system: systemRouter,
  auth: router({
    me: publicProcedure.query(opts => opts.ctx.user),
    logout: publicProcedure.mutation(({ ctx }) => {
      const cookieOptions = getSessionCookieOptions(ctx.req);
      ctx.res.clearCookie(COOKIE_NAME, { ...cookieOptions, maxAge: -1 });
      return {
        success: true,
      } as const;
    }),
  }),

  // Section 1: Chat procedures (Ask mode)
  chat: router({
    send: protectedProcedure
      .input(
        z.object({
          prompt: z.string().min(1),
          conversationId: z.number().optional(),
          enableSearch: z.boolean().default(false),
          enableThinking: z.boolean().default(false),
          history: z
            .array(
              z.object({
                role: z.enum(["user", "assistant"]),
                content: z.string(),
              })
            )
            .default([]),
        })
      )
      .mutation(async ({ ctx, input }) => {
        // Check rate limit
        const userId = ctx.user?.id?.toString() || "anonymous";
        const rateLimitCheck = checkRateLimit(userId, "chat");

        if (!rateLimitCheck.allowed) {
          throw new Error(
            `Rate limit exceeded. Try again in ${Math.ceil(rateLimitCheck.resetIn / 1000)} seconds.`
          );
        }

        try {
          let searchResults = "";

          // Search online if enabled
          if (input.enableSearch) {
            const sanitizedQuery = sanitizeSearchQuery(input.prompt);
            const results = await searchOnline(sanitizedQuery, 5);
            searchResults = formatSearchResults(results);
          }

          // Generate response with optional reasoning
          const response = await generateResponseWithReasoning(
            input.prompt,
            searchResults,
            input.enableThinking,
            input.history
          );

          return {
            success: true,
            response: response.response,
            reasoning: response.reasoning,
            model: response.model,
            tokensUsed: response.tokensUsed,
            searchResults: searchResults || undefined,
          };
        } catch (error) {
          console.error("Chat error:", error);
          throw new Error(
            error instanceof Error ? error.message : "Failed to generate response"
          );
        }
      }),
  }),

  // Section 1: Image generation procedures (Imagine mode)
  imagine: router({
    generate: protectedProcedure
      .input(
        z.object({
          prompt: z.string().min(1),
        })
      )
      .mutation(async ({ ctx, input }) => {
        // Check rate limit
        const userId = ctx.user?.id?.toString() || "anonymous";
        const rateLimitCheck = checkRateLimit(userId, "imagine");

        if (!rateLimitCheck.allowed) {
          throw new Error(
            `Rate limit exceeded. Try again in ${Math.ceil(rateLimitCheck.resetIn / 1000)} seconds.`
          );
        }

        try {
          const imageUrl = await generateImage(input.prompt);
          return {
            success: true,
            imageUrl,
            prompt: input.prompt,
          };
        } catch (error) {
          console.error("Image generation error:", error);
          throw new Error(
            error instanceof Error ? error.message : "Failed to generate image"
          );
        }
      }),
  }),

  // Section 1: Search procedure
  search: router({
    online: publicProcedure
      .input(
        z.object({
          query: z.string().min(1),
          maxResults: z.number().default(5),
        })
      )
      .query(async ({ input }) => {
        try {
          const sanitizedQuery = sanitizeSearchQuery(input.query);
          const results = await searchOnline(sanitizedQuery, input.maxResults);
          return {
            success: true,
            results,
            query: input.query,
          };
        } catch (error) {
          console.error("Search error:", error);
          throw new Error(
            error instanceof Error ? error.message : "Search failed"
          );
        }
      }),
  }),
});

export type AppRouter = typeof appRouter;