actions.ts

'use server'

import { revalidatePath } from 'next/cache'
import { redirect } from 'next/navigation'
import prisma from '@/lib/prisma'

import { auth } from '@/auth'
import { ServerActionResult, type ChatMetadata, type Chat } from '@/lib/types'
import { type ChatHistory } from '@/lib/store'
import { TokenUsage } from 'ai'
import { AIMODELS } from '@/lib/constants'

import { cache } from 'react'
import * as chatCache from '@/lib/chat/chat-cache'

export async function getChats(teamId?: string) {
  const session = await auth()
  if (!session?.user?.id) {
    return []
  }

  if (!teamId) {
    const user = await prisma.user.findUnique({
      where: {
        id: session.user.id
      }
    })
    if (!user) {
      return []
    }
    teamId = user.personalTeamId || ''
  }

  if (!teamId) {
    return []
  }
  console.log('Getting chats with team id', teamId)
  try {
    const folders = await prisma.folder.findMany({
      where: {
        teamId
      }
    })
    const chats = await prisma.chat.findMany({
      where: {
        teamId
      },
      orderBy: {
        createdAt: 'desc'
      }
    })

    return {
      folders,
      chats
    } as unknown as ChatHistory
  } catch (error) {
    console.error('Error fetching chats:', error)
    return []
  }
}

export const getChat = cache(async (id: string) => {
  const session = await auth()
  if (!session?.user?.id) {
    return null
  }
  console.log(`Getting chat ${id}`)
  // first, check in the cache
  const cachedChat = chatCache.get(id)
  if (cachedChat) {
    console.log(`Chat ${id} found in cache`)
    return cachedChat
  }
  // if not found, fetch via db
  const chat = await prisma.chat.findUnique({
    where: {
      id: id,
      userId: session.user.id
    },
    include: {
      messages: true
    }
  }) as unknown as Chat

  if (!chat) {
    console.log(`Chat ${id} not found`)
    return null
  }
  // also, save to cache
  chatCache.set(id, chat)

  console.log(`Chat ${id} found with ${chat.messages.length} messages`)

  return chat
})

export async function removeChat({ id, path }: { id: string; path: string }) {
  const session = await auth()

  if (!session?.user?.id) {
    return {
      error: 'Unauthorized'
    }
  }

  console.log('Removing chat', id)

  // Remove from cache
  chatCache.remove(id)

  const chat = await prisma.chat.findUnique({
    where: {
      id: id,
      userId: session.user.id
    }
  })

  if (!chat) {
    return {
      error: 'Chat not found or unauthorized'
    }
  }

  // Delete messages first
  await prisma.message.deleteMany({
    where: {
      chatId: id,
      chat: {
        userId: session.user.id
      }
    }
  })

  // Delete chat
  await prisma.chat.delete({
    where: {
      id: id,
      userId: session.user.id
    }
  })

  revalidatePath('/')
  return revalidatePath(path)
}

export async function clearChats(teamId?: string) {
  const session = await auth()

  if (!session?.user?.id) {
    return {
      error: 'Unauthorized'
    }
  }

  if (!teamId) {
    const user = await prisma.user.findUnique({
      where: {
        id: session.user.id
      }
    })
    if (!user) {
      return {
        error: 'User not found'
      }
    }
    teamId = user.personalTeamId || ''
  }

  if (!teamId) {
    return {
      error: 'Team not found'
    }
  }

  // Get a list of all chatIds
  const chatIds = await prisma.chat.findMany({
    where: {
      teamId
    }
  })

  // Delete all messages for the chatIds first
  await prisma.message.deleteMany({
    where: {
      chatId: {
        in: chatIds.map(chat => chat.id)
      }
    }
  })

  // Then delete all chats for the team from the database
  await prisma.chat.deleteMany({
    where: {
      teamId
    }
  })

  // Clear the chat cache for the user
  chatCache.removeAll(chatIds.map(chat => chat.id))

  revalidatePath('/')
  return redirect('/')
}

export async function getSharedChat(id: string) {
  const chat = await prisma.chat.findUnique({
    where: {
      id: id,
      sharePath: { not: null }
    }
  })

  if (!chat) {
    return null
  }

  return chat as unknown as Chat
}

export async function shareChat(id: string): Promise<ServerActionResult<ChatMetadata>> {
  const session = await auth()

  if (!session?.user?.id) {
    return {
      error: 'Unauthorized'
    }
  }

  const chat = await prisma.chat.findUnique({
    where: {
      id: id,
      userId: session.user.id
    }
  })

  if (!chat) {
    return {
      error: 'Something went wrong'
    }
  }

  const updatedChat = await prisma.chat.update({
    where: { id: chat.id },
    data: {
      sharePath: `/share/${chat.id}`
    }
  })

  return { ...updatedChat, messages: [] } as ChatMetadata
}

export async function saveChat(chat: Chat, teamId?: string) {
  const session = await auth()

  if (session && session.user) {
    console.log(`Saving chat ${chat.id} with ${chat.messages.length} messages`)
    // first, set in cache
    chatCache.set(chat.id, chat)
    // determine the team
    if (!teamId) {
      const user = await prisma.user.findUnique({
        where: {
          id: session.user.id
        }
      })
      teamId = user?.personalTeamId || ''
    }
    if (!teamId) {
      return {
        error: 'Team not found'
      }
    }
    // then, set in database
    try {
      console.log('Creating chat')
      await prisma.chat.upsert({
        where: { id: chat.id },
        update: {},
        create: {
          id: chat.id,
          title: chat.title,
          userId: chat.userId,
          sharePath: chat.sharePath,
          path: chat.path,
          teamId: teamId
        }
      });
    } catch (error) {
      console.error('Error creating chat')
    }
    console.log(`Chat saved ${chat.id} and ${chat.messages.length} messages`)
    // Create or update messages, saving the newest first
    if (chat.messages && chat.messages.length > 0) {
      const newestMessage = chat.messages[chat.messages.length - 1];
      await prisma.message.upsert({
        where: { id: newestMessage.id },
        update: {
          content: newestMessage.content as string,
          role: newestMessage.role,
        },
        create: {
          id: newestMessage.id,
          content: newestMessage.content as string,
          role: newestMessage.role,
          chatId: chat.id,
          modelKey: newestMessage.modelKey,
        }
      });
      console.log('Message created ' + newestMessage.id);

      // then, check and update older messages if needed
      for (let i = chat.messages.length - 2; i >= 0; i--) {
        const message = chat.messages[i];
        const existingMessage = await prisma.message.findUnique({
          where: { id: message.id }
        });

        if (!existingMessage || existingMessage.content !== message.content || existingMessage.role !== message.role) {
          await prisma.message.upsert({
            where: { id: message.id },
            update: {
              content: message.content as string,
              role: message.role,
            },
            create: {
              id: message.id,
              content: message.content as string,
              role: message.role,
              chatId: chat.id,
            }
          });
          console.log('Message updated ' + message.id);
        }
      }
    }
  } else {
    return
  }
}

export async function isWithinUsageBudget(userId: string) {
  const DAILY_COST_LIMIT_IN_NANODOLLARS = {
    'production': 20000000, // 2 cents
    'development': 1000000000, // 1 dollar
    'test': 0
  }
  const today = new Date()
  today.setHours(0, 0, 0, 0)
  const usage = await prisma.userDailyUsage.findUnique({
    where: {
      userId_date: {
        userId: userId,
        date: today
      }
    }
  })
  if (!usage) {
    return true
  }
  return usage.totalCostInNanoDollars <
    (DAILY_COST_LIMIT_IN_NANODOLLARS[process.env.NODE_ENV] || DAILY_COST_LIMIT_IN_NANODOLLARS.development)
}

/**
 * Given a `chatId`, `messageId`, and `usage`, update the message with the usage
 */
export async function saveMessageWithUsage(
  chatId: string,
  messageId: string,
  {
    modelKey,
    usage
  }: {
    modelKey: string
    usage: TokenUsage
  }
) {
  console.log('Saving message with usage', messageId, chatId, modelKey)
  const session = await auth()

  if (session && session.user) {
    const model = AIMODELS[modelKey]
    const { promptTokens, completionTokens, totalTokens } = usage;
    const today = new Date();
    today.setHours(0, 0, 0, 0);
    const cost = promptTokens * model.inputCostPerTokenInNanodollars + completionTokens * model.outputCostPerTokenInNanodollars

    // Update daily usage
    await prisma.userDailyUsage.upsert({
      where: {
        userId_date: {
          userId: session.user.id,
          date: today
        }
      },
      update: {
        promptTokensUsage: { increment: promptTokens },
        completionTokensUsage: { increment: completionTokens },
        totalTokensUsage: { increment: totalTokens },
        totalCostInNanoDollars: { increment: cost }
      },
      create: {
        userId: session.user.id,
        date: today,
        promptTokensUsage: promptTokens,
        completionTokensUsage: completionTokens,
        totalTokensUsage: totalTokens,
        totalCostInNanoDollars: cost
      }
    });

    // Update message with exact usage
    console.log('Updating message with usage', messageId, chatId)
    const message = {
      promptTokensUsage: promptTokens,
      completionTokensUsage: completionTokens,
      totalTokensUsage: totalTokens,
      totalCostInNanoDollars: cost,
      role: 'assistant',
      modelKey: modelKey
    }
    await prisma.message.upsert({
      where: {
        id: messageId,
        chatId: chatId
      },
      update: message,
      create: {
        id: messageId,
        chatId: chatId,
        ...message
      }
    });
  }
}

export async function refreshHistory(path: string) {
  redirect(path)
}

export async function getMissingKeys() {
  const keysRequired = ['OPENAI_API_KEY']
  return keysRequired
    .map(key => (process.env[key] ? '' : key))
    .filter(key => key !== '')
}

// Team actions
export async function createTeam(name: string, icon: string | null) {
  const session = await auth()
  if (!session?.user?.id) {
    return { error: 'Unauthorized' }
  }

  try {
    const team = await prisma.team.create({
      data: {
        name,
        icon,
        ownerId: session.user.id,
        members: {
          create: {
            userId: session.user.id,
            role: 'OWNER'
          }
        }
      }
    })
    return team
  } catch (error) {
    console.error('Error creating team:', error)
    return { error: 'Failed to create team' }
  }
}

export async function getTeam(id: string) {
  const session = await auth()
  if (!session?.user?.id) {
    return { error: 'Unauthorized' }
  }

  try {
    const team = await prisma.team.findUnique({
      where: { id },
      include: {
        members: true,
        chats: true,
        folders: true
      }
    })
    if (!team) {
      return { error: 'Team not found' }
    }
    const isMember = team.members.some(member => member.userId === session?.user?.id)
    if (!isMember) {
      return { error: 'Unauthorized' }
    }
    return team
  } catch (error) {
    console.error('Error getting team:', error)
    return { error: 'Failed to get team' }
  }
}

export async function updateTeam(id: string, data: { name?: string, icon?: string | null }) {
  const session = await auth()
  if (!session?.user?.id) {
    return { error: 'Unauthorized' }
  }

  try {
    const team = await prisma.team.findUnique({
      where: { id },
      include: { members: true }
    })
    if (!team) {
      return { error: 'Team not found' }
    }
    const userRole = team.members.find(member => member.userId === session?.user?.id)?.role
    if (userRole !== 'OWNER' && userRole !== 'ADMINISTRATOR') {
      return { error: 'Unauthorized' }
    }
    const updatedTeam = await prisma.team.update({
      where: { id },
      data
    })
    return updatedTeam
  } catch (error) {
    console.error('Error updating team:', error)
    return { error: 'Failed to update team' }
  }
}

export async function deleteTeam(id: string) {
  const session = await auth()
  if (!session?.user?.id) {
    return { error: 'Unauthorized' }
  }

  try {
    const team = await prisma.team.findUnique({
      where: { id },
      include: { members: true }
    })
    if (!team) {
      return { error: 'Team not found' }
    }
    const userRole = team.members.find(member => member.userId === session?.user?.id)?.role
    if (userRole !== 'OWNER') {
      return { error: 'Unauthorized' }
    }
    await prisma.team.delete({ where: { id } })
    return { success: true }
  } catch (error) {
    console.error('Error deleting team:', error)
    return { error: 'Failed to delete team' }
  }
}

// Folder actions
export async function createFolder(id: string, name: string, icon?: string, teamId?: string) {
  const session = await auth()
  if (!session?.user?.id) {
    return { error: 'Unauthorized' }
  }

  try {
    if (!teamId) {
      const user = await prisma.user.findUnique({
        where: {
          id: session.user.id
        }
      })
      if (!user) {
        return {
          error: 'User not found'
        }
      }
      teamId = user.personalTeamId || ''
    }
    if (!teamId) {
      return {
        error: 'Team not found'
      }
    }
    const team = await prisma.team.findUnique({
      where: { id: teamId },
      include: { members: true }
    })
    console.log('Found team', team)
    if (!team) {
      return { error: 'Team not found' }
    }
    const isMember = team.ownerId === session?.user?.id
      || team.members.some(member => member.userId === session?.user?.id)
    if (!isMember) {
      return { error: 'Unauthorized' }
    }
    const folder = await prisma.folder.create({
      data: {
        id,
        name,
        icon,
        teamId
      }
    })
    return folder
  } catch (error) {
    console.error('Error creating folder:', error)
    return { error: 'Failed to create folder' }
  }
}

export async function getFolder(id: string) {
  const session = await auth()
  if (!session?.user?.id) {
    return { error: 'Unauthorized' }
  }

  try {
    const folder = await prisma.folder.findUnique({
      where: { id },
      include: {
        team: {
          include: { members: true }
        },
        chats: true
      }
    })
    if (!folder) {
      return { error: 'Folder not found' }
    }
    const isMember = folder.team.members.some(member => member.userId === session?.user?.id)
    if (!isMember) {
      return { error: 'Unauthorized' }
    }
    return folder
  } catch (error) {
    console.error('Error getting folder:', error)
    return { error: 'Failed to get folder' }
  }
}

export async function updateFolder(id: string, data: { name?: string, icon?: string | null }) {
  const session = await auth()
  if (!session?.user?.id) {
    return { error: 'Unauthorized' }
  }

  try {
    const folder = await prisma.folder.findUnique({
      where: { id },
      include: {
        team: {
          include: { members: true }
        }
      }
    })
    if (!folder) {
      return { error: 'Folder not found' }
    }
    const isMember = folder.team.members.some(member => member.userId === session?.user?.id)
    if (!isMember) {
      return { error: 'Unauthorized' }
    }
    const updatedFolder = await prisma.folder.update({
      where: { id },
      data
    })
    return updatedFolder
  } catch (error) {
    console.error('Error updating folder:', error)
    return { error: 'Failed to update folder' }
  }
}

export async function deleteFolder(id: string) {
  const session = await auth()
  if (!session?.user?.id) {
    return { error: 'Unauthorized' }
  }

  try {
    const folder = await prisma.folder.findUnique({
      where: { id },
      include: {
        team: {
          include: { members: true }
        }
      }
    })
    if (!folder) {
      return { error: 'Folder not found' }
    }
    const isMember = folder.team.members.some(member => member.userId === session?.user?.id)
    if (!isMember) {
      return { error: 'Unauthorized' }
    }
    await prisma.folder.delete({ where: { id } })
    return { success: true }
  } catch (error) {
    console.error('Error deleting folder:', error)
    return { error: 'Failed to delete folder' }
  }
}

export async function moveChat(chatId: string, folderId?: string) {
  const session = await auth()
  if (!session?.user?.id) {
    return { error: 'Unauthorized' }
  }
}