package metrics

import (
	"database/sql"
	"strings"
	"sync"

	"github.com/prometheus/client_golang/prometheus"
)

type dbStatsProvider interface {
	Stats() sql.DBStats
}

type dbPoolCollector struct {
	pools sync.Map
	desc  map[string]*prometheus.Desc
}

func newDBPoolCollector() *dbPoolCollector {
	labels := []string{"pod", "pool"}
	return &dbPoolCollector{desc: map[string]*prometheus.Desc{
		"max_open":            prometheus.NewDesc("im_db_pool_max_open_connections", "Configured maximum open database connections.", labels, nil),
		"open":                prometheus.NewDesc("im_db_pool_open_connections", "Current open database connections.", labels, nil),
		"in_use":              prometheus.NewDesc("im_db_pool_in_use_connections", "Current database connections in use.", labels, nil),
		"idle":                prometheus.NewDesc("im_db_pool_idle_connections", "Current idle database connections.", labels, nil),
		"wait":                prometheus.NewDesc("im_db_pool_wait_total", "Total waits for a database connection.", labels, nil),
		"wait_duration":       prometheus.NewDesc("im_db_pool_wait_duration_seconds_total", "Total time spent waiting for database connections.", labels, nil),
		"max_idle_closed":     prometheus.NewDesc("im_db_pool_max_idle_closed_total", "Connections closed because the idle limit was exceeded.", labels, nil),
		"max_lifetime_closed": prometheus.NewDesc("im_db_pool_max_lifetime_closed_total", "Connections closed because their lifetime expired.", labels, nil),
	}}
}

func (c *dbPoolCollector) Describe(ch chan<- *prometheus.Desc) {
	for _, desc := range c.desc {
		ch <- desc
	}
}

func (c *dbPoolCollector) Collect(ch chan<- prometheus.Metric) {
	pod := metricPodName()
	c.pools.Range(func(key, value any) bool {
		pool, ok := key.(string)
		provider, providerOK := value.(dbStatsProvider)
		if !ok || !providerOK {
			return true
		}
		stats := provider.Stats()
		labels := []string{pod, pool}
		ch <- prometheus.MustNewConstMetric(c.desc["max_open"], prometheus.GaugeValue, float64(stats.MaxOpenConnections), labels...)
		ch <- prometheus.MustNewConstMetric(c.desc["open"], prometheus.GaugeValue, float64(stats.OpenConnections), labels...)
		ch <- prometheus.MustNewConstMetric(c.desc["in_use"], prometheus.GaugeValue, float64(stats.InUse), labels...)
		ch <- prometheus.MustNewConstMetric(c.desc["idle"], prometheus.GaugeValue, float64(stats.Idle), labels...)
		ch <- prometheus.MustNewConstMetric(c.desc["wait"], prometheus.CounterValue, float64(stats.WaitCount), labels...)
		ch <- prometheus.MustNewConstMetric(c.desc["wait_duration"], prometheus.CounterValue, stats.WaitDuration.Seconds(), labels...)
		ch <- prometheus.MustNewConstMetric(c.desc["max_idle_closed"], prometheus.CounterValue, float64(stats.MaxIdleClosed), labels...)
		ch <- prometheus.MustNewConstMetric(c.desc["max_lifetime_closed"], prometheus.CounterValue, float64(stats.MaxLifetimeClosed), labels...)
		return true
	})
}

var dbPools = newDBPoolCollector()

// RegisterDBPool exposes database/sql pool pressure for one logical database.
// Registering the same name again replaces the provider, which keeps tests and
// adapter reopen operations safe.
func RegisterDBPool(name string, provider dbStatsProvider) {
	if provider == nil {
		return
	}
	name = strings.TrimSpace(name)
	if name == "" {
		name = "default"
	}
	dbPools.pools.Store(name, provider)
}

func init() {
	Registry.MustRegister(dbPools)
}