import { AsyncLocalStorage } from 'async_hooks';
import { v4 as uuidv4 } from 'uuid';
import { logger } from '../../utils/logger';
import { NextFunction, Request, Response } from 'express';

/**
 * [BIZ_KER_128] 全链路追踪 ID 注入与上报 (Trace ID)
 * @description 核心逻辑：利用 AsyncLocalStorage 在异步调用链中传递 TraceID。
 * 自动为进入系统的每个请求注入唯一 TraceID，并透传至日志与下游服务调用（如 Axios）。
 * 遵循 Autocomplete-First (V31.5) 规范。
 */
export class GlobalTracingService {
  private static readonly storage = new AsyncLocalStorage<Map<string, string>>();
  private static readonly TRACE_ID_KEY = 'x-trace-id';

  /**
   * 获取当前上下文中的 TraceID
   * @returns 当前请求的 TraceID，若不存在则返回 null
   */
  static getTraceId(): string | null {
    const store = this.storage.getStore();
    return store ? store.get(this.TRACE_ID_KEY) || null : null;
  }

  /**
   * 启动追踪上下文中间件 (Express Middleware)
   * @param req Express Request
   * @param res Express Response
   * @param next Next Function
   */
  static middleware(req: Request, res: Response, next: NextFunction) {
    const traceId = (req.headers[GlobalTracingService.TRACE_ID_KEY] as string) || uuidv4();
    const store = new Map<string, string>();
    store.set(GlobalTracingService.TRACE_ID_KEY, traceId);

    // 将 TraceID 注入响应头
    res.setHeader(GlobalTracingService.TRACE_ID_KEY, traceId);

    GlobalTracingService.storage.run(store, () => {
      next();
    });
  }

  /**
   * 注入 TraceID 到 Axios 请求拦截器
   * @param config Axios Request Config
   * @returns 修改后的 Config
   */
  static injectToAxios(config: any) {
    const traceId = this.getTraceId();
    if (traceId) {
      config.headers[this.TRACE_ID_KEY] = traceId;
    }
    return config;
  }

  /**
   * 带 TraceID 的安全日志包装
   * @param level 日志级别
   * @param message 消息内容
   * @param data 附加数据
   */
  static logWithTrace(level: 'info' | 'warn' | 'error', message: string, data?: any) {
    const traceId = this.getTraceId();
    const logPayload = {
      traceId,
      ...data
    };
    (logger as any)[level](`[TRACED] ${message}`, logPayload);
  }
}
