package agent_tracker //nolint:stylecheck

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

	"github.com/redis/rueidis"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/tool/errz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/tool/logz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/tool/redistool"
	otelmetric "go.opentelemetry.io/otel/metric"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
	"google.golang.org/protobuf/proto"
	"k8s.io/apimachinery/pkg/util/wait"
)

const (
	refreshOverlap                       = 5 * time.Second
	connectedAgentsKey             int64 = 0
	connectionsByAgentIDHashName         = "connections_by_agent_id"
	connectionsByProjectIDHashName       = "connections_by_project_id"
	connectedAgentsHashName              = "connected_agents"
	agentVersionsHashName                = "agent_versions"
	// agentVersionKey is not actually used as a key. See `agentVersionsHashKey` function.
	agentVersionKey                   int64 = 0
	connectionsByAgentVersionHashName       = "connections_by_agent_version"
)

type ConnectedAgentInfoCallback func(*ConnectedAgentInfo) (done bool, err error)

type Registerer interface {
	// RegisterConnection registers connection with the tracker.
	RegisterConnection(ctx context.Context, info *ConnectedAgentInfo) error
	// UnregisterConnection unregisters connection with the tracker.
	UnregisterConnection(ctx context.Context, info *ConnectedAgentInfo) error
}

type ExpiringRegisterer interface {
	// RegisterExpiring registers connection with the tracker.
	// Registration will expire if not refreshed using this method.
	RegisterExpiring(ctx context.Context, info *ConnectedAgentInfo) error
}

type Querier interface {
	GetConnectionsByAgentID(ctx context.Context, agentID int64, cb ConnectedAgentInfoCallback) error
	GetConnectionsByProjectID(ctx context.Context, projectID int64, cb ConnectedAgentInfoCallback) error
	GetConnectedAgentsCount(ctx context.Context) (int64, error)
	CountAgentsByAgentVersions(ctx context.Context) (map[string]int64, error)
}

type Tracker interface {
	Registerer
	ExpiringRegisterer
	Querier
	Run(ctx context.Context) error
}

type RedisTracker struct {
	log           *zap.Logger
	errRep        errz.ErrReporter
	ttl           time.Duration
	refreshPeriod time.Duration
	gcPeriod      time.Duration

	// mu protects fields below
	mu                     sync.Mutex
	connectionsByAgentID   redistool.ExpiringHash[int64, int64] // agentID -> connectionId -> info
	connectionsByProjectID redistool.ExpiringHash[int64, int64] // projectID -> connectionId -> info
	connectedAgents        redistool.ExpiringHash[int64, int64] // hash name -> agentID -> ""

	// agentVersions keeps track of the list of agent versions that have active agents.
	agentVersions redistool.ExpiringHash[int64, string] // hash name -> agentVersions -> ""
	// connectionsByAgentVersion stores connections data for each agent versions.
	connectionsByAgentVersion redistool.ExpiringHash[string, int64] // agentVersion -> connectionId -> hash
}

func NewRedisTracker(log *zap.Logger, errRep errz.ErrReporter, client rueidis.Client, agentKeyPrefix string,
	ttl, refreshPeriod, gcPeriod time.Duration, m otelmetric.Meter) (*RedisTracker, error) {
	connectionsByAgentID, err := redistool.NewRedisExpiringHash(connectionsByAgentIDHashName, client, connectionsByAgentIDHashKey(agentKeyPrefix), int64ToStr, ttl, m, true)
	if err != nil {
		return nil, err
	}
	connectionsByProjectID, err := redistool.NewRedisExpiringHash(connectionsByProjectIDHashName, client, connectionsByProjectIDHashKey(agentKeyPrefix), int64ToStr, ttl, m, true)
	if err != nil {
		return nil, err
	}
	connectedAgents, err := redistool.NewRedisExpiringHash(connectedAgentsHashName, client, connectedAgentsHashKey(agentKeyPrefix), int64ToStr, ttl, m, false)
	if err != nil {
		return nil, err
	}
	agentVersions, err := redistool.NewRedisExpiringHash(agentVersionsHashName, client, agentVersionsHashKey(agentKeyPrefix), strToStr, ttl, m, false)
	if err != nil {
		return nil, err
	}
	connectionsByAgentVersion, err := redistool.NewRedisExpiringHash(connectionsByAgentVersionHashName, client, connectionsByAgentVersionHashKey(agentKeyPrefix), int64ToStr, ttl, m, false)
	if err != nil {
		return nil, err
	}

	return &RedisTracker{
		log:                       log,
		errRep:                    errRep,
		ttl:                       ttl,
		refreshPeriod:             refreshPeriod,
		gcPeriod:                  gcPeriod,
		connectionsByAgentID:      connectionsByAgentID,
		connectionsByProjectID:    connectionsByProjectID,
		connectedAgents:           connectedAgents,
		agentVersions:             agentVersions,
		connectionsByAgentVersion: connectionsByAgentVersion,
	}, nil
}

func (t *RedisTracker) Run(ctx context.Context) error {
	refreshTicker := time.NewTicker(t.refreshPeriod)
	defer refreshTicker.Stop()
	gcTicker := time.NewTicker(t.gcPeriod)
	defer gcTicker.Stop()
	done := ctx.Done()
	for {
		select {
		case <-done:
			return nil
		case <-refreshTicker.C:
			t.refreshRegistrations(ctx, time.Now().Add(t.refreshPeriod-refreshOverlap))
		case <-gcTicker.C:
			keysDeleted := t.runGC(ctx)
			if keysDeleted > 0 {
				t.log.Info("Deleted expired agent connections records", logz.RemovedHashKeys(keysDeleted))
			}
		}
	}
}

func (t *RedisTracker) RegisterConnection(ctx context.Context, info *ConnectedAgentInfo) error {
	infoBytes, err := proto.Marshal(info)
	if err != nil {
		// This should never happen
		return fmt.Errorf("failed to marshal object: %w", err)
	}
	exp := time.Now().Add(t.ttl)

	t.mu.Lock()
	defer t.mu.Unlock()
	var wg errgroup.Group
	wg.Go(func() error {
		return t.connectionsByProjectID.Set(ctx, info.ProjectId, info.ConnectionId, infoBytes)
	})
	wg.Go(func() error {
		return t.connectionsByAgentID.Set(ctx, info.AgentId, info.ConnectionId, infoBytes)
	})
	wg.Go(func() error {
		return t.connectedAgents.Set(ctx, connectedAgentsKey, info.AgentId, nil)
	})
	wg.Go(func() error {
		agentPodInfoBytes, err := proto.Marshal(&AgentPodInfo{
			AgentId: info.AgentId,
			PodId:   info.ConnectionId,
		})
		if err != nil {
			// This should never happen
			return fmt.Errorf("failed to marshal AgentPodInfo object: %w", err)
		}
		return t.connectionsByAgentVersion.SetEX(ctx, info.AgentMeta.Version, info.ConnectionId, agentPodInfoBytes, exp)
	})
	wg.Go(func() error {
		return t.agentVersions.SetEX(ctx, agentVersionKey, info.AgentMeta.Version, nil, exp)
	})
	return wg.Wait()
}

func (t *RedisTracker) UnregisterConnection(ctx context.Context, info *ConnectedAgentInfo) error {
	t.mu.Lock()
	defer t.mu.Unlock()
	var wg errgroup.Group
	wg.Go(func() error {
		return t.connectionsByProjectID.Unset(ctx, info.ProjectId, info.ConnectionId)
	})
	wg.Go(func() error {
		return t.connectionsByAgentID.Unset(ctx, info.AgentId, info.ConnectionId)
	})
	t.connectedAgents.Forget(connectedAgentsKey, info.AgentId)
	wg.Go(func() error {
		return t.connectionsByAgentVersion.Unset(ctx, info.AgentMeta.Version, info.ConnectionId)
	})
	return wg.Wait()
}

func (t *RedisTracker) RegisterExpiring(ctx context.Context, info *ConnectedAgentInfo) error {
	infoBytes, err := proto.Marshal(info)
	if err != nil {
		// This should never happen
		return fmt.Errorf("failed to marshal object: %w", err)
	}
	exp := time.Now().Add(t.ttl)
	var wg errgroup.Group
	wg.Go(func() error {
		return t.connectionsByProjectID.SetEX(ctx, info.ProjectId, info.ConnectionId, infoBytes, exp)
	})
	wg.Go(func() error {
		return t.connectionsByAgentID.SetEX(ctx, info.AgentId, info.ConnectionId, infoBytes, exp)
	})
	wg.Go(func() error {
		return t.connectedAgents.SetEX(ctx, connectedAgentsKey, info.AgentId, nil, exp)
	})
	wg.Go(func() error {
		agentPodInfoBytes, err := proto.Marshal(&AgentPodInfo{
			AgentId: info.AgentId,
			PodId:   info.ConnectionId,
		})
		if err != nil {
			// This should never happen
			return fmt.Errorf("failed to marshal AgentPodInfo object: %w", err)
		}
		return t.connectionsByAgentVersion.SetEX(ctx, info.AgentMeta.Version, info.ConnectionId, agentPodInfoBytes, exp)
	})
	wg.Go(func() error {
		return t.agentVersions.SetEX(ctx, agentVersionKey, info.AgentMeta.Version, nil, exp)
	})
	return wg.Wait()
}

func (t *RedisTracker) GetConnectionsByAgentID(ctx context.Context, agentID int64, cb ConnectedAgentInfoCallback) error {
	return t.getConnectionsByKey(ctx, t.connectionsByAgentID, agentID, cb)
}

func (t *RedisTracker) GetConnectionsByProjectID(ctx context.Context, projectID int64, cb ConnectedAgentInfoCallback) error {
	return t.getConnectionsByKey(ctx, t.connectionsByProjectID, projectID, cb)
}

func (t *RedisTracker) GetConnectedAgentsCount(ctx context.Context) (int64, error) {
	return t.connectedAgents.Len(ctx, connectedAgentsKey)
}

func (t *RedisTracker) CountAgentsByAgentVersions(ctx context.Context) (map[string]int64, error) {
	agentVersions, err := t.getAgentVersions(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to get agent versions from Redis: %w", err)
	}

	counts := make(map[string]int64, len(agentVersions))
	for _, version := range agentVersions {
		count, err := t.connectionsByAgentVersion.Len(ctx, version)
		if err != nil {
			return nil, fmt.Errorf("failed to get hash length from connectionsByAgentVersion in Redis: %w", err)
		}
		counts[version] = count
	}
	return counts, nil
}

func (t *RedisTracker) refreshRegistrations(ctx context.Context, nextRefresh time.Time) {
	t.mu.Lock()
	defer t.mu.Unlock()
	// Run refreshes concurrently to release mu ASAP.
	var wg wait.Group
	t.refreshHash(ctx, &wg, t.connectionsByProjectID, nextRefresh)
	t.refreshHash(ctx, &wg, t.connectionsByAgentID, nextRefresh)
	t.refreshHash(ctx, &wg, t.connectedAgents, nextRefresh)
	wg.Wait()
}

func (t *RedisTracker) refreshHash(ctx context.Context, wg *wait.Group, h redistool.ExpiringHash[int64, int64], nextRefresh time.Time) {
	wg.Start(func() {
		err := h.Refresh(ctx, nextRefresh)
		if err != nil {
			t.errRep.HandleProcessingError(ctx, t.log, fmt.Sprintf("Failed to refresh %s hash data in Redis", h.GetName()), err)
		}
	})
}

func (t *RedisTracker) runGC(ctx context.Context) int {
	type gcFunc struct {
		hashName string
		gc       func(context.Context) (int, error)
	}
	var gcFuncs []gcFunc // slice ensures deterministic iteration order
	func() {
		t.mu.Lock()
		defer t.mu.Unlock()
		gcFuncs = []gcFunc{
			{
				hashName: connectionsByProjectIDHashName,
				gc:       t.connectionsByProjectID.GC(),
			},
			{
				hashName: connectionsByAgentIDHashName,
				gc:       t.connectionsByAgentID.GC(),
			},
			{
				hashName: connectedAgentsHashName,
				gc:       t.connectedAgents.GC(),
			},
			{
				hashName: agentVersionsHashName,
				gc:       t.agentVersions.GC(), // First GC agentVersions
			},
			{
				hashName: connectionsByAgentVersionHashName,
				gc:       t.gcConnectionsByAgentVersion, // Then GC hashes based on what is in agentVersions
			},
		}
	}()
	keysDeleted := 0
	// No rush so run GC sequentially to not stress RAM/CPU/Redis/network.
	// We have more important work to do that we shouldn't impact.
	for _, gf := range gcFuncs {
		deleted, err := gf.gc(ctx)
		keysDeleted += deleted
		if err != nil {
			if errz.ContextDone(err) {
				t.log.Debug("Redis GC interrupted", logz.Error(err))
				break
			}
			t.errRep.HandleProcessingError(ctx, t.log, fmt.Sprintf("Failed to GC data in %s Redis hash", gf.hashName), err)
			// continue anyway
		}
	}
	return keysDeleted
}

func (t *RedisTracker) getConnectionsByKey(ctx context.Context, hash redistool.ExpiringHash[int64, int64], key int64, cb ConnectedAgentInfoCallback) error {
	err := hash.Scan(ctx, key, func(rawHashKey string, value []byte, err error) (bool, error) {
		if err != nil {
			t.errRep.HandleProcessingError(ctx, t.log, fmt.Sprintf("Redis %s hash scan", hash.GetName()), err)
			return false, nil
		}
		var info ConnectedAgentInfo
		err = proto.Unmarshal(value, &info)
		if err != nil {
			t.errRep.HandleProcessingError(ctx, t.log, fmt.Sprintf("Redis %s hash scan: proto.Unmarshal(ConnectedAgentInfo)", hash.GetName()), err)
			return false, nil
		}
		return cb(&info)
	})
	return err
}

func (t *RedisTracker) gcConnectionsByAgentVersion(ctx context.Context) (int, error) {
	// Get a list of agent versions.
	agentVersions, err1 := t.getAgentVersions(ctx)

	// GC connectionsByAgentVersion for agent versions that we got from agentVersions hash.
	deletedConnections, err2 := t.connectionsByAgentVersion.GCFor(agentVersions)(ctx)

	return deletedConnections, errors.Join(err1, err2)
}

func (t *RedisTracker) getAgentVersions(ctx context.Context) ([]string, error) {
	var agentVersions []string
	err := t.agentVersions.Scan(ctx, agentVersionKey, func(rawHashKey string, value []byte, err error) (bool, error) {
		if err != nil {
			t.errRep.HandleProcessingError(ctx, t.log, "getAgentVersions: failed to scan redis hash", err)
			return false, nil
		}

		agentVersions = append(agentVersions, rawHashKey)
		return false, nil
	})
	return agentVersions, err
}

// connectionsByAgentIDHashKey returns a key for agentID -> (connectionId -> marshaled ConnectedAgentInfo).
func connectionsByAgentIDHashKey(agentKeyPrefix string) redistool.KeyToRedisKey[int64] {
	prefix := agentKeyPrefix + ":conn_by_agent_id:"
	return func(agentID int64) string {
		return redistool.PrefixedInt64Key(prefix, agentID)
	}
}

// connectionsByProjectIDHashKey returns a key for projectID -> (agentID ->marshaled ConnectedAgentInfo).
func connectionsByProjectIDHashKey(agentKeyPrefix string) redistool.KeyToRedisKey[int64] {
	prefix := agentKeyPrefix + ":conn_by_project_id:"
	return func(projectID int64) string {
		return redistool.PrefixedInt64Key(prefix, projectID)
	}
}

// connectedAgentsHashKey returns the key for the hash of connected agents.
func connectedAgentsHashKey(agentKeyPrefix string) redistool.KeyToRedisKey[int64] {
	prefix := agentKeyPrefix + ":connected_agents"
	return func(_ int64) string {
		return prefix
	}
}

func connectionsByAgentVersionHashKey(agentKeyPrefix string) redistool.KeyToRedisKey[string] {
	prefix := agentKeyPrefix + ":conn_by_agent_version:"
	return func(agentVersion string) string {
		return prefix + agentVersion
	}
}

func agentVersionsHashKey(agentKeyPrefix string) redistool.KeyToRedisKey[int64] {
	prefix := agentKeyPrefix + ":agent_versions"
	return func(_ int64) string {
		return prefix
	}
}

type ConnectedAgentInfoCollector []*ConnectedAgentInfo

func (c *ConnectedAgentInfoCollector) Collect(info *ConnectedAgentInfo) (bool, error) {
	*c = append(*c, info)
	return false, nil
}

func int64ToStr(key int64) string {
	return strconv.FormatInt(key, 10)
}

func strToStr(s string) string {
	return s
}