/*!
 * Copyright 2024 Google LLC. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import {
  ATTR_OTEL_SCOPE_NAME,
  ATTR_OTEL_SCOPE_VERSION,
} from '@opentelemetry/semantic-conventions';

import {
  Span,
  SpanStatusCode,
  context,
  trace,
  INVALID_SPAN_CONTEXT,
  ROOT_CONTEXT,
  SpanAttributes,
  TimeInput,
  TracerProvider,
  Link,
  Exception,
  SpanContext,
  SpanStatus,
  SpanKind,
} from '@opentelemetry/api';

const optedInPII: boolean =
  process.env.SPANNER_ENABLE_EXTENDED_TRACING === 'true';

interface SQLStatement {
  sql: string;
}

/*
 * ObservabilityOptions defines the configuration to be used for Spanner OpenTelemetry Traces.
 * @property [tracerProvider] Sets the TracerProvider to be used for traces,
 * Global TracerProvider will be used as a fallback.
 * @property [enableExtendedTracing] Sets whether to enable extended OpenTelemetry tracing. Enabling this option will add the
 * following additional attributes to the traces that are generated by the client
 * db.statement: Contains the SQL statement that is being executed.
 * Alternatively, you could set environment variable `SPANNER_ENABLE_EXTENDED_TRACING=true`.
 */
interface ObservabilityOptions {
  tracerProvider: TracerProvider;
  enableExtendedTracing?: boolean;
  enableEndToEndTracing?: boolean;
}

export type {ObservabilityOptions};
export type {Span};

const TRACER_NAME = 'cloud.google.com/nodejs/spanner';
const TRACER_VERSION = require('../../package.json').version;

export {TRACER_NAME, TRACER_VERSION}; // Only exported for testing.

/**
 * getTracer fetches the tracer from the provided tracerProvider.
 * @param {TracerProvider} [tracerProvider] optional custom tracer provider
 * to use for fetching the tracer. If not provided, the global provider will be used.
 *
 * @returns {Tracer} The tracer instance associated with the provided or global provider.
 */
export function getTracer(tracerProvider?: TracerProvider) {
  if (tracerProvider) {
    return tracerProvider.getTracer(TRACER_NAME, TRACER_VERSION);
  }
  // Otherwise use the global tracer.
  return trace.getTracer(TRACER_NAME, TRACER_VERSION);
}

interface traceConfig {
  sql?: string | SQLStatement;
  tableName?: string;
  dbName?: string;
  transactionTag?: string | null;
  requestTag?: string | null;
  opts?: ObservabilityOptions;
}

const SPAN_NAMESPACE_PREFIX = 'CloudSpanner'; // TODO: discuss & standardize this prefix.
export {SPAN_NAMESPACE_PREFIX, traceConfig};

const {
  AsyncHooksContextManager,
} = require('@opentelemetry/context-async-hooks');

/*
 * This function ensures that async/await works correctly by
 * checking if context.active() returns an invalid/unset context
 * and if so, sets a global AsyncHooksContextManager otherwise
 * spans resulting from async/await invocations won't be correctly
 * associated in their respective hierarchies.
 */
function ensureInitialContextManagerSet() {
  if (!context['_contextManager'] || context.active() === ROOT_CONTEXT) {
    // If no context manager is currently set, or if the active context is the ROOT_CONTEXT,
    // trace context propagation cannot
    // function correctly with async/await for OpenTelemetry
    // See {@link https://opentelemetry.io/docs/languages/js/context/#active-context}
    context.disable(); // Disable any prior contextManager.
    const contextManager = new AsyncHooksContextManager();
    contextManager.enable();
    context.setGlobalContextManager(contextManager);
  }
}

export {ensureInitialContextManagerSet};

/**
 * startTrace begins an active span in the current active context
 * and passes it back to the set callback function. Each span will
 * be prefixed with "cloud.google.com/nodejs/spanner". It is the
 * responsibility of the caller to invoke [span.end] when finished tracing.
 *
 * @returns {Span} The created span.
 */
export function startTrace<T>(
  spanNameSuffix: string,
  config: traceConfig | undefined,
  cb: (span: Span) => T,
): T {
  if (!config) {
    config = {} as traceConfig;
  }

  return getTracer(config.opts?.tracerProvider).startActiveSpan(
    SPAN_NAMESPACE_PREFIX + '.' + spanNameSuffix,
    {kind: SpanKind.CLIENT},
    span => {
      span.setAttribute(ATTR_OTEL_SCOPE_NAME, TRACER_NAME);
      span.setAttribute(ATTR_OTEL_SCOPE_VERSION, TRACER_VERSION);
      span.setAttribute('gcp.client.service', 'spanner');
      span.setAttribute('gcp.client.version', TRACER_VERSION);
      span.setAttribute('gcp.client.repo', 'googleapis/nodejs-spanner');

      if (config.tableName) {
        span.setAttribute('db.sql.table', config.tableName);
      }
      if (config.dbName) {
        span.setAttribute('db.name', config.dbName);
      }
      if (config.requestTag) {
        span.setAttribute('request.tag', config.requestTag);
      }
      if (config.transactionTag) {
        span.setAttribute('transaction.tag', config.transactionTag);
      }

      const allowExtendedTracing =
        optedInPII || config.opts?.enableExtendedTracing;
      if (config.sql && allowExtendedTracing) {
        const sql = config.sql;
        if (typeof sql === 'string') {
          span.setAttribute('db.statement', sql as string);
        } else {
          const stmt = sql as SQLStatement;
          span.setAttribute('db.statement', stmt.sql);
        }
      }

      // If at all the invoked function throws an exception,
      // record the exception and then end this span.
      try {
        return cb(span);
      } catch (e) {
        setSpanErrorAndException(span, e as Error);
        span.end();
        // Finally re-throw the exception.
        throw e;
      }
    },
  );
}

/**
 * Sets the span status with err, if non-null onto the span with
 * status.code=ERROR and the message of err.toString()
 *
 * @returns {boolean} to signify if the status was set.
 */
export function setSpanError(span: Span, err: Error | String): boolean {
  if (!err || !span) {
    return false;
  }

  let message = '';
  if (typeof err === 'object' && 'message' in err) {
    message = err.message as string;
  } else {
    message = err.toString();
  }
  span.setStatus({
    code: SpanStatusCode.ERROR,
    message: message,
  });
  return true;
}

/**
 * Sets err, if non-null onto the span with
 * status.code=ERROR and the message of err.toString()
 * as well as recording an exception on the span.
 * @param {Span} [span] the subject span
 * @param {Error} [err] the error whose message to use to record
 * the span error and the span exception.
 *
 * @returns {boolean} to signify if the status and exception were set.
 */
export function setSpanErrorAndException(
  span: Span,
  err: Error | String,
): boolean {
  if (setSpanError(span, err)) {
    span.recordException(err as Error);
    return true;
  }
  return false;
}

/**
 * getActiveOrNoopSpan queries the global tracer for the currently active
 * span and returns it, otherwise if there is no active span available, it'll
 * simply create a NoopSpan. This is important in the cases where we don't
 * want to create a new span, such as in sensitive and frequently called code
 * for which the new spans would be too many and thus pollute the trace,
 * but yet we'd like to record an important annotation.
 *
 * @returns {Span} the non-null span.
 */
export function getActiveOrNoopSpan(): Span {
  const span = trace.getActiveSpan();
  if (span) {
    return span;
  }
  return new noopSpan();
}

/**
 * noopSpan is a pass-through Span that does nothing and shall not
 * be exported, nor added into any context. It serves as a placeholder
 * to allow calls in sensitive areas like sessionPools to transparently
 * add attributes to spans without lots of ugly null checks.
 *
 * It exists because OpenTelemetry-JS does not seem to export the NoopSpan.
 */
class noopSpan implements Span {
  constructor() {}

  spanContext(): SpanContext {
    return INVALID_SPAN_CONTEXT;
  }

  setAttribute(key: string, value: unknown): this {
    return this;
  }

  setAttributes(attributes: SpanAttributes): this {
    return this;
  }

  addEvent(name: string, attributes?: SpanAttributes): this {
    return this;
  }

  addLink(link: Link): this {
    return this;
  }

  addLinks(links: Link[]): this {
    return this;
  }

  setStatus(status: SpanStatus): this {
    return this;
  }

  end(endTime?: TimeInput): void {}

  isRecording(): boolean {
    return false;
  }

  recordException(exc: Exception, timeAt?: TimeInput): void {}

  updateName(name: string): this {
    return this;
  }
}
