package logging

import (
	"context"
	"log/slog"
)

type contextFieldsKey struct{}

// ContextFields contains request metadata attached to logs emitted with a
// request context.
type ContextFields struct {
	RequestID    string
	TraceID      string
	SpanID       string
	ParentSpanID string
	DeviceID     string
	ClientIP     string
	HTTPMethod   string
	HTTPPath     string
}

func WithContextFields(ctx context.Context, fields *ContextFields) context.Context {
	return context.WithValue(ctx, contextFieldsKey{}, fields)
}

func ContextFieldsFrom(ctx context.Context) *ContextFields {
	if ctx == nil {
		return nil
	}
	fields, _ := ctx.Value(contextFieldsKey{}).(*ContextFields)
	return fields
}

type contextHandler struct{ handler slog.Handler }

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

func (h *contextHandler) Handle(ctx context.Context, record slog.Record) error {
	if fields := ContextFieldsFrom(ctx); fields != nil {
		record.AddAttrs(
			slog.String("request_id", fields.RequestID),
			slog.String("trace_id", fields.TraceID),
			slog.String("device_id", fields.DeviceID),
			slog.String("client_ip", fields.ClientIP),
			slog.String("http_method", fields.HTTPMethod),
			slog.String("http_path", fields.HTTPPath),
		)
		if fields.SpanID != "" {
			record.AddAttrs(slog.String("span_id", fields.SpanID))
		}
		if fields.ParentSpanID != "" {
			record.AddAttrs(slog.String("parent_span_id", fields.ParentSpanID))
		}
	}
	return h.handler.Handle(ctx, record)
}

func (h *contextHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	return &contextHandler{handler: h.handler.WithAttrs(attrs)}
}

func (h *contextHandler) WithGroup(name string) slog.Handler {
	return &contextHandler{handler: h.handler.WithGroup(name)}
}