import { context } from "@opentelemetry/api";
import { registerInstrumentations } from '@opentelemetry/instrumentation';
import { AwsLambdaInstrumentation } from '@opentelemetry/instrumentation-aws-lambda';
import { AlwaysOnSampler } from "@opentelemetry/sdk-trace-base";
import { NodeTracerProvider } from '@opentelemetry/sdk-trace-node';
import { B3_INFO_KEY, UoaB3Propagator } from "./UoaB3Propagator";

const provider = new NodeTracerProvider({sampler: new AlwaysOnSampler()});
let infoHeader: string | undefined;

export function initializeTracing() {
    provider.register({
        propagator: new UoaB3Propagator()
    });

    registerInstrumentations({
        instrumentations: [
            new AwsLambdaInstrumentation({
                requestHook: (span, {event, context}) => {
                    span.setAttribute('faas.name', context.functionName);
                    infoHeader = undefined; //reset header value in case lambda execution environment maintained since last invocation
                },
                responseHook: (span, {err, res}) => {
                    if (err instanceof Error) span.setAttribute('faas.error', err.message);
                    if (res) span.setAttribute('faas.res', res);
                }
            })
        ]
    });
}

export function getTraceInfoHeader(): string {
    if (infoHeader) {
        return infoHeader;
    } else {
        return <string>context.active().getValue(B3_INFO_KEY);
    }
}

export function setTraceInfoHeader(info: string) {
    infoHeader = info;
}