simstudioai--sim/apps/sim/app/api/knowledge/utils.ts

import { db } from '@sim/db'
import { document, embedding, knowledgeBase } from '@sim/db/schema'
import { and, eq, isNull } from 'drizzle-orm'
import { getUserEntityPermissions } from '@/lib/workspaces/permissions/utils'

interface KnowledgeBaseData {
  id: string
  userId: string
  workspaceId?: string | null
  name: string
  description?: string | null
  tokenCount: number
  embeddingModel: string
  embeddingDimension: number
  chunkingConfig: unknown
  deletedAt?: Date | null
  createdAt: Date
  updatedAt: Date
}

interface DocumentData {
  id: string
  knowledgeBaseId: string
  filename: string
  fileUrl: string
  fileSize: number
  mimeType: string
  chunkCount: number
  tokenCount: number
  characterCount: number
  processingStatus: string
  processingStartedAt?: Date | null
  processingCompletedAt?: Date | null
  processingError?: string | null
  enabled: boolean
  deletedAt?: Date | null
  uploadedAt: Date
  // Text tags
  tag1?: string | null
  tag2?: string | null
  tag3?: string | null
  tag4?: string | null
  tag5?: string | null
  tag6?: string | null
  tag7?: string | null
  // Number tags (5 slots)
  number1?: number | null
  number2?: number | null
  number3?: number | null
  number4?: number | null
  number5?: number | null
  // Date tags (2 slots)
  date1?: Date | null
  date2?: Date | null
  // Boolean tags (3 slots)
  boolean1?: boolean | null
  boolean2?: boolean | null
  boolean3?: boolean | null
  // Connector fields
  connectorId?: string | null
  sourceUrl?: string | null
  externalId?: string | null
}

interface EmbeddingData {
  id: string
  knowledgeBaseId: string
  documentId: string
  chunkIndex: number
  chunkHash: string
  content: string
  contentLength: number
  tokenCount: number
  embedding?: number[] | null
  embeddingModel: string
  startOffset: number
  endOffset: number
  // Text tags
  tag1?: string | null
  tag2?: string | null
  tag3?: string | null
  tag4?: string | null
  tag5?: string | null
  tag6?: string | null
  tag7?: string | null
  // Number tags (5 slots)
  number1?: number | null
  number2?: number | null
  number3?: number | null
  number4?: number | null
  number5?: number | null
  // Date tags (2 slots)
  date1?: Date | null
  date2?: Date | null
  // Boolean tags (3 slots)
  boolean1?: boolean | null
  boolean2?: boolean | null
  boolean3?: boolean | null
  enabled: boolean
  createdAt: Date
  updatedAt: Date
}

export interface KnowledgeBaseAccessResult {
  hasAccess: true
  knowledgeBase: Pick<
    KnowledgeBaseData,
    'id' | 'userId' | 'workspaceId' | 'name' | 'embeddingModel'
  >
}

interface KnowledgeBaseAccessDenied {
  hasAccess: false
  notFound?: boolean
  reason?: string
}

export type KnowledgeBaseAccessCheck = KnowledgeBaseAccessResult | KnowledgeBaseAccessDenied

interface DocumentAccessResult {
  hasAccess: true
  document: DocumentData
  knowledgeBase: Pick<
    KnowledgeBaseData,
    'id' | 'userId' | 'workspaceId' | 'name' | 'embeddingModel'
  >
}

interface DocumentAccessDenied {
  hasAccess: false
  notFound?: boolean
  reason: string
}

export type DocumentAccessCheck = DocumentAccessResult | DocumentAccessDenied

interface ChunkAccessResult {
  hasAccess: true
  chunk: EmbeddingData
  document: DocumentData
  knowledgeBase: Pick<
    KnowledgeBaseData,
    'id' | 'userId' | 'workspaceId' | 'name' | 'embeddingModel'
  >
}

interface ChunkAccessDenied {
  hasAccess: false
  notFound?: boolean
  reason: string
}

export type ChunkAccessCheck = ChunkAccessResult | ChunkAccessDenied

/**
 * Resolve knowledge-base access for a user, gated by read or write permission.
 *
 * Read (`requireWrite: false`) grants on any workspace permission; write
 * (`requireWrite: true`) requires `write`/`admin`. Legacy non-workspace KBs grant
 * to the owning user in both modes.
 */
async function resolveKnowledgeBaseAccess(
  knowledgeBaseId: string,
  userId: string,
  requireWrite: boolean
): Promise<KnowledgeBaseAccessCheck> {
  const kb = await db
    .select({
      id: knowledgeBase.id,
      userId: knowledgeBase.userId,
      workspaceId: knowledgeBase.workspaceId,
      name: knowledgeBase.name,
      embeddingModel: knowledgeBase.embeddingModel,
    })
    .from(knowledgeBase)
    .where(and(eq(knowledgeBase.id, knowledgeBaseId), isNull(knowledgeBase.deletedAt)))
    .limit(1)

  if (kb.length === 0) {
    return { hasAccess: false, notFound: true }
  }

  const kbData = kb[0]

  if (kbData.workspaceId) {
    // Workspace KB: use workspace permissions only
    const userPermission = await getUserEntityPermissions(userId, 'workspace', kbData.workspaceId)
    const permitted = requireWrite
      ? userPermission === 'write' || userPermission === 'admin'
      : userPermission !== null
    return permitted ? { hasAccess: true, knowledgeBase: kbData } : { hasAccess: false }
  }

  // Legacy non-workspace KB: allow owner access
  if (kbData.userId === userId) {
    return { hasAccess: true, knowledgeBase: kbData }
  }

  return { hasAccess: false }
}

/**
 * Check if a user has read access to a knowledge base.
 */
export async function checkKnowledgeBaseAccess(
  knowledgeBaseId: string,
  userId: string
): Promise<KnowledgeBaseAccessCheck> {
  return resolveKnowledgeBaseAccess(knowledgeBaseId, userId, false)
}

/**
 * Check if a user has write access to a knowledge base.
 *
 * Write access is granted if:
 * 1. KB has a workspace: user has write or admin permissions on that workspace
 * 2. KB has no workspace (legacy): user owns the KB directly
 */
export async function checkKnowledgeBaseWriteAccess(
  knowledgeBaseId: string,
  userId: string
): Promise<KnowledgeBaseAccessCheck> {
  return resolveKnowledgeBaseAccess(knowledgeBaseId, userId, true)
}

/**
 * Resolve document access within a knowledge base, gated by read or write
 * permission on the KB (see {@link resolveKnowledgeBaseAccess}).
 */
async function resolveDocumentAccess(
  knowledgeBaseId: string,
  documentId: string,
  userId: string,
  requireWrite: boolean
): Promise<DocumentAccessCheck> {
  const kbAccess = await resolveKnowledgeBaseAccess(knowledgeBaseId, userId, requireWrite)

  if (!kbAccess.hasAccess) {
    return {
      hasAccess: false,
      notFound: kbAccess.notFound,
      reason: kbAccess.notFound ? 'Knowledge base not found' : 'Unauthorized knowledge base access',
    }
  }

  const doc = await db
    .select()
    .from(document)
    .where(
      and(
        eq(document.id, documentId),
        eq(document.knowledgeBaseId, knowledgeBaseId),
        eq(document.userExcluded, false),
        isNull(document.archivedAt),
        isNull(document.deletedAt)
      )
    )
    .limit(1)

  if (doc.length === 0) {
    return { hasAccess: false, notFound: true, reason: 'Document not found' }
  }

  return {
    hasAccess: true,
    document: doc[0] as DocumentData,
    knowledgeBase: kbAccess.knowledgeBase!,
  }
}

/**
 * Check if a user has read access to a document within a knowledge base.
 */
export async function checkDocumentAccess(
  knowledgeBaseId: string,
  documentId: string,
  userId: string
): Promise<DocumentAccessCheck> {
  return resolveDocumentAccess(knowledgeBaseId, documentId, userId, false)
}

/**
 * Check if a user has write access to a specific document.
 * Write access is granted if user has write access to the knowledge base.
 */
export async function checkDocumentWriteAccess(
  knowledgeBaseId: string,
  documentId: string,
  userId: string
): Promise<DocumentAccessCheck> {
  return resolveDocumentAccess(knowledgeBaseId, documentId, userId, true)
}

/**
 * Resolve chunk access within a document/knowledge base, gated by read or write
 * permission on the KB. The document must exist and be fully processed
 * (`processingStatus === 'completed'`) before its chunks are accessible.
 */
async function resolveChunkAccess(
  knowledgeBaseId: string,
  documentId: string,
  chunkId: string,
  userId: string,
  requireWrite: boolean
): Promise<ChunkAccessCheck> {
  const kbAccess = await resolveKnowledgeBaseAccess(knowledgeBaseId, userId, requireWrite)

  if (!kbAccess.hasAccess) {
    return {
      hasAccess: false,
      notFound: kbAccess.notFound,
      reason: kbAccess.notFound ? 'Knowledge base not found' : 'Unauthorized knowledge base access',
    }
  }

  const doc = await db
    .select()
    .from(document)
    .where(
      and(
        eq(document.id, documentId),
        eq(document.knowledgeBaseId, knowledgeBaseId),
        eq(document.userExcluded, false),
        isNull(document.archivedAt),
        isNull(document.deletedAt)
      )
    )
    .limit(1)

  if (doc.length === 0) {
    return { hasAccess: false, notFound: true, reason: 'Document not found' }
  }

  const docData = doc[0] as DocumentData

  // Chunks are only accessible once the document has finished processing.
  if (docData.processingStatus !== 'completed') {
    return {
      hasAccess: false,
      reason: `Document is not ready for access (status: ${docData.processingStatus})`,
    }
  }

  const chunk = await db
    .select()
    .from(embedding)
    .where(and(eq(embedding.id, chunkId), eq(embedding.documentId, documentId)))
    .limit(1)

  if (chunk.length === 0) {
    return { hasAccess: false, notFound: true, reason: 'Chunk not found' }
  }

  return {
    hasAccess: true,
    chunk: chunk[0] as EmbeddingData,
    document: docData,
    knowledgeBase: kbAccess.knowledgeBase!,
  }
}

/**
 * Check if a user has read access to a chunk within a document and knowledge base.
 */
export async function checkChunkAccess(
  knowledgeBaseId: string,
  documentId: string,
  chunkId: string,
  userId: string
): Promise<ChunkAccessCheck> {
  return resolveChunkAccess(knowledgeBaseId, documentId, chunkId, userId, false)
}

/**
 * Check if a user has write access to a chunk.
 *
 * Mirrors {@link checkChunkAccess} but requires write/admin on the knowledge
 * base's workspace (or KB ownership for legacy KBs), matching the permission
 * needed to create chunks. Used for chunk mutation (update and delete) so those
 * operations require the same permission as creation rather than read.
 */
export async function checkChunkWriteAccess(
  knowledgeBaseId: string,
  documentId: string,
  chunkId: string,
  userId: string
): Promise<ChunkAccessCheck> {
  return resolveChunkAccess(knowledgeBaseId, documentId, chunkId, userId, true)
}