package main

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"time"

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

func startAgentMetrics(address string, run *agentRun) (func(), error) {
	if address == "" {
		return func() {}, nil
	}
	listener, err := net.Listen("tcp", address)
	if err != nil {
		return nil, fmt.Errorf("listen for agent metrics on %s: %w", address, err)
	}
	labels := prometheus.Labels{"machine": run.machine}
	registry := prometheus.NewRegistry()
	registry.MustRegister(
		collectors.NewGoCollector(),
		collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}),
		prometheus.NewGaugeFunc(prometheus.GaugeOpts{Name: "im_load_agent_ready_connections", Help: "Current ready WebSocket connections.", ConstLabels: labels}, func() float64 {
			return float64(run.ready.Load())
		}),
		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_publish_attempted_total", Help: "Publish commands attempted.", ConstLabels: labels}, func() float64 {
			return float64(run.attempted.Load())
		}),
		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_publish_accepted_total", Help: "Publish commands acknowledged by the server.", ConstLabels: labels}, func() float64 {
			run.mu.Lock()
			defer run.mu.Unlock()
			return float64(len(run.accepted))
		}),
		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_publish_rejected_total", Help: "Publish commands rejected or timed out.", ConstLabels: labels}, func() float64 {
			return float64(run.rejected.Load())
		}),
		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_publish_403_total", Help: "Publish commands rejected with 403.", ConstLabels: labels}, func() float64 { return float64(run.publish403.Load()) }),
		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_publish_502_total", Help: "Publish commands rejected with 502.", ConstLabels: labels}, func() float64 { return float64(run.publish502.Load()) }),
		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_publish_503_total", Help: "Publish commands rejected with 503.", ConstLabels: labels}, func() float64 { return float64(run.publish503.Load()) }),
		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_publish_ctrl_timeout_total", Help: "Publish commands with control-response timeout.", ConstLabels: labels}, func() float64 { return float64(run.publishTimeout.Load()) }),
		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_publish_socket_error_total", Help: "Publish commands with socket errors.", ConstLabels: labels}, func() float64 { return float64(run.publishSocket.Load()) }),
		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_received_total", Help: "Unique matching data deliveries received.", ConstLabels: labels}, func() float64 {
			received, _, _, _ := run.recorder.metricCounts()
			return float64(received)
		}),
		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_duplicates_total", Help: "Duplicate matching deliveries received.", ConstLabels: labels}, func() float64 {
			_, duplicates, _, _ := run.recorder.metricCounts()
			return float64(duplicates)
		}),
		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_out_of_order_total", Help: "Per-connection topic sequence regressions.", ConstLabels: labels}, func() float64 {
			_, _, outOfOrder, _ := run.recorder.metricCounts()
			return float64(outOfOrder)
		}),
		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_disconnects_total", Help: "Unexpected WebSocket disconnects.", ConstLabels: labels}, func() float64 {
			_, _, _, disconnects := run.recorder.metricCounts()
			return float64(disconnects)
		}),
	)
	mux := http.NewServeMux()
	mux.Handle("/metrics", promhttp.HandlerFor(registry, promhttp.HandlerOpts{}))
	server := &http.Server{Handler: mux, ReadHeaderTimeout: 5 * time.Second}
	go func() { _ = server.Serve(listener) }()
	return func() {
		ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
		defer cancel()
		_ = server.Shutdown(ctx)
	}, nil
}