/**
 * Class extended from OpenTelemetry B3MultiPropagator so that we can:
 * 1. Propagate & log the X-B3-Info header
 * 2. Propagate the correct X-B3-ParentSpanId
 *
 * See {@link https://github.com/open-telemetry/opentelemetry-js/blob/main/packages/opentelemetry-propagator-b3/src/B3MultiPropagator.ts}
 */
import {
    Context,
    createContextKey,
    isSpanContextValid,
    TextMapGetter,
    TextMapPropagator,
    TextMapSetter,
    trace,
    TraceFlags
} from '@opentelemetry/api';
import { isTracingSuppressed } from '@opentelemetry/core';
import crypto from 'crypto';
import { getTraceInfoHeader } from "./tracing";


export const X_B3_TRACE_ID = 'x-b3-traceid';
export const X_B3_SPAN_ID = 'x-b3-spanid';
export const X_B3_SAMPLED = 'x-b3-sampled';
export const X_B3_PARENT_SPAN_ID = 'x-b3-parentspanid';
export const X_B3_FLAGS = 'x-b3-flags';
export const X_B3_INFO = 'x-b3-info';
export const B3_DEBUG_FLAG_KEY = createContextKey(
    'B3 Debug Flag'
);
export const B3_INFO_KEY = createContextKey(
    'B3 Info Header'
);
export const B3_TRACE_ID_LENGTH_KEY = createContextKey(
    'B3 TraceId Header Length'
);
export const B3_PARENT_SPAN_ID_KEY = createContextKey(
    'B3 Parent Span Id'
);

const VALID_SAMPLED_VALUES = new Set([true, 'true', 'True', '1', 1]);
const VALID_UNSAMPLED_VALUES = new Set([false, 'false', 'False', '0', 0]);

function parseHeader(header: unknown) {
    return Array.isArray(header) ? header[0] : header;
}

function getHeaderValue(carrier: unknown, getter: TextMapGetter, key: string) {
    const header = getter.get(carrier, key);
    return parseHeader(header);
}

function getTraceIdLength(carrier: unknown, getter: TextMapGetter): number {
    const traceId = getHeaderValue(carrier, getter, X_B3_TRACE_ID);
    if (typeof traceId === 'string') {
        if (traceId.length <= 16 || traceId.length > 32) {
            return 16;
        } else {
            return 32;
        }
    }
    return 16;
}

function getTraceId(carrier: unknown, getter: TextMapGetter): string {
    const traceId = getHeaderValue(carrier, getter, X_B3_TRACE_ID);
    if (typeof traceId === 'string') {
        return traceId.padStart(32, '0');
    }
    return '';
}

function getSpanId(carrier: unknown, getter: TextMapGetter): string {
    const spanId = getHeaderValue(carrier, getter, X_B3_SPAN_ID);
    if (typeof spanId === 'string') {
        return spanId;
    }
    return '';
}

function getParentSpanId(carrier: unknown, getter: TextMapGetter): string {
    const parentSpanId = getHeaderValue(carrier, getter, X_B3_PARENT_SPAN_ID);
    if (typeof parentSpanId === 'string') {
        return parentSpanId;
    }
    return '';
}

function getDebug(carrier: unknown, getter: TextMapGetter): string | undefined {
    const debug = getHeaderValue(carrier, getter, X_B3_FLAGS);
    return debug === '1' ? '1' : undefined;
}

function getTraceFlags(
    carrier: unknown,
    getter: TextMapGetter
): TraceFlags | undefined {
    const traceFlags = getHeaderValue(carrier, getter, X_B3_SAMPLED);
    const debug = getDebug(carrier, getter);
    if (debug === '1' || VALID_SAMPLED_VALUES.has(traceFlags)) {
        return TraceFlags.SAMPLED;
    }
    if (traceFlags === undefined || VALID_UNSAMPLED_VALUES.has(traceFlags)) {
        return TraceFlags.NONE;
    }
    // This indicates to isValidSampledValue that this is not valid
    return;
}

function getInfo(carrier: unknown, getter: TextMapGetter): string | undefined {
    const info = getHeaderValue(carrier, getter, X_B3_INFO);
    if (typeof info === 'string') {
        return info;
    }
    return undefined;
}

export class UoaB3Propagator implements TextMapPropagator {
    inject(context: Context, carrier: any, setter: TextMapSetter<any>): void {
        const spanContext = trace.getSpanContext(context);
        if (
            !spanContext ||
            !isSpanContextValid(spanContext) ||
            isTracingSuppressed(context)
        )
            return;

        const debug = context.getValue(B3_DEBUG_FLAG_KEY);
        const traceIdLength: number = <number>context.getValue(B3_TRACE_ID_LENGTH_KEY);
        setter.set(carrier, X_B3_TRACE_ID, spanContext.traceId.slice(spanContext.traceId.length - traceIdLength));
        setter.set(carrier, X_B3_SPAN_ID, crypto.randomBytes(8).toString('hex'));
        setter.set(carrier, X_B3_PARENT_SPAN_ID, spanContext.spanId);
        const info = context.getValue(B3_INFO_KEY);
        if (getTraceInfoHeader()) {
            setter.set(carrier, X_B3_INFO, getTraceInfoHeader());
        } else if (info && typeof info === 'string') {
            setter.set(carrier, X_B3_INFO, info.toString());
        }
        // According to the B3 spec, if the debug flag is set,
        // the sampled flag shouldn't be propagated as well.
        if (debug === '1') {
            setter.set(carrier, X_B3_FLAGS, debug);
        } else if (spanContext.traceFlags !== undefined) {
            // We set the header only if there is an existing sampling decision.
            // Otherwise we will omit it => Absent.
            setter.set(
                carrier,
                X_B3_SAMPLED,
                (TraceFlags.SAMPLED & spanContext.traceFlags) === TraceFlags.SAMPLED
                    ? '1'
                    : '0'
            );
        }
    }

    extract(context: Context, carrier: any, getter: TextMapGetter<any>): Context {
        const traceIdLength = getTraceIdLength(carrier, getter);
        context = context.setValue(B3_TRACE_ID_LENGTH_KEY, traceIdLength);

        const traceId = getTraceId(carrier, getter);
        const spanId = getSpanId(carrier, getter);
        const parentSpanId = getParentSpanId(carrier, getter);
        context = context.setValue(B3_PARENT_SPAN_ID_KEY, parentSpanId);

        const traceFlags = getTraceFlags(carrier, getter) as TraceFlags;
        const debug = getDebug(carrier, getter);
        const info = getInfo(carrier, getter);

        context = context.setValue(B3_DEBUG_FLAG_KEY, debug);
        if (info) {
            context = context.setValue(B3_INFO_KEY, info);
        }
        return trace.setSpanContext(context, {
            traceId,
            spanId,
            isRemote: true,
            traceFlags,
        });
    }

    fields(): string[] {
        return [
            X_B3_TRACE_ID,
            X_B3_SPAN_ID,
            X_B3_FLAGS,
            X_B3_SAMPLED,
            X_B3_PARENT_SPAN_ID,
            X_B3_INFO
        ];
    }

}