import { query, mutation } from "./_generated/server";
import { v } from "convex/values";

// Add the missing getChats function
export const getChats = query({
  args: {
    userId: v.string(),
  },
  returns: v.array(
    v.object({
      _id: v.id("Chat"),
      _creationTime: v.number(),
      agentId: v.optional(v.string()),
      aiAvatarUrl: v.optional(v.string()),
      chatId: v.string(),
      chatName: v.optional(v.string()),
      createdAt: v.float64(),
      messages: v.array(
        v.object({
          content: v.string(),
          images: v.optional(v.array(v.string())),
          role: v.string(),
          timestamp: v.optional(v.float64()),
          videos: v.optional(v.array(v.string())),
        })
      ),
      model: v.optional(v.string()),
      updatedAt: v.float64(),
      userId: v.string(),
    })
  ),
  handler: async (ctx, args) => {
    return await ctx.db
      .query("Chat")
      .withIndex("by_userId", (q) => q.eq("userId", args.userId))
      .collect();
  },
});

export const getChatById = query({
  args: {
    chatId: v.string(), // Changed from 'id' to 'chatId' to match schema
  },
  returns: v.union(
    v.object({
      _id: v.id("Chat"),
      _creationTime: v.number(),
      agentId: v.optional(v.string()),
      aiAvatarUrl: v.optional(v.string()),
      chatId: v.string(),
      chatName: v.optional(v.string()),
      createdAt: v.float64(),
      messages: v.array(
        v.object({
          content: v.string(),
          images: v.optional(v.array(v.string())),
          role: v.string(),
          timestamp: v.optional(v.float64()),
          videos: v.optional(v.array(v.string())),
        })
      ),
      model: v.optional(v.string()),
      updatedAt: v.float64(),
      userId: v.string(),
    }),
    v.null()
  ),
  handler: async (ctx, args) => {
    return await ctx.db
      .query("Chat")
      .withIndex("by_chatId", (q) => q.eq("chatId", args.chatId))
      .first();
  },
});

export const getChatByUserId = query({
  args: {
    userId: v.string(),
  },
  returns: v.array(
    v.object({
      _id: v.id("Chat"),
      _creationTime: v.number(),
      agentId: v.optional(v.string()),
      aiAvatarUrl: v.optional(v.string()),
      chatId: v.string(),
      chatName: v.optional(v.string()),
      createdAt: v.float64(),
      messages: v.array(
        v.object({
          content: v.string(),
          images: v.optional(v.array(v.string())),
          role: v.string(),
          timestamp: v.optional(v.float64()),
          videos: v.optional(v.array(v.string())),
        })
      ),
      model: v.optional(v.string()),
      updatedAt: v.float64(),
      userId: v.string(),
    })
  ),
  handler: async (ctx, args) => {
    return await ctx.db
      .query("Chat")
      .withIndex("by_userId", (q) => q.eq("userId", args.userId))
      .collect();
  },
});

export const getChatByChatId = query({
  args: {
    chatId: v.string(),
  },
  returns: v.union(
    v.object({
      _id: v.id("Chat"),
      _creationTime: v.number(),
      agentId: v.optional(v.string()),
      aiAvatarUrl: v.optional(v.string()),
      chatId: v.string(),
      chatName: v.optional(v.string()),
      createdAt: v.float64(),
      messages: v.array(
        v.object({
          content: v.string(),
          images: v.optional(v.array(v.string())),
          role: v.string(),
          timestamp: v.optional(v.float64()),
          videos: v.optional(v.array(v.string())),
        })
      ),
      model: v.optional(v.string()),
      updatedAt: v.float64(),
      userId: v.string(),
    }),
    v.null()
  ),
  handler: async (ctx, args) => {
    return await ctx.db
      .query("Chat")
      .withIndex("by_chatId", (q) => q.eq("chatId", args.chatId))
      .first();
  },
});

export const addMessages = mutation({
  args: {
    chatId: v.string(),
    messages: v.array(
      v.object({
        content: v.string(),
        images: v.optional(v.array(v.string())),
        role: v.string(),
        timestamp: v.optional(v.float64()),
        videos: v.optional(v.array(v.string())),
      })
    ),
  },
  returns: v.union(
    v.object({
      _id: v.id("Chat"),
      _creationTime: v.number(),
      agentId: v.optional(v.string()),
      aiAvatarUrl: v.optional(v.string()),
      chatId: v.string(),
      chatName: v.optional(v.string()),
      createdAt: v.float64(),
      messages: v.array(
        v.object({
          content: v.string(),
          images: v.optional(v.array(v.string())),
          role: v.string(),
          timestamp: v.optional(v.float64()),
          videos: v.optional(v.array(v.string())),
        })
      ),
      model: v.optional(v.string()),
      updatedAt: v.float64(),
      userId: v.string(),
    }),
    v.null()
  ),
  handler: async (ctx, args) => {
    const chat = await ctx.db
      .query("Chat")
      .withIndex("by_chatId", (q) => q.eq("chatId", args.chatId))
      .first();

    if (!chat) {
      throw new Error(`Chat with chatId ${args.chatId} not found`);
    }

    const updatedMessages = [...chat.messages, ...args.messages];

    await ctx.db.patch(chat._id, {
      messages: updatedMessages,
      updatedAt: Date.now(),
    });

    return await ctx.db.get(chat._id);
  },
});

export const createChat = mutation({
  args: {
    chatId: v.string(),
    chatName: v.optional(v.string()), // Made optional to match schema
    userId: v.string(),
    agentId: v.optional(v.string()),
    aiAvatarUrl: v.optional(v.string()),
    model: v.optional(v.string()),
    initialMessage: v.optional(
      v.object({
        content: v.string(),
        images: v.optional(v.array(v.string())),
        role: v.string(),
        timestamp: v.optional(v.float64()),
        videos: v.optional(v.array(v.string())),
      })
    ),
  },
  returns: v.id("Chat"),
  handler: async (ctx, args) => {
    const messages = args.initialMessage ? [args.initialMessage] : [];
    const now = Date.now();

    return await ctx.db.insert("Chat", {
      agentId: args.agentId,
      aiAvatarUrl: args.aiAvatarUrl,
      chatId: args.chatId,
      chatName: args.chatName,
      createdAt: now,
      messages: messages,
      model: args.model,
      updatedAt: now,
      userId: args.userId,
    });
  },
});