package tunclient

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tunnel/info"
	"google.golang.org/grpc/reflection"
	"k8s.io/apimachinery/pkg/util/wait"
)

type state int8

const (
	// zero value is invalid to catch initialization bugs.
	_ state = iota
	idle
	active
	stopped
)

// ConnectionFactory helps to inject fake connections for testing.
type ConnectionFactory func(onActive, onIdle func(ConnectionInterface)) ConnectionInterface

type connectionInfo struct {
	pollCancel context.CancelFunc
	lastActive time.Time
	state      state
}

// ConnectionManager manages a pool of connections and their lifecycles.
type ConnectionManager struct {
	mu          sync.Mutex // protects connections,idleConnections,activeConnections
	connections map[ConnectionInterface]connectionInfo
	// Counters to track connections in those states. There may be timedOut connections in the map too.
	idleConnections   int32
	activeConnections int32

	wg wait.Group

	// minIdleConnections is the minimum number of connections that are not streaming a request.
	minIdleConnections int32
	// maxConnections is the maximum number of connections (idle and active).
	maxConnections int32
	// scaleUpStep is the number of new connections to start when below minIdleConnections.
	scaleUpStep int32
	// maxIdleTime is the maximum duration of time a connection can stay in an idle state.
	maxIdleTime       time.Duration
	connectionFactory ConnectionFactory
}

func NewConnectionManager(minIdleConnections, maxConnections, scaleUpStep int32, maxIdleTime time.Duration,
	connectionFactory ConnectionFactory) *ConnectionManager {
	return &ConnectionManager{
		connections:        make(map[ConnectionInterface]connectionInfo),
		minIdleConnections: minIdleConnections,
		maxConnections:     maxConnections,
		scaleUpStep:        scaleUpStep,
		maxIdleTime:        maxIdleTime,
		connectionFactory:  connectionFactory,
	}
}

func (m *ConnectionManager) Run(ctx context.Context) {
	defer m.wg.Wait() // blocks here until ctx is done and all connections exit
	m.mu.Lock()
	defer m.mu.Unlock()
	for m.idleConnections < m.minIdleConnections {
		m.startConnectionLocked(ctx)
	}
}

func (m *ConnectionManager) startConnectionLocked(rootCtx context.Context) {
	m.idleConnections++
	c := m.connectionFactory(
		func(c ConnectionInterface) {
			m.onActive(rootCtx, c)
		},
		m.onIdle,
	)
	pollCtx, pollCancel := context.WithCancel(rootCtx)
	m.connections[c] = connectionInfo{
		pollCancel: pollCancel,
		state:      idle,
	}
	m.wg.StartWithContext(rootCtx, func(rootCtx context.Context) {
		defer m.onStop(c)
		c.Run(rootCtx, pollCtx)
	})
}

func (m *ConnectionManager) onActive(rootCtx context.Context, c ConnectionInterface) {
	m.mu.Lock()
	defer m.mu.Unlock()
	i := m.connections[c]
	switch i.state {
	case idle: // idle -> active transition
		i.state = active
		m.connections[c] = i
		m.idleConnections--
		m.activeConnections++
		if m.idleConnections < m.minIdleConnections {
			// Not enough idle connections. Must scale up the number of connections.
			// Ensure we don't go above the limit.
			scaleBy := m.scaleUpStep
			haveConnections := m.idleConnections + m.activeConnections
			canSpawnConnections := m.maxConnections - haveConnections
			if scaleBy > canSpawnConnections {
				scaleBy = canSpawnConnections
			}
			for ; scaleBy > 0; scaleBy-- {
				m.startConnectionLocked(rootCtx)
			}
		}
	case active:
		panic(errors.New("connection is already active"))
	case stopped:
		panic(errors.New("invalid state: stopped"))
	default:
		panic(fmt.Errorf("unknown state: %d", i.state))
	}
}

func (m *ConnectionManager) onIdle(c ConnectionInterface) {
	m.mu.Lock()
	defer m.mu.Unlock()
	i := m.connections[c]
	switch i.state {
	case idle:
		// Already in the idle state. This can happen if a gRPC connection was established, but never
		// transitioned into the active state.
		if m.idleConnections > m.minIdleConnections && time.Since(i.lastActive) > m.maxIdleTime {
			// Too many idle connections, can stop this one.
			i.pollCancel()
			// this counter must be decremented right when the connection is stopped so that other connections
			// don't stop too while it is off by one. i.e. cannot do this in onStop().
			m.idleConnections--
			i.state = stopped
			m.connections[c] = i
		}
	case active: // active -> idle transition
		i.state = idle
		i.lastActive = time.Now()
		m.connections[c] = i
		m.idleConnections++
		m.activeConnections--
	case stopped:
		panic(errors.New("invalid state: stopped"))
	default:
		panic(fmt.Errorf("unknown state: %d", i.state))
	}
}

func (m *ConnectionManager) onStop(c ConnectionInterface) {
	m.mu.Lock()
	defer m.mu.Unlock()
	i := m.connections[c]
	delete(m.connections, c)
	if i.state != stopped {
		// onIdle() decrements this field if maxIdleTime has been reached.
		// It's decremented here too to handle context done situation.
		m.idleConnections--
	}
}

func APIDescriptor(server reflection.ServiceInfoProvider) *info.APIDescriptor {
	serverInfo := server.GetServiceInfo()
	services := make([]*info.Service, 0, len(serverInfo))
	for svcName, svcInfo := range serverInfo {
		methods := make([]*info.Method, 0, len(svcInfo.Methods))
		for _, mInfo := range svcInfo.Methods {
			methods = append(methods, &info.Method{
				Name: mInfo.Name,
			})
		}
		services = append(services, &info.Service{
			Name:    svcName,
			Methods: methods,
		})
	}
	return &info.APIDescriptor{
		Services: services,
	}
}
