import { context, propagation, Span, SpanKind, trace } from '@opentelemetry/api';

import { Interceptor } from '@connectrpc/connect';

export type Carrier = { 
    traceparent?: string;
    tracestate?: string;
}

const tracer = trace.getTracer('connect-rpc');

type InterceptorResponse = Awaited<ReturnType<ReturnType<Interceptor>>>;

const handleSpanEnd = (response: InterceptorResponse, span: Span) => {
    if (!response.stream) {
        span.end();

        return response;
    }

    async function* stream() {
        if (!response.stream) {
            throw new Error('Stream is not defined');
        }

        const activeContext = context.active();

        // eslint-disable-next-line fp/no-loops
        for await (const message of response.message) {
            yield context.with(activeContext, () => message);
        }

        span.end();
    }
    
    return {
        ...response,
        message: stream(),
    };
};

export const serverInterceptor: Interceptor[] = [
    (next) => async (request) => {
        const parentContext = propagation.extract(
            context.active(),
            Object.fromEntries(request.header.entries()),
        );
        const span = tracer.startSpan(
            `${request.service.typeName}/${request.method.name}`,
            {
                kind: SpanKind.SERVER,
            },
            parentContext,
        );

        return context.with(trace.setSpan(parentContext, span), async () => {
            const activeContext = context.active();
            context.bind(activeContext, request);
            context.bind(activeContext, next);
            const response = await next(request);
            context.bind(activeContext, response);
            
            return handleSpanEnd(response, span);
        });
    },
];

export const clientInterceptor: Interceptor[] = [ 
    (next) => async (request) => {
        const span = tracer.startSpan(
            `${request.service.typeName}/${request.method.name}`, {
                kind: SpanKind.CLIENT,
            });
        const parentContext = context.active();

        return context.with(trace.setSpan(parentContext, span), async () => {
            const activeContext = context.active();
            context.bind(activeContext, request);
            const output: Carrier = {};
            propagation.inject(activeContext, output);
            Object.entries(output).forEach(([key, value]) => {
                key && value && request.header.set(key, value);
            });

            context.bind(activeContext, next);
            const response = await next(request);
            context.bind(activeContext, response);

            return handleSpanEnd(response, span);
        });      
    },
];