import { TRPCError } from "@trpc/server";
import { z } from "zod";
import { COOKIE_NAME } from "@shared/const";
import { getSessionCookieOptions } from "./_core/cookies";
import { invokeLLM, type FileContent, type ImageContent, type Message, type TextContent } from "./_core/llm";
import { transcribeAudio } from "./_core/voiceTranscription";
import { systemRouter } from "./_core/systemRouter";
import { protectedProcedure, publicProcedure, router } from "./_core/trpc";
import {
  addAttachment,
  addMessage,
  createConversation,
  getConversation,
  listConversations,
  listMessages,
  renameConversation,
} from "./db";
import { storagePut } from "./storage";

const attachmentInput = z.object({
  fileName: z.string().min(1).max(255),
  mimeType: z.string().min(1).max(160),
  fileSize: z.number().int().positive().max(8 * 1024 * 1024),
  data: z.string().min(1),
});

const conversationInput = z.object({ conversationId: z.number().int().positive() });

function requestOrigin(req: { protocol: string; get: (name: string) => string | undefined }) {
  const host = req.get("host");
  return `${req.protocol}://${host}`;
}

function safeFileName(fileName: string) {
  return fileName.replace(/[^a-zA-Z0-9._-]+/g, "-").slice(0, 120) || "attachment";
}

function getTextContent(content: string | Array<TextContent | ImageContent | FileContent>) {
  if (typeof content === "string") return content;
  return content
    .filter((part): part is TextContent => part.type === "text")
    .map((part) => part.text)
    .join("\n")
    .trim();
}

export const appRouter = router({
  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;
    }),
  }),

  conversations: router({
    list: protectedProcedure.query(({ ctx }) => listConversations(ctx.user.id)),
    get: protectedProcedure.input(conversationInput).query(({ ctx, input }) =>
      listMessages(ctx.user.id, input.conversationId),
    ),
    create: protectedProcedure
      .input(z.object({ title: z.string().trim().min(1).max(160).optional() }).optional())
      .mutation(({ ctx, input }) => createConversation(ctx.user.id, input?.title)),
    rename: protectedProcedure
      .input(conversationInput.extend({ title: z.string().trim().min(1).max(160) }))
      .mutation(({ ctx, input }) => renameConversation(ctx.user.id, input.conversationId, input.title)),
  }),

  chat: router({
    send: protectedProcedure
      .input(
        z.object({
          conversationId: z.number().int().positive(),
          content: z.string().trim().min(1).max(12000),
          attachments: z.array(attachmentInput).max(5).default([]),
        }),
      )
      .mutation(async ({ ctx, input }) => {
        const conversation = await getConversation(ctx.user.id, input.conversationId);
        if (!conversation) {
          throw new TRPCError({ code: "NOT_FOUND", message: "Conversation not found." });
        }

        const uploaded = [] as Array<{
          fileName: string;
          mimeType: string;
          fileSize: number;
          storageKey: string;
          storageUrl: string;
        }>;

        for (const attachment of input.attachments) {
          const bytes = Buffer.from(attachment.data, "base64");
          if (bytes.length !== attachment.fileSize) {
            throw new TRPCError({ code: "BAD_REQUEST", message: `Attachment size mismatch for ${attachment.fileName}.` });
          }
          const key = `chat/${ctx.user.id}/${Date.now()}-${safeFileName(attachment.fileName)}`;
          const stored = await storagePut(key, bytes, attachment.mimeType);
          uploaded.push({
            fileName: attachment.fileName,
            mimeType: attachment.mimeType,
            fileSize: attachment.fileSize,
            storageKey: stored.key,
            storageUrl: stored.url,
          });
        }

        const userMessageId = await addMessage(ctx.user.id, input.conversationId, "user", input.content);
        for (const attachment of uploaded) {
          await addAttachment({
            userId: ctx.user.id,
            conversationId: input.conversationId,
            messageId: userMessageId,
            ...attachment,
          });
        }

        const persistedMessages = await listMessages(ctx.user.id, input.conversationId);
        const origin = requestOrigin(ctx.req);
        const llmMessages: Message[] = [
          {
            role: "system",
            content:
              "You are Luma, a calm, thoughtful AI workspace assistant. Be concise but warm. Use Markdown when it improves clarity. Never claim to have accessed a file unless its contents are actually available to you.",
          },
          ...persistedMessages.map((message) => {
            if (message.role === "assistant") {
              return { role: "assistant", content: message.content } as Message;
            }

            const textPart: TextContent = {
              type: "text",
              text: message.content,
            };
            const parts: Array<TextContent | ImageContent | FileContent> = [textPart];
            for (const attachment of message.attachments) {
              const absoluteUrl = `${origin}${attachment.storageUrl}`;
              if (attachment.mimeType.startsWith("image/")) {
                parts.push({ type: "image_url", image_url: { url: absoluteUrl, detail: "auto" } });
              } else if (["application/pdf", "audio/mpeg", "audio/wav", "audio/mp4", "video/mp4"].includes(attachment.mimeType)) {
                parts.push({
                  type: "file_url",
                  file_url: {
                    url: absoluteUrl,
                    mime_type: attachment.mimeType as FileContent["file_url"]["mime_type"],
                  },
                });
              } else {
                textPart.text += `\n[Attached file: ${attachment.fileName} (${attachment.mimeType})]`;
              }
            }
            return { role: "user", content: parts } as Message;
          }),
        ];

        try {
          const response = await invokeLLM({ messages: llmMessages, maxTokens: 1400 });
          const rawContent = response.choices[0]?.message?.content;
          const assistantContent = rawContent ? getTextContent(rawContent) : "I’m here with you. Could you try sending that once more?";
          const assistantMessageId = await addMessage(ctx.user.id, input.conversationId, "assistant", assistantContent);
          return { conversationId: input.conversationId, userMessageId, assistantMessageId, content: assistantContent };
        } catch (error) {
          console.error("[Chat] LLM request failed:", error);
          throw new TRPCError({ code: "INTERNAL_SERVER_ERROR", message: "The assistant could not respond right now. Please try again." });
        }
      }),
  }),

  voice: router({
    transcribe: protectedProcedure
      .input(
        z.object({
          data: z.string().min(1),
          mimeType: z.string().min(1).max(120).default("audio/webm"),
          fileSize: z.number().int().positive().max(16 * 1024 * 1024),
          language: z.string().length(2).optional(),
        }),
      )
      .mutation(async ({ ctx, input }) => {
        const bytes = Buffer.from(input.data, "base64");
        if (bytes.length !== input.fileSize) {
          throw new TRPCError({ code: "BAD_REQUEST", message: "The audio recording size could not be verified." });
        }
        const extension = input.mimeType.split("/")[1]?.split(";")[0] || "webm";
        const stored = await storagePut(`voice/${ctx.user.id}/${Date.now()}.${extension}`, bytes, input.mimeType);
        const result = await transcribeAudio({
          audioUrl: `${requestOrigin(ctx.req)}${stored.url}`,
          language: input.language,
          prompt: "Transcribe the user's voice to text. Preserve the user's wording and punctuation.",
        });
        if ("error" in result) {
          throw new TRPCError({ code: "BAD_REQUEST", message: result.error });
        }
        return { text: result.text };
      }),
  }),
});

export type AppRouter = typeof appRouter;