package postgres

import (
	"context"
	"strings"
	"time"

	"route256/loms/internal/infra/grpc/metrics"

	"github.com/jackc/pgx/v5"
)

type startKey struct{}

type promTracer struct{}

func (promTracer) TraceQueryStart(ctx context.Context, _ *pgx.Conn, _ pgx.TraceQueryStartData) context.Context {
	return context.WithValue(ctx, startKey{}, time.Now())
}

func (promTracer) TraceQueryEnd(ctx context.Context, _ *pgx.Conn, data pgx.TraceQueryEndData) {
	start, _ := ctx.Value(startKey{}).(time.Time)

	sql := strings.TrimSpace(strings.ToLower(data.CommandTag.String()))
	parts := strings.SplitN(sql, " ", 2)
	category := "unknown"
	if len(parts) > 0 {
		category = parts[0]
	}

	errLabel := "none"
	if data.Err != nil {
		errLabel = "error"
	}

	metrics.IncDBQueryCount(category, errLabel)
	metrics.StoreDBQueryDuration(category, errLabel, time.Since(start))
}