package nettool

import (
	"context"
	"fmt"
	"log/slog"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/errz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/logz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/mathz"
)

const (
	// gcPeriodJitterPercentage is the percentage of the configured GC period that's added as jitter.
	gcPeriodJitterPercentage = 5
)

var (
	_ net.Listener = (*listenerMetricsWrapper)(nil)
	_ net.Conn     = (*connMetricsWrapper)(nil)
)

type ConnectionSet interface {
	Set(connID int64) error
	Unset(connID int64) error
	GC(context.Context) error
}

type NewConnectionSet func(name string, ttl time.Duration) (ConnectionSet, error)

type connectionSetHolder struct {
	name string
	set  ConnectionSet
}

type ListenerMetrics struct {
	Log      *slog.Logger
	NewSet   NewConnectionSet
	GCPeriod time.Duration
	ErrRep   errz.ErrReporter

	mu   sync.Mutex // protects sets
	sets []connectionSetHolder
}

func (m *ListenerMetrics) RunGC(ctx context.Context) {
	gcPeriodWithJitter := mathz.DurationWithPositiveJitter(m.GCPeriod, gcPeriodJitterPercentage)
	ticker := time.NewTicker(gcPeriodWithJitter)
	defer ticker.Stop()
	done := ctx.Done()
	for {
		select {
		case <-done:
			return
		case <-ticker.C:
			m.runGC() //nolint: contextcheck
		}
	}
}

func (m *ListenerMetrics) runGC() {
	m.mu.Lock()
	defer m.mu.Unlock()

	// Iterating with the mutex held since nothing else should be using it at this stage.

	for _, setH := range m.sets {
		func() {
			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
			defer cancel()
			err := setH.set.GC(ctx)
			if err != nil {
				m.ErrRep.HandleProcessingError(context.Background(), m.Log, "Connection set GC failed", fmt.Errorf("%s: %w", setH.name, err))
			}
		}()
	}
}

// Wrap wraps the provided listener.
// Can be called concurrently.
func (m *ListenerMetrics) Wrap(l net.Listener, name string, ttl time.Duration) (net.Listener, error) {
	set, err := m.NewSet(name, ttl)
	if err != nil {
		return nil, err
	}
	m.mu.Lock()
	m.sets = append(m.sets, connectionSetHolder{
		name: name,
		set:  set,
	})
	m.mu.Unlock()
	wrapper := &listenerMetricsWrapper{
		Listener: l,
		log:      m.Log.With(logz.ListenerName(name)),
		set:      set,
		errRep:   m.ErrRep,
	}
	return wrapper, nil
}

type listenerMetricsWrapper struct {
	connCounter atomic.Int64
	net.Listener
	log    *slog.Logger
	set    ConnectionSet
	errRep errz.ErrReporter
}

func (w *listenerMetricsWrapper) Accept() (net.Conn, error) {
	conn, err := w.Listener.Accept()
	if err != nil {
		return nil, err
	}
	connID := w.connCounter.Add(1)
	err = w.set.Set(connID)
	if err != nil {
		w.errRep.HandleProcessingError(context.Background(), w.log, "conn accept Set failed", err)
		// continue
	}
	return &connMetricsWrapper{
		Conn: conn,
		closeOnce: sync.OnceFunc(func() {
			// only want to unset once
			w.unset(connID)
		}),
	}, nil
}

func (w *listenerMetricsWrapper) unset(connID int64) {
	err := w.set.Unset(connID)
	if err != nil {
		w.errRep.HandleProcessingError(context.Background(), w.log, "conn close Unset failed", err)
	}
}

type connMetricsWrapper struct {
	net.Conn
	closeOnce func()
}

func (w *connMetricsWrapper) Close() error {
	w.closeOnce()
	return w.Conn.Close()
}
