package kvtrace

import (
	"context"

	"github.com/royalcat/kv"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"
)

var tracer = otel.Tracer("github.com/royalcat/kv/tracer")

type traceSrtore[K, V any] struct {
	kv    kv.Store[K, V]
	attrs []attribute.KeyValue
}

func WrapTracing[K, V any](kv kv.Store[K, V], attrs ...attribute.KeyValue) kv.Store[K, V] {
	return &traceSrtore[K, V]{
		kv:    kv,
		attrs: attrs,
	}
}

// Close implements kv.Store.
func (m *traceSrtore[K, V]) Close(ctx context.Context) error {
	ctx, span := tracer.Start(ctx, "Close", trace.WithAttributes(m.attrs...))
	defer span.End()

	return m.kv.Close(ctx)
}

// Delete implements kv.Store.
func (m *traceSrtore[K, V]) Delete(ctx context.Context, k K) error {
	ctx, span := tracer.Start(ctx, "Delete", trace.WithAttributes(m.attrs...))
	defer span.End()

	return m.kv.Delete(ctx, k)
}

// Get implements kv.Store.
func (m *traceSrtore[K, V]) Get(ctx context.Context, k K) (v V, found bool, err error) {
	ctx, span := tracer.Start(ctx, "Get", trace.WithAttributes(m.attrs...))
	defer span.End()

	return m.kv.Get(ctx, k)
}

// Range implements kv.Store.
func (m *traceSrtore[K, V]) Range(ctx context.Context, iter kv.Iter[K, V]) error {
	ctx, span := tracer.Start(ctx, "Range", trace.WithAttributes(m.attrs...))
	defer span.End()

	count := 0
	iterCount := func(k K, v V) bool {
		count++
		return iter(k, v)
	}

	err := m.kv.Range(ctx, iterCount)
	span.SetAttributes(attribute.Int("count", count))
	return err
}

// RangeWithPrefix implements kv.Store.
func (m *traceSrtore[K, V]) RangeWithPrefix(ctx context.Context, k K, iter kv.Iter[K, V]) error {
	ctx, span := tracer.Start(ctx, "RangeWithPrefix", trace.WithAttributes(m.attrs...))
	defer span.End()

	count := 0
	iterCount := func(k K, v V) bool {
		count++
		return iter(k, v)
	}

	err := m.kv.Range(ctx, iterCount)
	span.SetAttributes(attribute.Int("count", count))
	return err
}

// Set implements kv.Store.
func (m *traceSrtore[K, V]) Set(ctx context.Context, k K, v V) error {
	ctx, span := tracer.Start(ctx, "Set", trace.WithAttributes(m.attrs...))
	defer span.End()

	return m.kv.Set(ctx, k, v)
}

var _ kv.Store[any, any] = (*traceSrtore[any, any])(nil)