package agent

import (
	"context"
	"log/slog"
	"sync"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/logz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/syncz"
)

// agentKeyHandler wraps a slog.Handler to add agent key field if agent key is available.
type agentKeyHandler struct {
	mu       sync.Mutex
	delegate slog.Handler
	agentKey *syncz.ValueHolder[api.AgentKey] // if nil, agent key has been added to the delegate.
}

func (h *agentKeyHandler) Enabled(ctx context.Context, level slog.Level) bool {
	h.mu.Lock()
	d := h.delegate
	h.mu.Unlock()
	return d.Enabled(ctx, level)
}

func (h *agentKeyHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	d, agentKey := h.tryToGetAgentKey()
	d = d.WithAttrs(attrs)
	if agentKey == nil { // d already has agent key attr
		return d
	}
	return &agentKeyHandler{
		delegate: d,
		agentKey: h.agentKey,
	}
}

func (h *agentKeyHandler) WithGroup(name string) slog.Handler {
	d, agentKey := h.tryToGetAgentKey()
	d = d.WithGroup(name)
	if agentKey == nil { // d already has agent key attr
		return d
	}
	return &agentKeyHandler{
		delegate: d,
		agentKey: h.agentKey,
	}
}

func (h *agentKeyHandler) Handle(ctx context.Context, r slog.Record) error {
	d, _ := h.tryToGetAgentKey()
	return d.Handle(ctx, r)
}

func (h *agentKeyHandler) tryToGetAgentKey() (slog.Handler, *syncz.ValueHolder[api.AgentKey]) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if h.agentKey == nil {
		return h.delegate, nil
	}
	agentKey, ok := h.agentKey.TryGet()
	if ok {
		h.delegate = h.delegate.WithAttrs([]slog.Attr{logz.AgentKey(agentKey)})
		h.agentKey = nil
	}
	return h.delegate, h.agentKey
}
