package logging

import (
	"context"
	"errors"
	"log/slog"

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

type contextHandler struct {
	handler       slog.Handler
	hasDomain     bool
	hasStacktrace bool
	flush         func() error
}

func WrapHandler(handler slog.Handler, flush ...func() error) slog.Handler {
	if handler == nil {
		return nil
	}
	var flushFunc func() error
	if len(flush) > 0 {
		flushFunc = flush[0]
	}
	return contextHandler{handler: handler, flush: flushFunc}
}

func (handler contextHandler) Enabled(ctx context.Context, level slog.Level) bool {
	return handler.handler.Enabled(ctx, level)
}

func (handler contextHandler) Handle(ctx context.Context, record slog.Record) error {
	hasDomain := handler.hasDomain || recordHasAttr(record, "domain")
	if !hasDomain {
		record.AddAttrs(slog.String("domain", "system"))
	}
	if requestID := RequestID(ctx); requestID != "" {
		record.AddAttrs(slog.String("request_id", requestID))
	}
	if spanContext := trace.SpanContextFromContext(ctx); spanContext.IsValid() {
		record.AddAttrs(
			slog.String("trace_id", spanContext.TraceID().String()),
			slog.String("span_id", spanContext.SpanID().String()),
		)
	}
	if record.Level >= slog.LevelError && !handler.hasStacktrace && !recordHasAttr(record, "stacktrace") {
		record.AddAttrs(slog.String("stacktrace", boundedStacktrace()))
	}
	err := handler.handler.Handle(ctx, record)
	if record.Level >= slog.LevelError && handler.flush != nil {
		err = errors.Join(err, handler.flush())
	}
	return err
}

func (handler contextHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	return contextHandler{
		handler:       handler.handler.WithAttrs(attrs),
		hasDomain:     handler.hasDomain || attrsHaveKey(attrs, "domain"),
		hasStacktrace: handler.hasStacktrace || attrsHaveKey(attrs, "stacktrace"),
		flush:         handler.flush,
	}
}

func (handler contextHandler) WithGroup(name string) slog.Handler {
	return contextHandler{handler: handler.handler.WithGroup(name), hasDomain: handler.hasDomain, hasStacktrace: handler.hasStacktrace, flush: handler.flush}
}

func recordHasAttr(record slog.Record, key string) bool {
	found := false
	record.Attrs(func(attr slog.Attr) bool {
		if attr.Key == key {
			found = true
			return false
		}
		return true
	})
	return found
}

func attrsHaveKey(attrs []slog.Attr, key string) bool {
	for _, attr := range attrs {
		if attr.Key == key {
			return true
		}
	}
	return false
}