'use node';

import { v } from 'convex/values';

import { isRecord, getBoolean, getString } from '../../lib/utils/type-guards';
import { internal } from '../_generated/api';
import { action } from '../_generated/server';
import { orgSlugFromIdOrNull } from '../lib/helpers/org_slug';
import { ragFetch } from '../lib/helpers/rag_config';
import { getAuthUserIdentity } from '../lib/rls/auth/get_auth_user_identity';

/**
 * Check RAG indexing status for a list of files and update fileMetadata.
 *
 * Called by the frontend on an interval while files are in queued/running
 * state. Stops being called when the user leaves the page — no wasted
 * server-side scheduled actions.
 *
 * Auth: caller must be authenticated AND the supplied storageIds must
 * belong to fileMetadata rows whose organizationId is one of the caller's
 * org memberships. Without this gate, an anonymous attacker can flip any
 * org's `ragStatus` to `failed` via `expireStaleRagQueue` (DoS), and a
 * member of org A can poke org B's RAG status. (Pre-existing on `main`
 * but in scope for this branch's RAG-auth surface.)
 */
export const checkFileRagStatuses = action({
  args: {
    storageIds: v.array(v.id('_storage')),
  },
  returns: v.null(),
  handler: async (ctx, args): Promise<null> => {
    if (args.storageIds.length === 0) return null;

    const authUser = await getAuthUserIdentity(ctx);
    if (!authUser) {
      console.warn('[checkFileRagStatuses] unauthenticated caller — refused');
      return null;
    }
    const callerId = authUser.userId;

    // Filter storageIds down to ones the caller is authorized to see, and
    // get the org for each so we can call RAG (which is now per-org) with
    // the correct X-Tale-Org header per group.
    const allowed = await ctx.runQuery(
      internal.file_metadata.internal_queries.filterStorageIdsByCallerOrg,
      { storageIds: args.storageIds, userId: callerId },
    );
    if (allowed.length === 0) {
      console.warn(
        '[checkFileRagStatuses] no authorized storage ids for caller — refused',
      );
      return null;
    }

    // Group authorized storage ids by org so we can issue one RAG call
    // per distinct org. The cache means each org slug is resolved once
    // even when many files belong to the same org.
    const orgIdsToFiles = new Map<
      string,
      Array<(typeof args.storageIds)[number]>
    >();
    for (const { storageId, organizationId } of allowed) {
      const bucket = orgIdsToFiles.get(organizationId);
      if (bucket) bucket.push(storageId);
      else orgIdsToFiles.set(organizationId, [storageId]);
    }

    // Give RAG 90s to have ingested a newly-queued upload. If we're still
    // getting null after that window, the upload never reached RAG (likely
    // the scheduled action was dropped before it ran) — mark failed so the
    // client stops polling. Threshold is measured against `ragQueuedAt` on
    // the fileMetadata row, so re-queues reset the clock.
    const STALE_QUEUE_MS = 90_000;

    const mergedStatuses: Record<string, unknown> = {};
    // Track storageIds whose org bucket queried RAG SUCCESSFULLY. Only
    // these are eligible for the post-loop `expireStaleRagQueue` sweep —
    // without this guard, a transient RAG outage in one org's request
    // (or that org's slug going missing) made the loop `continue`, the
    // org's storageIds never landed in `mergedStatuses`, and the sweep
    // permanently marked them `failed` ("RAG service did not receive
    // the upload") even though the uploads were healthy. Cross-org
    // failure propagation.
    const eligibleForStaleSweep = new Set<string>();
    for (const [organizationId, storageIds] of orgIdsToFiles) {
      const orgSlug = await orgSlugFromIdOrNull(ctx, organizationId);
      if (orgSlug === null) {
        console.warn(
          `[checkFileRagStatuses] org ${organizationId} unresolvable; skipping status fetch (its storageIds will NOT be marked failed)`,
        );
        continue;
      }
      try {
        const response = await ragFetch('/api/v1/documents/statuses', {
          method: 'POST',
          headers: { 'Content-Type': 'application/json' },
          body: JSON.stringify({ file_ids: storageIds }),
          timeoutMs: 10_000,
          orgSlug,
        });
        if (!response.ok) {
          console.warn(
            `[checkFileRagStatuses] RAG returned ${response.status} for org ${orgSlug}`,
          );
          continue;
        }
        const body: unknown = await response.json();
        if (!isRecord(body) || !isRecord(body.statuses)) continue;
        Object.assign(mergedStatuses, body.statuses);
        for (const id of storageIds) eligibleForStaleSweep.add(id);
      } catch (error) {
        console.warn(
          `[checkFileRagStatuses] Failed to fetch statuses for org ${orgSlug}:`,
          error,
        );
      }
    }

    const statuses = mergedStatuses;
    const allAuthorizedStorageIds = allowed.map((a) => a.storageId);

    for (const storageId of allAuthorizedStorageIds) {
      const docStatus = statuses[storageId];
      if (!isRecord(docStatus)) {
        if (eligibleForStaleSweep.has(storageId)) {
          await ctx.runMutation(
            internal.file_metadata.internal_mutations.expireStaleRagQueue,
            { storageId, staleAfterMs: STALE_QUEUE_MS },
          );
        }
        continue;
      }

      const status = getString(docStatus, 'status');
      const error = getString(docStatus, 'error');
      const progressPhase = getString(docStatus, 'progress_phase');
      const progressDetail = getString(docStatus, 'progress_detail');

      const ragProgress =
        progressPhase && progressDetail
          ? `${progressPhase} ${progressDetail}`
          : progressPhase || undefined;

      if (status === 'completed') {
        const ocrApplied = getBoolean(docStatus, 'ocr_applied');
        await ctx.runMutation(
          internal.file_metadata.internal_mutations.updateFileRagStatus,
          {
            storageId,
            ragStatus: 'completed',
            ...(ocrApplied != null && { ocrApplied }),
          },
        );
      } else if (status === 'failed') {
        await ctx.runMutation(
          internal.file_metadata.internal_mutations.updateFileRagStatus,
          {
            storageId,
            ragStatus: 'failed',
            ragError: error || 'Unknown error',
          },
        );
      } else if (status === 'processing') {
        await ctx.runMutation(
          internal.file_metadata.internal_mutations.updateFileRagStatus,
          { storageId, ragStatus: 'running', ragProgress },
        );
      }
    }

    return null;
  },
});
