// Copyright Abridged, Inc. 2021,2024. All Rights Reserved.
// Node module: @collabland/aws
// This file is licensed under the MIT License.
// License text available at https://opensource.org/licenses/MIT

import {
  PutTraceSegmentsRequest,
  PutTraceSegmentsResult,
  XRay,
  XRayClientConfig,
} from '@aws-sdk/client-xray';
import {BuildMiddleware, Client, MetadataBearer} from '@aws-sdk/types';
import {
  AnyError,
  FetchInterceptor,
  getEnvVarAsBoolean,
  loggers,
} from '@collabland/common';
import crypto from 'crypto';
import {URL} from 'url';
import {getDefaultConfig} from '../common/types.js';
import {Segment, Subsegment, TraceId} from './types.js';

const {trace} = loggers('collabland:tracing:xray');
// const {trace} = loggers('collabland:tracing:xray');

export class XRayServiceClient {
  readonly xray: XRay;

  constructor(options?: XRayClientConfig) {
    this.xray = new XRay({
      ...getDefaultConfig(),
      ...options,
    });
  }

  async send(...segments: XRaySubsegment[]) {
    const args: PutTraceSegmentsRequest = {
      TraceSegmentDocuments: segments.map(s => JSON.stringify(s)),
    };
    trace('Sending traces: %O', args);
    const res = await this.xray.putTraceSegments(args);
    trace('Response: %O', res);
    return res;
  }

  startSegment(segment: Partial<XRaySegment>): Promise<PutTraceSegmentsResult> {
    segment.id = segment.id ?? generateXRayId();
    segment.trace_id = segment.trace_id ?? new XRayTraceId().toString();
    segment.start_time = segment.start_time ?? Date.now() / 1000;
    segment.in_progress = true;
    trace('Starting segment: %j', segment);
    return this.send(segment as XRaySegment);
  }

  endSegment(segment: XRaySubsegment) {
    segment.end_time = Date.now() / 1000;
    segment.in_progress = false;
    trace('Ending segment: %j', segment);
    return this.send(segment);
  }

  async trace<T>(fn: () => T | Promise<T>, segment: XRaySubsegment | string) {
    if (typeof segment === 'string') {
      segment = new XRaySegment(this, segment);
    }
    try {
      trace('Starting segment: %j', segment);
      await segment.start();
      const result = await fn();
      return result;
    } catch (err: AnyError) {
      trace('Checking error for segment: %j', segment);
      segment.checkHttpRes(err);
      throw err;
    } finally {
      trace('Ending segment: %j', segment);
      await segment.end();
    }
  }

  getFetchInterceptor() {
    if (!isXRayEnabled()) return undefined;
    const fetchInterceptor: FetchInterceptor = (next, info, init) => {
      const url =
        info instanceof URL
          ? info.toString()
          : typeof info === 'string'
            ? info
            : info.url;
      const method = init.method;
      const segment = new XRaySegment(this, 'fetch');
      segment.http = {
        request: {
          method,
          url,
        },
        response: {},
      };
      return this.trace(async () => {
        const result = await next();
        segment.checkHttpRes({statusCode: result.status});
        segment.http!.response!.content_length =
          result.headers.get('content-length') ?? undefined;
        return result;
      }, segment);
    };
    return fetchInterceptor;
  }

  addMiddleware<Input extends object, Output extends MetadataBearer>(
    client: Client<Input, Output, unknown>,
    name?: string,
  ): BuildMiddleware<Input, Output> | undefined {
    if (!isXRayEnabled()) return undefined;
    const mw: BuildMiddleware<Input, Output> =
      (next, context) => async args => {
        name = name ?? client.constructor.name;
        trace('Invoking XRay middleware for %s', name);
        const segment = new XRaySegment(this, name);
        const fn = async () => {
          const result = await next(args);
          segment.checkHttpRes({
            statusCode: result.output.$metadata.httpStatusCode,
          });
          return result;
        };
        const result = await this.trace(fn, segment);
        trace('XRay middleware for %s is executed', name);
        return result;
      };
    client.middlewareStack.add(mw, {
      step: 'build',
      name: 'x-ray-tracing',
      tags: ['TRACING', 'X-RAY'],
    });
    return mw;
  }
}

export function generateXRayId() {
  return crypto.randomBytes(8).toString('hex');
}

/**
 * Subsegment backed by XRay
 */
export class XRaySubsegment implements Subsegment {
  trace_id?: string;
  parent_id?: string;
  id: string;
  name: string;
  start_time: number;
  end_time?: number;
  in_progress?: boolean;
  origin = 'AWS::ElasticBeanstalk::Environment';
  http?: {
    request?: Record<string, string | number | boolean | undefined>;
    response?: Record<string, string | number | boolean | undefined>;
  };
  readonly subsegments: Subsegment[] = [];
  [name: string]: unknown;

  constructor(
    readonly xray: XRayServiceClient,
    name?: string,
  ) {
    this.name = name ?? '';
    this.id = generateXRayId();
  }

  toJSON() {
    const json: Record<string, unknown> = {...this};
    for (const p in json) {
      if (typeof json[p] === 'function') {
        delete json[p];
      }
    }
    delete json.parent;
    delete json.xray;
    return json;
  }

  start(): Promise<PutTraceSegmentsResult> {
    return this.xray.startSegment(this);
  }

  end() {
    return this.xray.endSegment(this);
  }

  checkHttpRes(res: {statusCode?: number}) {
    if (typeof res.statusCode !== 'number') return res;
    if (res.statusCode >= 200 && res.statusCode < 300) {
      return undefined;
    }
    if (res.statusCode >= 400 && res.statusCode < 500) {
      if (res.statusCode === 429) {
        this.throttle = true;
      } else {
        this.error = true;
      }
    } else if (res.statusCode >= 500) {
      this.fault = true;
    } else {
      this.fault = true;
    }
    return res;
  }

  /**
   * Add and start a subsegment
   * @param subsegment - Subsegment object or name
   */
  startSubsegment(subsegment: XRaySubsegment | string) {
    subsegment = this.addSubsegment(subsegment);
    return subsegment.start();
  }

  addSubsegment(subsegment: string | XRaySubsegment) {
    if (typeof subsegment === 'string') {
      subsegment = new XRaySubsegment(this.xray, subsegment);
    }
    subsegment.parent_id = this.id;
    subsegment.trace_id = this.trace_id;
    return subsegment;
  }
}

/**
 * Segment backed by XRay
 */
export class XRaySegment extends XRaySubsegment implements Segment {
  trace_id: string;

  constructor(
    xray: XRayServiceClient,
    name: string,
    traceId?: string,
    parentId?: string,
  ) {
    super(xray, name);
    this.id = generateXRayId();
    this.trace_id = traceId ?? XRayTraceId.generate();
    if (parentId != null) this.parent_id = parentId;
  }

  start() {
    return this.xray.startSegment(this);
  }

  end() {
    return this.xray.endSegment(this);
  }
}

/**
 * Class describing an AWS X-Ray trace ID.
 * @see https://docs.aws.amazon.com/xray/latest/devguide/xray-concepts.html#xray-concepts-traces
 */
export class XRayTraceId implements TraceId {
  readonly version = 1;
  timestamp: string;
  id: string;
  /**
   * Constructs a new trace ID using the current time.
   * @constructor
   */
  constructor() {
    this.version = 1;
    this.timestamp = Math.round(Date.now() / 1000).toString(16);
    this.id = crypto.randomBytes(12).toString('hex');
  }

  static generate() {
    return new XRayTraceId().toString();
  }

  /**
   * Constructs a new trace ID from provided string. If no string is provided
   * or the provided string is invalid,
   * log an error but a new trace ID still returned. This can be used as a
   * trace ID string validator.
   * @param rawId - string to create a Trace ID object from.
   */
  static from(rawId?: string) {
    const traceId = new XRayTraceId();

    if (!rawId || typeof rawId !== 'string') {
      return traceId;
    }

    const parts = rawId.trim().split('-');
    if (parts.length !== 3) {
      return traceId;
    }

    const version = parseInt(parts[0]);
    if (isNaN(version) || version < 1) {
      return traceId;
    }

    const timestamp = parseInt(parts[1], 16).toString(16);
    if (timestamp === 'NaN') {
      return traceId;
    }

    traceId.timestamp = timestamp;
    traceId.id = parts[2];

    return traceId;
  }

  /**
   * Returns a string representation of the trace ID
   */
  toString() {
    return `${this.version}-${this.timestamp}-${this.id}`;
  }
}

/**
 * Check if XRay tracing is enabled
 */
export function isXRayEnabled() {
  return getEnvVarAsBoolean('COLLABLAND_XRAY_TRACING');
}