/**
 * "use cache" Runtime
 *
 * Provides the runtime wrapper for functions marked with "use cache" directive.
 * The Vite transform plugin wraps exports with registerCachedFunction().
 *
 * On cache miss: executes the function, serializes the result via RSC Flight
 * protocol, captures handle data if tainted ctx is detected, and stores in
 * the SegmentCacheStore.
 *
 * On cache hit: deserializes the cached result, restores handle data if present.
 *
 * On stale hit: returns stale data immediately, triggers background
 * re-execution via waitUntil().
 */

/// <reference types="@vitejs/plugin-rsc/types" />

import {
  encodeReply,
  createClientTemporaryReferenceSet,
} from "@vitejs/plugin-rsc/rsc";
import { getRequestContext } from "../server/request-context.js";
import {
  isTainted,
  CACHED_FN_SYMBOL,
  isCachedFunction,
  stampCacheExec,
  unstampCacheExec,
} from "./taint.js";

export { isCachedFunction };
import { serializeResult, deserializeResult } from "./segment-codec.js";
import { createHandleStore } from "../server/handle-store.js";
import { restoreHandles } from "./handle-snapshot.js";
import { startHandleCapture, type HandleCapture } from "./handle-capture.js";
import { sortedSearchString } from "./cache-key-utils.js";
import { runBackground } from "./background-task.js";

/**
 * Convert encodeReply result to a stable string key.
 * encodeReply may return string or FormData — normalize to string.
 */
async function replyToCacheKey(encoded: string | FormData): Promise<string> {
  if (typeof encoded === "string") return encoded;
  // FormData: convert to Response body, then to string for deterministic key
  const text = await new Response(encoded).text();
  return text;
}

// ============================================================================
// Core: registerCachedFunction
// ============================================================================

/**
 * Register a function as a cached function.
 * Called by the Vite transform for each "use cache" function.
 *
 * @param fn - The original async function
 * @param id - Stable identifier (module path + export name)
 * @param profileName - Cache profile name (from "use cache: profileName" or "default")
 */
export function registerCachedFunction<T extends (...args: any[]) => any>(
  fn: T,
  id: string,
  profileName: string,
): T {
  const wrapped = async function (this: any, ...args: any[]): Promise<any> {
    const requestCtx = getRequestContext();
    const store = requestCtx?._cacheStore;
    const resolvedProfileName = profileName || "default";

    // Bypass: no store or no getItem support
    if (!store?.getItem) {
      return fn.apply(this, args);
    }

    // Resolve profile strictly from request-scoped config (set by the
    // active router via createRequestContext). No global fallback —
    // global profile state is only for DSL-time cache("profileName").
    const profile = requestCtx?._cacheProfiles?.[resolvedProfileName];

    if (!profile) {
      throw new Error(
        `[use cache] "${id}" uses unknown cache profile "${resolvedProfileName}". ` +
          `Define it in createRouter({ cacheProfiles: { "${resolvedProfileName}": { ttl: ... } } }).`,
      );
    }

    // Separate tainted args (ctx, env, req) from key-generating args.
    // For tainted objects that carry route context (params, pathname,
    // searchParams), extract serializable values into the key so
    // different routes, param combinations, and query variants produce
    // distinct cache entries.
    const keyArgs: unknown[] = [];
    let hasTaintedArgs = false;
    for (const arg of args) {
      if (isTainted(arg)) {
        hasTaintedArgs = true;
        const ctx = arg as any;
        if (ctx.params && typeof ctx.params === "object") {
          // Include host to prevent cross-host cache collisions (same
          // pattern as route-level cache-scope.ts key generation).
          if (ctx.url?.host) {
            keyArgs.push(ctx.url.host);
          }
          // Include route name to prevent collisions when the same cached
          // function is reused across routes with identical pathname/params
          // but different local reverse() scope.
          if (ctx._routeName) {
            keyArgs.push(ctx._routeName);
          }
          keyArgs.push(ctx.pathname, ctx.params);
          if (ctx._responseType) {
            keyArgs.push(ctx._responseType);
          }
          // Include user-facing search params (exclude internal _rsc*/__ params)
          if (ctx.searchParams instanceof URLSearchParams) {
            const normalized = sortedSearchString(ctx.searchParams);
            if (normalized) {
              keyArgs.push(normalized);
            }
          }
        }
      } else {
        keyArgs.push(arg);
      }
    }

    // If tainted args are present, we need the handle store for capture/restore.
    // During late streaming (Suspense boundary resolution), ALS context may be
    // gone. Throw early rather than silently dropping handle side effects.
    if (hasTaintedArgs && !requestCtx?._handleStore) {
      throw new Error(
        `[use cache] "${id}" receives a tainted argument (ctx/env/req) but the ` +
          `HandleStore is not available. This typically happens when a "use cache" ` +
          `function with ctx runs outside the request context (e.g., during late ` +
          `streaming after AsyncLocalStorage context is lost). Move the "use cache" ` +
          `directive to a function that does not receive request-scoped objects, or ` +
          `use the route-level cache() DSL instead.`,
      );
    }

    // Generate cache key
    let cacheKey: string;
    try {
      if (keyArgs.length > 0) {
        const tempRefs = createClientTemporaryReferenceSet();
        const encoded = await encodeReply(keyArgs as unknown[], {
          temporaryReferences: tempRefs,
        });
        const argsKey = await replyToCacheKey(encoded);
        cacheKey = `use-cache:${id}:${argsKey}`;
      } else {
        cacheKey = `use-cache:${id}`;
      }
    } catch {
      // Non-serializable args: run uncached
      return fn.apply(this, args);
    }

    // Cache lookup
    const cached = await store.getItem(cacheKey);

    if (cached && !cached.shouldRevalidate) {
      // Fresh hit: deserialize and return
      try {
        const result = await deserializeResult(cached.value);
        // Restore handle data if present
        if (cached.handles && hasTaintedArgs) {
          const handleStore = requestCtx?._handleStore;
          if (handleStore) {
            restoreHandles(cached.handles, handleStore);
          }
        }
        return result;
      } catch {
        // Deserialization failed, fall through to fresh execution
      }
    }

    if (cached?.shouldRevalidate) {
      // Stale hit: return stale value, revalidate in background
      try {
        const result = await deserializeResult(cached.value);
        if (cached.handles && hasTaintedArgs) {
          const handleStore = requestCtx?._handleStore;
          if (handleStore) {
            restoreHandles(cached.handles, handleStore);
          }
        }
        // Background revalidation — must capture handles if tainted args present.
        // Use an isolated handle store so background pushes don't pollute the
        // live response or throw LateHandlePushError on the completed store.
        // Same isolation pattern as route-level background-revalidation.ts.
        runBackground(requestCtx, async () => {
          // Reuse closure-captured requestCtx instead of calling
          // getRequestContext() — ALS context may be gone inside waitUntil.
          let originalHandleStore:
            | ReturnType<typeof createHandleStore>
            | undefined;
          if (hasTaintedArgs && requestCtx) {
            originalHandleStore = requestCtx._handleStore;
            requestCtx._handleStore = createHandleStore();
          }
          const bgHandleStore = hasTaintedArgs
            ? requestCtx?._handleStore
            : undefined;
          let bgCapture: HandleCapture | undefined;
          let bgStopCapture: (() => void) | undefined;
          if (bgHandleStore) {
            const c = startHandleCapture(bgHandleStore);
            bgCapture = c.capture;
            bgStopCapture = c.stop;
          }

          // Stamp tainted ARGS only — not requestCtx. The args stamp guards
          // direct ctx method calls (ctx.set, ctx.header, ctx.onResponse, etc.)
          // which is sufficient for correctness.
          //
          // We intentionally skip stamping requestCtx here because:
          // 1. runBackground starts the async task synchronously (before the
          //    first await), so stampCacheExec would pollute the shared
          //    requestCtx while the foreground pipeline is still running.
          //    This causes assertNotInsideCacheExec to fire when cache-store
          //    later calls requestCtx.onResponse().
          // 2. requestCtx methods are closure-bound to the original ctx, so
          //    neither Object.create() nor a proxy can isolate the stamp.
          // 3. The foreground miss path already stamps requestCtx and catches
          //    cookies()/headers() misuse on first execution. The background
          //    re-runs the same function with the same request.
          const bgTaintedArgs: unknown[] = [];
          for (const arg of args) {
            if (isTainted(arg)) {
              stampCacheExec(arg as object);
              bgTaintedArgs.push(arg);
            }
          }

          try {
            const freshResult = await fn.apply(this, args);
            bgStopCapture?.();
            const serialized = await serializeResult(freshResult);
            if (serialized !== null) {
              await store.setItem!(cacheKey, serialized, {
                handles: bgCapture?.data,
                ttl: profile.ttl,
                swr: profile.swr,
                tags: profile.tags,
              });
            }
          } catch (bgError) {
            bgStopCapture?.();
            requestCtx?._reportBackgroundError?.(bgError, "stale-revalidation");
          } finally {
            for (const arg of bgTaintedArgs) {
              unstampCacheExec(arg as object);
            }
            // Restore original handle store
            if (originalHandleStore && requestCtx) {
              requestCtx._handleStore = originalHandleStore;
            }
          }
        });
        return result;
      } catch {
        // Deserialization of stale value failed, fall through
      }
    }

    // Cache miss: execute, serialize, store
    const handleStore = hasTaintedArgs ? requestCtx?._handleStore : undefined;
    let capture: HandleCapture | undefined;
    let stopCapture: (() => void) | undefined;
    if (handleStore && hasTaintedArgs) {
      const c = startHandleCapture(handleStore);
      capture = c.capture;
      stopCapture = c.stop;
    }

    // Stamp tainted args so ctx.set(), ctx.header(), etc. throw if called
    // inside the cached function body (those side effects are lost on hit).
    // Uses ref-counted stamp/unstamp so overlapping executions
    // sharing the same ctx don't clear each other's guards.
    const taintedArgs: unknown[] = [];
    for (const arg of args) {
      if (isTainted(arg)) {
        stampCacheExec(arg as object);
        taintedArgs.push(arg);
      }
    }
    // Always stamp the ALS RequestContext so cookies()/headers() guards fire
    // even when the cached function receives no tainted args. The guard in
    // cookie-store.ts checks RequestContext, not function args.
    if (requestCtx) {
      stampCacheExec(requestCtx as object);
    }

    let result: any;
    try {
      result = await fn.apply(this, args);
    } finally {
      // Decrement ref count; symbol is deleted when it reaches zero
      for (const arg of taintedArgs) {
        unstampCacheExec(arg as object);
      }
      if (requestCtx) {
        unstampCacheExec(requestCtx as object);
      }
      // Remove this capture token (order-independent, safe for concurrent use)
      stopCapture?.();
    }

    // Serialize and store — fully non-blocking when waitUntil is available.
    // The response does not need to wait for serialization or the store write.
    const cacheWrite = async () => {
      try {
        const serialized = await serializeResult(result);
        if (serialized !== null) {
          await store.setItem!(cacheKey, serialized, {
            handles: capture?.data,
            ttl: profile.ttl,
            swr: profile.swr,
            tags: profile.tags,
          });
        }
      } catch (writeError) {
        requestCtx?._reportBackgroundError?.(writeError, "cache-write");
      }
    };

    await runBackground(requestCtx, cacheWrite, true);

    return result;
  };

  // Brand the wrapper so it can be detected at runtime (e.g., to prevent
  // accidental use as middleware).
  (wrapped as any)[CACHED_FN_SYMBOL] = true;

  return wrapped as unknown as T;
}