package nettool

import (
	"context"
	"crypto/tls"
	"fmt"
	"log/slog"
	"net"
	"sync"
	"sync/atomic"
	"syscall"
	"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, expectTLS bool) (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,
		expectTLS: expectTLS,
	}
	return wrapper, nil
}

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

type cleartextConn interface {
	net.Conn
	syscall.Conn
}

type tlsConn interface {
	net.Conn
	ConnectionState() tls.ConnectionState
	HandshakeContext(context.Context) error
}

func (w *listenerMetricsWrapper) Accept() (net.Conn, error) {
	conn, err := w.Listener.Accept()
	if err != nil {
		return nil, err
	}

	if w.expectTLS {
		tlsC, ok := conn.(tlsConn)
		if !ok {
			_ = conn.Close()
			return nil, fmt.Errorf("expected TLS connection, got %T", conn)
		}
		return &tlsConnMetricsWrapper{
			tlsConn:   tlsC,
			closeOnce: w.connClose(),
		}, nil
	} else {
		c, ok := conn.(cleartextConn)
		if !ok {
			_ = conn.Close()
			return nil, fmt.Errorf("expected syscall.Conn connection, got %T", conn)
		}
		return &connMetricsWrapper{
			cleartextConn: c,
			closeOnce:     w.connClose(),
		}, nil
	}
}

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

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 {
	cleartextConn
	closeOnce func()
}

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

type tlsConnMetricsWrapper struct {
	tlsConn
	closeOnce func()
}

func (w *tlsConnMetricsWrapper) Close() error {
	w.closeOnce()
	return w.tlsConn.Close()
}
