/**
 * FrontMcpContextStorage - AsyncLocalStorage wrapper for unified context
 *
 * Provides concurrent-safe context propagation using Node.js AsyncLocalStorage.
 * Access through DI only - never use static imports to access the storage directly.
 *
 * @example
 * ```typescript
 * // In a flow or middleware
 * const storage = this.get(FrontMcpContextStorage);
 * await storage.runFromHeaders(request.headers, {
 *   sessionId: sessionId,
 *   scopeId: scope.id,
 * }, async () => {
 *   // All code here can access the context via DI
 *   const ctx = this.get(FRONTMCP_CONTEXT);
 * });
 * ```
 */

import { type AuthInfo } from '@frontmcp/protocol';
import { AsyncLocalStorage } from '@frontmcp/utils';

import { Provider } from '../common/decorators';
import { ProviderScope } from '../common/metadata';
import { RequestContextNotAvailableError } from '../errors/mcp.error';
import { FrontMcpContext, type FrontMcpContextArgs } from './frontmcp-context';
import { extractMetadata } from './metadata.utils';
import { parseTraceContext } from './trace-context';

/**
 * Module-level AsyncLocalStorage instance.
 *
 * This is the ONLY place where the storage is created.
 * Access should be through DI, not through static imports.
 */
const frontmcpContextStorage = new AsyncLocalStorage<FrontMcpContext>();

/**
 * FrontMcpContextStorage provides unified context via AsyncLocalStorage.
 *
 * This is a GLOBAL-scoped provider because it manages the storage itself,
 * not the per-context data. The actual FrontMcpContext is accessed via
 * the FRONTMCP_CONTEXT token which is CONTEXT-scoped.
 */
@Provider({
  name: 'FrontMcpContextStorage',
  description: 'Manages unified context via AsyncLocalStorage',
  scope: ProviderScope.GLOBAL,
})
export class FrontMcpContextStorage {
  /**
   * Run a callback with a new FrontMcpContext.
   *
   * @param args - Arguments to create the context
   * @param fn - Async function to run with the context
   * @returns Result of the callback
   */
  run<T>(args: FrontMcpContextArgs, fn: () => T | Promise<T>): T | Promise<T> {
    const context = new FrontMcpContext(args);
    return frontmcpContextStorage.run(context, fn);
  }

  /**
   * Run with context extracted from HTTP headers.
   *
   * Automatically parses trace context from headers using W3C Trace Context
   * specification with fallback to x-frontmcp-trace-id.
   *
   * @param headers - HTTP headers
   * @param args - Additional context args (sessionId, scopeId)
   * @param fn - Async function to run
   * @returns Result of the callback
   */
  runFromHeaders<T>(
    headers: Record<string, unknown>,
    args: Omit<FrontMcpContextArgs, 'traceContext' | 'metadata'>,
    fn: () => T | Promise<T>,
  ): T | Promise<T> {
    const traceContext = parseTraceContext(headers);
    const metadata = extractMetadata(headers);
    const context = new FrontMcpContext({
      ...args,
      traceContext,
      metadata,
    });
    return frontmcpContextStorage.run(context, fn);
  }

  /**
   * Run with an existing FrontMcpContext.
   *
   * Useful when you need to propagate an existing context to a new async scope.
   *
   * @param context - Existing FrontMcpContext
   * @param fn - Async function to run
   * @returns Result of the callback
   */
  runWithContext<T>(context: FrontMcpContext, fn: () => T | Promise<T>): T | Promise<T> {
    return frontmcpContextStorage.run(context, fn);
  }

  /**
   * Get the current FrontMcpContext.
   *
   * @returns Current context or undefined if not in a context scope
   */
  getStore(): FrontMcpContext | undefined {
    return frontmcpContextStorage.getStore();
  }

  /**
   * Get the current FrontMcpContext, throwing if not available.
   *
   * @throws Error if not in a context scope
   */
  getStoreOrThrow(): FrontMcpContext {
    const ctx = this.getStore();
    if (!ctx) {
      throw new RequestContextNotAvailableError(
        'FrontMcpContext not available. Ensure operation runs within context scope.',
      );
    }
    return ctx;
  }

  /**
   * Check if currently running within a context.
   *
   * @returns True if a FrontMcpContext is available
   */
  hasContext(): boolean {
    return frontmcpContextStorage.getStore() !== undefined;
  }

  /**
   * Update the authInfo in the current context.
   *
   * This mutates the existing context in place to preserve internal state
   * (marks, store, sessionMetadata) while updating auth info.
   *
   * @param authInfo - Auth info fields to set/update (merged with existing)
   * @param fn - Function to run after update
   * @returns Result of the callback
   */
  updateAuthInfo<T>(authInfo: Partial<AuthInfo>, fn: () => T | Promise<T>): T | Promise<T> {
    const current = this.getStoreOrThrow();
    current.updateAuthInfo(authInfo);
    return fn();
  }
}
