import { AsyncLocalStorage } from 'async_hooks';
import { GalileoLogger } from './utils/galileo-logger';
import type {
  GalileoLoggerConfig,
  StartSessionOptions
} from './types/logging/logger.types';
import type { LocalMetricConfig } from './types/metrics.types';
import { StepWithChildSpans } from './types/logging/span.types';

/**
 * Context information that is automatically propagated through async execution chains.
 */
export interface ExperimentContext {
  /** (Optional) The project ID */
  projectId?: string;
  /** (Optional) The experiment ID currently being executed */
  experimentId?: string;
  /** (Optional) The project name for the current experiment */
  projectName?: string;
  /** The parent stack for nested spans (trace, workflow, agent spans) */
  parentStack?: StepWithChildSpans[];
  /** The session ID for the current logger context */
  sessionId?: string;
  /** The log stream name for the current logger context */
  logStreamName?: string;
}

/**
 * AsyncLocalStorage instance for propagating experiment context through async execution chains.
 */
export const experimentContext = new AsyncLocalStorage<ExperimentContext>();

/**
 * Logger context information that is automatically propagated through async execution chains.
 * This allows logger state (parent stack, session) to be available across async boundaries.
 */
export interface LoggerContext {
  /** The parent stack for nested spans (trace, workflow, agent spans) */
  parentStack?: StepWithChildSpans[];
  /** The session ID for the current logger context */
  sessionId?: string;
  /** The log stream name for the current logger context */
  logStreamName?: string;
}

/**
 * AsyncLocalStorage instance for propagating logger context through async execution chains.
 * This ensures that parent stack and session information are available across async boundaries.
 */
export const loggerContext = new AsyncLocalStorage<LoggerContext>();

/**
 * Options for identifying a logger by its key (project, logStreamName/experimentId, mode).
 */
export interface LoggerKeyOptions {
  /** (Optional) The project name */
  projectName?: string;
  /** (Optional) The project ID */
  projectId?: string;
  /** (Optional) The log stream name (used when experimentId is not provided) */
  logStreamName?: string;
  /** (Optional) The experiment ID (takes precedence over logStreamName) */
  experimentId?: string;
  /** (Optional) The logger mode (defaults to 'batch') */
  mode?: string;
}

/**
 * Extends LoggerKeyOptions with localMetrics to configure new logger instances.
 */
export interface GetLoggerOptions extends LoggerKeyOptions {
  /** (Optional) Local metrics to run on traces/spans (only used when initializing a new logger) */
  localMetrics?: LocalMetricConfig[];
}

/**
 * A singleton class that manages a collection of GalileoLogger instances.
 */
export class GalileoSingleton {
  private static instance: GalileoSingleton;
  private galileoLoggers: Map<string, GalileoLogger> = new Map();
  private lastAvailableLogger: GalileoLogger | null = null;

  private constructor() {}

  /**
   * Gets the singleton instance of GalileoSingleton.
   * @returns The singleton instance
   */
  public static getInstance(): GalileoSingleton {
    if (!GalileoSingleton.instance) {
      GalileoSingleton.instance = new GalileoSingleton();
    }
    return GalileoSingleton.instance;
  }

  /**
   * Returns the last available logger instance, or creates a new one if no logger is available.
   * @deprecated Use getLogger() method instead. This method is kept for backwards compatibility.
   * @returns An instance of GalileoLogger
   */
  public getClient(): GalileoLogger {
    return this.lastAvailableLogger ?? this.getLogger();
  }

  /**
   * Generates a key string based on project, logStreamName/experimentId, and mode parameters.
   * @param projectName - (Optional) The project name
   * @param logStreamName - (Optional) The log stream name (used when experimentId is not provided)
   * @param experimentId - (Optional) The experiment ID (takes precedence over logStreamName)
   * @param mode - (Optional) The logger mode (defaults to "batch")
   * @returns A string key used for caching
   */
  private static _getKey(
    projectName?: string,
    logStreamName?: string,
    experimentId?: string,
    mode?: string
  ): string {
    // Get context from AsyncLocalStorage
    const context = experimentContext.getStore();

    // Apply fallbacks: explicit parameter -> context -> environment variable -> default
    const finalProjectName =
      projectName ??
      context?.projectName ??
      process.env.GALILEO_PROJECT ??
      process.env.GALILEO_PROJECT_NAME ??
      'default';
    const finalLogStream =
      logStreamName ??
      process.env.GALILEO_LOG_STREAM ??
      process.env.GALILEO_LOG_STREAM_NAME ??
      'default';

    // Use experimentId if provided, otherwise check context, otherwise use log_stream
    const identifier = experimentId ?? context?.experimentId ?? finalLogStream;

    // Return a string key: "project:identifier:mode"
    return `${finalProjectName}:${identifier}:${mode || 'batch'}`;
  }

  /**
   * Retrieves an existing GalileoLogger or creates a new one if it does not exist.
   * @param options - Configuration options
   * @param options.projectName - (Optional) The project name
   * @param options.logStreamName - (Optional) The log stream name (used when experimentId is not provided)
   * @param options.experimentId - (Optional) The experiment ID (takes precedence over logStreamName)
   * @param options.mode - (Optional) The logger mode (defaults to "batch")
   * @param options.localMetrics - (Optional) Local metrics to run on traces/spans (only used when initializing a new logger)
   * @returns An instance of GalileoLogger corresponding to the key
   */
  public getLogger(options: GetLoggerOptions = {}): GalileoLogger {
    // Get context from AsyncLocalStorage for fallback values
    const context = experimentContext.getStore();

    // Compute the key based on provided parameters, context, or environment variables
    const key = GalileoSingleton._getKey(
      options.projectName,
      options.logStreamName,
      options.experimentId,
      options.mode
    );

    // First check if logger already exists
    if (this.galileoLoggers.has(key)) {
      return this.galileoLoggers.get(key)!;
    }

    // Create new logger
    // Prepare initialization arguments, using context as fallback if not provided
    const config: GalileoLoggerConfig = {
      projectName:
        options.projectName ??
        context?.projectName ??
        process.env.GALILEO_PROJECT ??
        process.env.GALILEO_PROJECT_NAME,
      projectId: options.projectId ?? context?.projectId,
      logStreamName:
        options.logStreamName ??
        process.env.GALILEO_LOG_STREAM ??
        process.env.GALILEO_LOG_STREAM_NAME,
      experimentId: options.experimentId ?? context?.experimentId,
      localMetrics: options.localMetrics,
      mode: options.mode,
      onTerminate: (terminatedLogger) =>
        this.cleanupLogger(key, terminatedLogger as GalileoLogger)
    };

    const logger = new GalileoLogger(config);

    // Cache the newly created logger
    this.galileoLoggers.set(key, logger);
    this.lastAvailableLogger = logger;
    return logger;
  }

  /**
   * Retrieve a copy of the map containing all active loggers.
   *
   * @returns A map of keys to GalileoLogger instances
   */
  public getAllLoggers(): Map<string, GalileoLogger> {
    // Return a shallow copy to prevent external modifications
    return new Map(this.galileoLoggers);
  }

  /**
   * Removes a logger from internal tracking. Idempotent — only removes the
   * map entry / clears lastAvailableLogger if they still reference the given
   * logger, so it can be safely called from both the onTerminate hook and the
   * reset() backstop.
   */
  private cleanupLogger(key: string, logger: GalileoLogger): void {
    if (this.galileoLoggers.get(key) === logger) {
      this.galileoLoggers.delete(key);
    }
    if (this.lastAvailableLogger === logger) {
      this.lastAvailableLogger = null;
    }
  }

  /**
   * Resets (terminates and removes) a GalileoLogger instance.
   * @param options - Configuration options to identify which logger to reset
   * @param options.projectName - (Optional) The project name
   * @param options.logStreamName - (Optional) The log stream name
   * @param options.experimentId - (Optional) The experiment ID
   * @param options.mode - (Optional) The logger mode
   * @returns A promise that resolves when the logger is reset
   */
  public async reset(options: LoggerKeyOptions = {}): Promise<void> {
    const key = GalileoSingleton._getKey(
      options.projectName,
      options.logStreamName,
      options.experimentId,
      options.mode
    );

    const logger = this.galileoLoggers.get(key);
    if (logger) {
      // terminate() triggers the onTerminate callback which calls cleanupLogger().
      // The defensive call below is a backstop for cases where onTerminate doesn't
      // fire (e.g. legacy setClient() loggers that lack the hook).
      await logger.terminate();
      this.cleanupLogger(key, logger);
    }
  }

  /**
   * Resets (terminates and removes) all GalileoLogger instances.
   * @returns A promise that resolves when all loggers are reset
   */
  public async resetAll(): Promise<void> {
    // Snapshot values first — terminate() mutates galileoLoggers via the onTerminate callback.
    const loggers = Array.from(this.galileoLoggers.values());
    await Promise.all(loggers.map((logger) => logger.terminate()));
    // Defensive backstop: clear any loggers that lacked an onTerminate hook
    // (e.g., ones added via the legacy setClient path prior to this change).
    this.galileoLoggers.clear();
    this.lastAvailableLogger = null;
  }

  /**
   * Flushes (uploads) a GalileoLogger instance.
   * @param options - Configuration options to identify which logger to flush
   * @param options.projectName - (Optional) The project name
   * @param options.logStreamName - (Optional) The log stream name
   * @param options.experimentId - (Optional) The experiment ID
   * @param options.mode - (Optional) The logger mode
   * @returns A promise that resolves when the logger is flushed
   */
  public async flush(options: LoggerKeyOptions = {}): Promise<void> {
    const key = GalileoSingleton._getKey(
      options.projectName,
      options.logStreamName,
      options.experimentId,
      options.mode
    );

    const logger = this.galileoLoggers.get(key);
    if (logger) {
      await logger.flush();
    }
  }

  /**
   * Flushes (uploads) all GalileoLogger instances.
   * @returns A promise that resolves when all loggers are flushed
   */
  public async flushAll(): Promise<void> {
    const flushPromises = Array.from(this.galileoLoggers.values()).map(
      (logger) => logger.flush()
    );
    await Promise.all(flushPromises);
  }

  // Legacy methods for backward compatibility

  /**
   * Sets a client logger instance.
   * @deprecated Use getLogger() method instead. This maintains backward compatibility.
   * @param client - The GalileoLogger instance to set
   */
  public setClient(client: GalileoLogger): void {
    // Store with default key
    const key = GalileoSingleton._getKey();
    this.galileoLoggers.set(key, client);
    this.lastAvailableLogger = client;
  }
}

/**
 * Gets or creates a logger and optionally initializes a session.
 *
 * Note: When startNewSession is true and an existing session is found via externalId, it is returned as-is.
 * All other session options (sessionName, previousSessionId, metadata) are ignored in that case -
 * they only apply when creating a new session. To update an existing session, use an explicit update method.
 *
 * @param options - Configuration options to initialize the logger
 * @param options.projectName - (Optional) The project name
 * @param options.logStreamName - (Optional) The log stream name
 * @param options.experimentId - (Optional) The experiment ID
 * @param options.mode - (Optional) The logger mode
 * @param options.localMetrics - (Optional) Local metrics to run on traces/spans (only used when initializing a new logger)
 * @param options.sessionId - (Optional) The session ID
 * @param options.startNewSession - (Optional) Whether to start a new session
 * @param options.sessionName - (Optional) The name of the session. Only applied when creating a new session.
 * @param options.previousSessionId - (Optional) The ID of a previous session to link to. Creates a reference only; does not inherit metadata. Only applied when creating a new session.
 * @param options.externalId - (Optional) An external identifier for the session. If a session with this external ID already exists, it will be reused.
 * @param options.metadata - (Optional) User metadata for the session. Only applied when creating a new session.
 * @returns A promise that resolves when initialization is complete
 */
export const init = async (
  options: GetLoggerOptions & {
    sessionId?: string;
    startNewSession?: boolean;
    sessionName?: string;
    previousSessionId?: string;
    externalId?: string;
    metadata?: Record<string, string>;
  } = {}
) => {
  const singleton = GalileoSingleton.getInstance();
  const logger = singleton.getLogger({
    projectName: options.projectName,
    projectId: options.projectId,
    logStreamName: options.logStreamName,
    experimentId: options.experimentId,
    mode: options.mode,
    localMetrics: options.localMetrics
  });

  if (options.startNewSession) {
    await logger.startSession({
      name: options.sessionName,
      previousSessionId: options.previousSessionId,
      externalId: options.externalId,
      metadata: options.metadata
    });
  }
};

/**
 * Retrieves an existing GalileoLogger or creates a new one if it does not exist.
 * @param options - Configuration options
 * @param options.projectName - (Optional) The project name
 * @param options.logStreamName - (Optional) The log stream name (used when experimentId is not provided)
 * @param options.experimentId - (Optional) The experiment ID (takes precedence over logStreamName)
 * @param options.mode - (Optional) The logger mode (defaults to "batch")
 * @param options.localMetrics - (Optional) Local metrics to run on traces/spans (only used when initializing a new logger)
 * @returns An instance of GalileoLogger corresponding to the key
 */
export const getLogger = (options: GetLoggerOptions = {}) => {
  return GalileoSingleton.getInstance().getLogger(options);
};

/**
 * Returns the logger for the current async context (experimentContext + loggerContext).
 * Used by session helpers and the log() wrapper so they target the same logger.
 */
const getLoggerFromContext = (): GalileoLogger => {
  const exp = experimentContext.getStore();
  const logStore = loggerContext.getStore();
  return GalileoSingleton.getInstance().getLogger({
    projectName: exp?.projectName,
    experimentId: exp?.experimentId,
    logStreamName: logStore?.logStreamName ?? exp?.logStreamName
  });
};

/**
 * Starts a new session on the logger for the current context.
 *
 * @param options - (Optional) Session options.
 * @param options.name - (Optional) The session name.
 * @param options.previousSessionId - (Optional) The previous session ID to link to.
 * @param options.externalId - (Optional) External ID for the session.
 * @param options.metadata - (Optional) User metadata for the session as key-value string pairs. Only applied when creating a new session.
 * @returns A promise that resolves to the session ID.
 */
export const startSession = async (
  options?: StartSessionOptions
): Promise<string> => {
  const logger = getLoggerFromContext();
  return logger.startSession(options);
};

/**
 * Sets the session ID on the logger for the current context.
 * Traces created via log() are associated with this session.
 *
 * @param sessionId - The session ID to set.
 * @returns Nothing.
 */
export const setSession = (sessionId: string): void => {
  getLoggerFromContext().setSessionId(sessionId);
};

/**
 * Clears the current session ID on the logger for the current context.
 * Subsequent traces are not associated with a session until startSession or setSession is called.
 *
 * @returns Nothing.
 */
export const clearSession = (): void => {
  getLoggerFromContext().clearSession();
};

/**
 * Retrieves a shallow copy of the map containing all active loggers.
 * @returns A map of keys to GalileoLogger instances
 */
export const getAllLoggers = () => {
  return GalileoSingleton.getInstance().getAllLoggers();
};

/**
 * Resets (terminates and removes) a specific logger instance.
 * @param options - Configuration options to identify which logger to reset
 * @param options.projectName - (Optional) The project name
 * @param options.logStreamName - (Optional) The log stream name
 * @param options.experimentId - (Optional) The experiment ID
 * @param options.mode - (Optional) The logger mode
 * @returns A promise that resolves when the logger is reset
 */
export const reset = async (options: LoggerKeyOptions = {}) => {
  await GalileoSingleton.getInstance().reset(options);
};

/**
 * Resets (terminates and removes) all logger instances.
 * @returns A promise that resolves when all loggers are reset
 */
export const resetAll = async () => {
  await GalileoSingleton.getInstance().resetAll();
};

/**
 * Flushes (uploads) traces from a specific logger to the Galileo platform.
 * @param options - Configuration options to identify which logger to flush
 * @param options.projectName - (Optional) The project name
 * @param options.logStreamName - (Optional) The log stream name
 * @param options.experimentId - (Optional) The experiment ID
 * @param options.mode - (Optional) The logger mode
 * @returns A promise that resolves when the logger is flushed
 */
export const flush = async (options: LoggerKeyOptions = {}) => {
  await GalileoSingleton.getInstance().flush(options);
};

/**
 * Flushes (uploads) all captured traces to the Galileo platform.
 * @returns A promise that resolves when all loggers are flushed
 */
export const flushAll = async () => {
  await GalileoSingleton.getInstance().flushAll();
};

/**
 * Lifecycle and context API for Galileo logging.
 * Groups init, flush, reset, and session methods for ergonomic use.
 */
export const galileoContext = {
  init,
  flush,
  flushAll,
  reset,
  resetAll,
  startSession,
  setSession,
  clearSession
};