package utils

import (
	"context"
	"log/slog"
	"os"

	"go.opentelemetry.io/otel/trace"
)

type JSONHandlerAppender struct {
	*slog.JSONHandler
}

func (h *JSONHandlerAppender) Handle(ctx context.Context, r slog.Record) error {
	spanCtx := trace.SpanContextFromContext(ctx)
	if spanCtx.HasTraceID() {
		rCloned := r.Clone()
		rCloned.AddAttrs(slog.String("trace_id", spanCtx.SpanID().String()))
		return h.JSONHandler.Handle(ctx, rCloned)
	}
	return h.JSONHandler.Handle(ctx, r)
}

func SetDefaultSlog(level string) {
	var l slog.Level
	if err := (&l).UnmarshalText([]byte(level)); err != nil {
		l = slog.LevelInfo
	}
	handler := slog.NewJSONHandler(os.Stderr, &slog.HandlerOptions{AddSource: true, Level: l})
	slog.SetDefault(slog.New(&JSONHandlerAppender{JSONHandler: handler}))
}