import type { QueryCtx } from '../_generated/server';
import { isActiveDocument } from './_helpers';

/**
 * Get RAG-indexed file storage IDs scoped to a agent's knowledge config.
 *
 * Access layers (in order):
 * 1. Agent-specific files — always included (passed in directly as knowledgeFileIds)
 * 2. Team documents — included when includeTeamKnowledge !== false and agent has teams
 * 3. Org-wide documents — included when includeOrgKnowledge is true
 * 4. Project documents — included when chatting inside a project (agentProjectIds set).
 *    Per the projects mutual-exclusivity rule, a project doc has projectId set and
 *    teamId cleared — so this branch never double-counts with the team layer.
 *
 * Only returns documents whose current blob is RAG-completed
 * (fileMetadata.ragStatus === 'completed') and that have a valid fileId.
 */
export async function getAgentScopedFileIds(
  ctx: QueryCtx,
  args: {
    organizationId: string;
    agentTeamId?: string;
    agentTeamIds?: string[];
    includeTeamKnowledge?: boolean;
    includeOrgKnowledge?: boolean;
    knowledgeFileIds?: string[];
    /**
     * Project IDs whose documents should be unioned into the file set.
     * Typically a single-element array (one project per chat).
     */
    agentProjectIds?: string[];
  },
): Promise<string[]> {
  const fileIdSet = new Set<string>(args.knowledgeFileIds ?? []);

  // Build effective team set: prefer agentTeamIds, fall back to single agentTeamId
  const agentTeamIdSet = new Set<string>();
  if (args.agentTeamIds) {
    for (const id of args.agentTeamIds) agentTeamIdSet.add(id);
  } else if (args.agentTeamId) {
    agentTeamIdSet.add(args.agentTeamId);
  }

  const agentProjectIdSet = new Set<string>(args.agentProjectIds ?? []);

  const needsTeamDocs =
    args.includeTeamKnowledge !== false && agentTeamIdSet.size > 0;
  const needsOrgDocs = args.includeOrgKnowledge === true;
  const needsProjectDocs = agentProjectIdSet.size > 0;

  if (!needsTeamDocs && !needsOrgDocs && !needsProjectDocs) {
    return [...fileIdSet];
  }

  // RAG completion now lives on fileMetadata.ragStatus (canonical). Enumerate
  // completed Document-Hub fileMetadata for the org and resolve each to its
  // document to apply lifecycle + team/project/org scoping. The `.gt(documentId,
  // undefined)` bound seeks past chat-upload / transcript rows (documentId
  // absent), so the scan stays bounded by Hub-doc count instead of the org's
  // whole completed corpus.
  const completedFiles = ctx.db
    .query('fileMetadata')
    .withIndex('by_organizationId_and_ragStatus_and_documentId', (q) =>
      q
        .eq('organizationId', args.organizationId)
        .eq('ragStatus', 'completed')
        .gt('documentId', undefined),
    );

  for await (const fm of completedFiles) {
    // Defensive — the index range already excludes documentId-absent rows.
    if (!fm.documentId) continue;
    const fileId = String(fm.storageId);
    if (fileIdSet.has(fileId)) continue;

    const doc = await ctx.db.get(fm.documentId);
    if (!doc || !doc.fileId) continue;
    // Only the document's CURRENT blob counts — a re-indexed doc can leave a
    // stale completed fileMetadata on its previous blob.
    if (String(doc.fileId) !== fileId) continue;
    // Skip trashed/soft-deleted docs (e.g. WebDAV DELETE) — they must not
    // remain in agent RAG retrieval scope.
    if (!isActiveDocument(doc)) continue;

    if (
      needsProjectDocs &&
      doc.projectId &&
      agentProjectIdSet.has(doc.projectId)
    ) {
      fileIdSet.add(fileId);
      continue;
    }

    if (
      needsTeamDocs &&
      doc.teamId &&
      !doc.projectId &&
      agentTeamIdSet.has(doc.teamId)
    ) {
      fileIdSet.add(fileId);
      continue;
    }

    if (needsOrgDocs && !doc.teamId && !doc.projectId) {
      fileIdSet.add(fileId);
    }
  }

  return [...fileIdSet];
}
