package httpgateway

import (
	"context"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/tidbcloud/pkgv2/log"
	"go.uber.org/zap"
)

const gatewayComponent = "http_gateway"

var (
	HTTPGatewayHistogram = prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Subsystem: gatewayComponent,
			Name:      "connection_seconds",
			Help:      "Histogram of HTTP gateway connection",
			Buckets:   prometheus.DefBuckets,
		},
		[]string{"status"},
	)

	HTTPGatewaySessionGauge = prometheus.NewDesc(
		prometheus.BuildFQName("", gatewayComponent, "session_number"),
		"HTTP gateway session number",
		[]string{"type", "token"},
		nil)

	HTTPGatewaySessionStatsGauge = prometheus.NewDesc(
		prometheus.BuildFQName("", gatewayComponent, "session_stats"),
		"HTTP gateway session stats",
		[]string{"session", "status"},
		nil)
)

func RegisterMetrics(gateway *HttpGateway) {
	prometheus.MustRegister(HTTPGatewayHistogram)
	prometheus.MustRegister(NewHTTPGatewaySessionCollector(gateway))
}

type HTTPGatewaySessionCollector struct {
	HttpGateway *HttpGateway
}

func NewHTTPGatewaySessionCollector(httpGateway *HttpGateway) *HTTPGatewaySessionCollector {
	return &HTTPGatewaySessionCollector{
		HttpGateway: httpGateway,
	}
}

func (c *HTTPGatewaySessionCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- HTTPGatewaySessionGauge
	ch <- HTTPGatewaySessionStatsGauge
}

type sessionStats struct {
	inUse int
	idle  int
}

func (c *HTTPGatewaySessionCollector) Collect(ch chan<- prometheus.Metric) {
	ctx := context.Background()
	logger := log.WithContextL(ctx)
	logger.Debug("start to collect session metrics")

	txnCount := make(map[string]int)
	c.HttpGateway.sessionStore.RangeTxn(func(session string, txn *TxnSession) bool {
		txnCount[txn.cred.Username]++
		return true
	})
	logger.Debug("txn session count", zap.Int("txncount", len(txnCount)))
	for username, number := range txnCount {
		ch <- prometheus.MustNewConstMetric(HTTPGatewaySessionGauge,
			prometheus.GaugeValue, float64(number), "tx", username)
	}

	nonTxnCount := make(map[string]int)
	c.HttpGateway.sessionStore.RangeNonTxn(func(session string, nonTxn *NonTxnSession) bool {
		nonTxnCount[nonTxn.cred.Username]++
		return true
	})
	logger.Debug("non-txn session count", zap.Int("nonTxnCount", len(nonTxnCount)))
	for username, number := range nonTxnCount {
		ch <- prometheus.MustNewConstMetric(HTTPGatewaySessionGauge,
			prometheus.GaugeValue, float64(number), "non_tx", username)
	}

	sessionStatsMap := make(map[string]sessionStats)
	c.HttpGateway.poolManager.Range(func(token string, dbPool *DBPool) bool {
		userName := dbPool.cred.Username
		stats := dbPool.db.Stats()
		current := sessionStatsMap[userName]
		current.inUse += stats.InUse
		current.idle += stats.Idle
		sessionStatsMap[userName] = current
		return true
	})
	for session, stats := range sessionStatsMap {
		ch <- prometheus.MustNewConstMetric(HTTPGatewaySessionStatsGauge,
			prometheus.GaugeValue, float64(stats.inUse), session, "inuse")
		ch <- prometheus.MustNewConstMetric(HTTPGatewaySessionStatsGauge,
			prometheus.GaugeValue, float64(stats.idle), session, "idle")
	}

	logger.Debug("finish collect session metrics")
}