package metrics

import (
	"github.com/prometheus/client_golang/prometheus"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
)

type DBPoolSnapshotProvider interface {
	Snapshot() coredb.DBPoolSnapshot
}

type dbPoolCollector struct {
	provider DBPoolSnapshotProvider

	openConnections    *prometheus.Desc
	inUse              *prometheus.Desc
	idle               *prometheus.Desc
	waitCount          *prometheus.Desc
	waitDurationSecond *prometheus.Desc
}

func RegisterDBPool(reg prometheus.Registerer, provider DBPoolSnapshotProvider) error {
	if provider == nil {
		return nil
	}
	return reg.Register(newDBPoolCollector(provider))
}

func newDBPoolCollector(provider DBPoolSnapshotProvider) *dbPoolCollector {
	return &dbPoolCollector{
		provider: provider,
		openConnections: prometheus.NewDesc(
			"agentteam_db_pool_open_connections",
			"Current number of open database connections.",
			nil,
			nil,
		),
		inUse: prometheus.NewDesc(
			"agentteam_db_pool_in_use",
			"Current number of database connections in use.",
			nil,
			nil,
		),
		idle: prometheus.NewDesc(
			"agentteam_db_pool_idle",
			"Current number of idle database connections.",
			nil,
			nil,
		),
		waitCount: prometheus.NewDesc(
			"agentteam_db_pool_wait_count_total",
			"Total number of database connection waits.",
			nil,
			nil,
		),
		waitDurationSecond: prometheus.NewDesc(
			"agentteam_db_pool_wait_duration_seconds_total",
			"Total duration spent waiting for database connections.",
			nil,
			nil,
		),
	}
}

func (collector *dbPoolCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- collector.openConnections
	ch <- collector.inUse
	ch <- collector.idle
	ch <- collector.waitCount
	ch <- collector.waitDurationSecond
}

func (collector *dbPoolCollector) Collect(ch chan<- prometheus.Metric) {
	snapshot := collector.provider.Snapshot()
	stats := snapshot.Stats

	ch <- prometheus.MustNewConstMetric(collector.openConnections, prometheus.GaugeValue, float64(stats.OpenConnections))
	ch <- prometheus.MustNewConstMetric(collector.inUse, prometheus.GaugeValue, float64(stats.InUse))
	ch <- prometheus.MustNewConstMetric(collector.idle, prometheus.GaugeValue, float64(stats.Idle))
	ch <- prometheus.MustNewConstMetric(collector.waitCount, prometheus.CounterValue, float64(stats.WaitCount))
	ch <- prometheus.MustNewConstMetric(collector.waitDurationSecond, prometheus.CounterValue, stats.WaitDuration.Seconds())
}