package metrics

import (
	"log/slog"
	"net"
	"sync"
	"time"

	"golang.getoutline.org/tunnel-server/ipinfo"
	outline_prometheus "golang.getoutline.org/tunnel-server/prometheus"
	"golang.getoutline.org/tunnel-server/service"
	svcmetrics "golang.getoutline.org/tunnel-server/service/metrics"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/collectors"
)

// ServerMetrics tracks server-level Prometheus metrics.
type ServerMetrics struct {
	buildInfo         *prometheus.GaugeVec
	accessKeys        prometheus.Gauge
	ports             prometheus.Gauge
	addedNatEntries   prometheus.Counter
	removedNatEntries prometheus.Counter
}

var _ prometheus.Collector = (*ServerMetrics)(nil)
var _ service.NATMetrics = (*ServerMetrics)(nil)

func NewServerMetrics() *ServerMetrics {
	return &ServerMetrics{
		buildInfo: prometheus.NewGaugeVec(prometheus.GaugeOpts{
			Name: "build_info",
			Help: "Information on the shroud build",
		}, []string{"version"}),
		accessKeys: prometheus.NewGauge(prometheus.GaugeOpts{
			Name: "keys",
			Help: "Count of access keys",
		}),
		ports: prometheus.NewGauge(prometheus.GaugeOpts{
			Name: "ports",
			Help: "Count of open ports",
		}),
		addedNatEntries: prometheus.NewCounter(prometheus.CounterOpts{
			Subsystem: "udp",
			Name:      "nat_entries_added",
			Help:      "Entries added to the UDP NAT table",
		}),
		removedNatEntries: prometheus.NewCounter(prometheus.CounterOpts{
			Subsystem: "udp",
			Name:      "nat_entries_removed",
			Help:      "Entries removed from the UDP NAT table",
		}),
	}
}

func (m *ServerMetrics) Describe(ch chan<- *prometheus.Desc) {
	m.buildInfo.Describe(ch)
	m.accessKeys.Describe(ch)
	m.ports.Describe(ch)
	m.addedNatEntries.Describe(ch)
	m.removedNatEntries.Describe(ch)
}

func (m *ServerMetrics) Collect(ch chan<- prometheus.Metric) {
	m.buildInfo.Collect(ch)
	m.accessKeys.Collect(ch)
	m.ports.Collect(ch)
	m.addedNatEntries.Collect(ch)
	m.removedNatEntries.Collect(ch)
}

func (m *ServerMetrics) SetVersion(version string) {
	m.buildInfo.WithLabelValues(version).Set(1)
}

func (m *ServerMetrics) SetNumAccessKeys(numKeys int, ports int) {
	m.accessKeys.Set(float64(numKeys))
	m.ports.Set(float64(ports))
}

func (m *ServerMetrics) AddNATEntry()    { m.addedNatEntries.Inc() }
func (m *ServerMetrics) RemoveNATEntry() { m.removedNatEntries.Inc() }

// TransferTracker wraps ServiceMetrics to also track per-key byte transfer
// for the REST API's /metrics/transfer endpoint.
type TransferTracker struct {
	service.ServiceMetrics
	mu       sync.Mutex
	transfer map[string]int64 // key ID -> cumulative outbound bytes
}

func NewTransferTracker(inner service.ServiceMetrics) *TransferTracker {
	return &TransferTracker{
		ServiceMetrics: inner,
		transfer:       make(map[string]int64),
	}
}

func (t *TransferTracker) AddOpenTCPConnection(conn net.Conn) service.TCPConnMetrics {
	inner := t.ServiceMetrics.AddOpenTCPConnection(conn)
	return &trackingTCPConnMetrics{TCPConnMetrics: inner, tracker: t}
}

// GetTransferByKey returns cumulative byte transfer per key and resets counters.
func (t *TransferTracker) GetTransferByKey() map[string]int64 {
	t.mu.Lock()
	defer t.mu.Unlock()
	result := t.transfer
	t.transfer = make(map[string]int64)
	return result
}

func (t *TransferTracker) addBytes(keyID string, bytes int64) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.transfer[keyID] += bytes
}

type trackingTCPConnMetrics struct {
	service.TCPConnMetrics
	tracker   *TransferTracker
	accessKey string
}

func (m *trackingTCPConnMetrics) AddAuthentication(accessKey string) {
	m.accessKey = accessKey
	m.TCPConnMetrics.AddAuthentication(accessKey)
}

func (m *trackingTCPConnMetrics) AddClose(status string, data svcmetrics.ProxyMetrics, duration time.Duration) {
	if m.accessKey != "" {
		m.tracker.addBytes(m.accessKey, data.ProxyTarget+data.TargetProxy)
	}
	m.TCPConnMetrics.AddClose(status, data, duration)
}

// SetupRegistry creates a Prometheus registry with all metric collectors.
func SetupRegistry(
	ip2info ipinfo.IPInfoMap,
	serverMetrics *ServerMetrics,
	version string,
	nodeCollectors []string,
	logger *slog.Logger,
) (*prometheus.Registry, *TransferTracker, *ServerMetrics, error) {
	registry := prometheus.NewRegistry()

	// Go runtime metrics.
	registry.MustRegister(collectors.NewGoCollector())
	registry.MustRegister(collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}))

	// Shadowsocks service metrics.
	serviceMetrics, err := outline_prometheus.NewServiceMetrics(ip2info)
	if err != nil {
		return nil, nil, nil, err
	}

	serverMetrics.SetVersion(version)

	r := prometheus.WrapRegistererWithPrefix("shadowsocks_", registry)
	r.MustRegister(serverMetrics, serviceMetrics)

	transferTracker := NewTransferTracker(serviceMetrics)

	// Node exporter collectors.
	if len(nodeCollectors) > 0 {
		if err := registerNodeCollectors(registry, logger, nodeCollectors); err != nil {
			logger.Warn("Failed to register node_exporter collectors, continuing without them.", "err", err)
		}
	}

	return registry, transferTracker, serverMetrics, nil
}