package kk_stage

import (
	"context"
	"go.opentelemetry.io/otel/trace"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"log/slog"

	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/sdk/resource"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	"go.opentelemetry.io/otel/semconv/v1.21.0"
)

var tracerProviderShutdown func(context.Context) error
var grpcConn *grpc.ClientConn

type Tracer struct {
	Tracer  trace.Tracer
	Context context.Context
}

func NewTracer(serviceName string, ctx context.Context) *Tracer {
	tracer := &Tracer{
		Tracer:  otel.Tracer(serviceName),
		Context: ctx,
	}
	return tracer
}

// Initialize a gRPC connection to be used by both the tracer and meter providers.
func initGrpcConn(stage *Stage, grpcTarget string) {
	conn, err := grpc.NewClient(grpcTarget,
		grpc.WithTransportCredentials(insecure.NewCredentials()),
		grpc.WithBlock(),
	)
	if err != nil {
		slog.Error("failed to create gRPC connection to collector", NewLog(stage).Error(err).Args()...)
		panic(nil)
	} else {
		grpcConn = conn
	}
}

// Initializes an OTLP exporter, and configures the corresponding trace
// provider.
func initTracerProvider(stage *Stage, serviceName string) {
	ctx := context.Background()
	// Set up a trace exporter
	traceExporter, err := otlptracegrpc.New(ctx, otlptracegrpc.WithGRPCConn(grpcConn))
	if err != nil {
		slog.Error("failed to create trace exporter", NewLog(stage).Error(err).Args()...)
		panic(nil)
	}

	res, err := resource.New(ctx,
		resource.WithAttributes(semconv.ServiceName(serviceName)))
	if err != nil {
		slog.Error("failed to create resource", NewLog(stage).Error(err).Args()...)
		panic(nil)
	}

	// Register the trace exporter with a TracerProvider, using a batch
	// span processor to aggregate spans before export.
	bsp := sdktrace.NewBatchSpanProcessor(traceExporter)
	tracerProvider := sdktrace.NewTracerProvider(
		sdktrace.WithSampler(sdktrace.AlwaysSample()),
		sdktrace.WithResource(res),
		sdktrace.WithSpanProcessor(bsp),
	)
	otel.SetTracerProvider(tracerProvider)

	// set global propagator to traceContext (the default is no-op).
	otel.SetTextMapPropagator(propagation.TraceContext{})

	// Shutdown will flush any remaining spans and shut down the exporter.
	tracerProviderShutdown = tracerProvider.Shutdown

}

var EnableTracer = false

func InitTracer(stage *Stage, enabled bool, serviceName string, grpcTarget string) {
	EnableTracer = enabled
	if EnableTracer {
		initGrpcConn(stage, grpcTarget)
		initTracerProvider(stage, serviceName)
	}
}

func CloseTracer(stage *Stage) {
	if tracerProviderShutdown != nil {
		err := tracerProviderShutdown(context.Background())
		if err != nil {
			slog.Error("failed to shutdown tracer provider", NewLog(stage).Error(err).Args()...)
		}
	}
	if grpcConn != nil {
		err := grpcConn.Close()
		if err != nil {
			slog.Error("failed to close grpc connection", NewLog(stage).Error(err).Args()...)
		}
	}
}