// MIT License

// Copyright (c) 2020 GameChanger Media

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

import { SpanOptions } from "dd-trace";
import DDTags from "dd-trace/ext/tags";
import env from "@server/env";
import tracer from "./tracer";
import * as Tracing from "./tracer";

type DDTag = typeof DDTags[keyof typeof DDTags];

type Tags = {
  [tag in DDTag]?: any;
} & {
  [key: string]: any;
};

interface Constructor {
  new (...args: any[]): any;
}

interface TraceConfig {
  className?: string;
  methodName?: string;
  serviceName?: string;
  spanName?: string;
  resourceName?: string;
  isRoot?: boolean;
  /** Cause the span to show up in trace search and analytics */
  makeSearchable?: boolean;
  tags?: Tags;
}

/**
 * This decorator will cause an individual function to be traced by the APM.
 *
 * @param config Optional configuration for the span that will be created for this trace.
 */
export const traceFunction = (config: TraceConfig) => <
  F extends (...args: any[]) => any,
  P extends Parameters<F>,
  R extends ReturnType<F>
>(
  target: F
): F =>
  env.ENVIRONMENT === "test"
    ? target
    : (function wrapperFn(this: any, ...args: P): R {
        const { className, methodName = target.name, tags } = config;
        const childOf = config.isRoot
          ? undefined
          : tracer.scope().active() || undefined;

        const spanName = config.spanName || className || "DEFAULT_SPAN_NAME";

        const resourceName = config.resourceName
          ? config.resourceName
          : methodName;
        const spanOptions: SpanOptions = {
          childOf,
          tags: {
            [DDTags.RESOURCE_NAME]: resourceName,
            ...tags,
          },
        };

        const span = tracer.startSpan(spanName, spanOptions);

        if (!span) {
          return target.call(this, ...args);
        }

        if (config.serviceName) {
          span.setTag(
            DDTags.SERVICE_NAME,
            `${env.DD_SERVICE}-${config.serviceName}`
          );
        }

        if (config.makeSearchable) {
          span.setTag(DDTags.ANALYTICS, true);
        }

        // The callback fn needs to be wrapped in an arrow fn as the activate fn clobbers `this`
        return tracer.scope().activate(span, () => {
          const output = target.call(this, ...args);

          if (output && typeof output.then === "function") {
            output
              .catch((error: Error) => {
                Tracing.setError(error, span);
              })
              .finally(() => {
                span.finish();
              });
          } else {
            span.finish();
          }

          return output;
        });
      } as F);

const traceMethod = (config?: TraceConfig) =>
  function <R, A extends any[], F extends (...args: A) => R>(
    target: any,
    _propertyKey: string,
    descriptor: PropertyDescriptor
  ): TypedPropertyDescriptor<F> {
    const wrappedFn = descriptor.value;

    if (wrappedFn) {
      const className = target.name || target.constructor.name; // target.name is needed if the target is the constructor itself
      const methodName = wrappedFn.name;
      descriptor.value = traceFunction({ ...config, className, methodName })(
        wrappedFn
      );
    }

    return descriptor;
  };

const traceClass = (config?: TraceConfig) =>
  function <T extends Constructor>(constructor: T): void {
    const protoKeys = Reflect.ownKeys(constructor.prototype);
    protoKeys.forEach((key) => {
      if (key === "constructor") {
        return;
      }

      const descriptor = Object.getOwnPropertyDescriptor(
        constructor.prototype,
        key
      );

      // eslint-disable-next-line no-undef
      if (typeof key === "string" && typeof descriptor?.value === "function") {
        Object.defineProperty(
          constructor.prototype,
          key,
          traceMethod(config)(constructor, key, descriptor)
        );
      }
    });

    const staticKeys = Reflect.ownKeys(constructor);
    staticKeys.forEach((key) => {
      const descriptor = Object.getOwnPropertyDescriptor(constructor, key);

      // eslint-disable-next-line no-undef
      if (typeof key === "string" && typeof descriptor?.value === "function") {
        Object.defineProperty(
          constructor,
          key,
          traceMethod(config)(constructor, key, descriptor)
        );
      }
    });
  };

/**
 * This decorator will cause the methods of a class, or an individual method, to be traced by the APM.
 *
 * @param config Optional configuration for the span that will be created for this trace.
 */
// Going to rely on inferrence do its thing for this function
// eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
export function trace(config?: TraceConfig) {
  function traceDecorator(target: Constructor): void;
  function traceDecorator<T>(
    target: Record<string, any>,
    propertyKey: string | symbol,
    descriptor: TypedPropertyDescriptor<T>
  ): void;
  function traceDecorator(
    a: Constructor | Record<string, any>,
    b?: any,
    c?: any
  ): void {
    if (typeof a === "function") {
      // Need to cast as there is no safe runtime way to check if a function is a constructor
      traceClass(config)(a as Constructor);
    } else {
      traceMethod(config)(a, b, c);
    }
  }

  return traceDecorator;
}
