package router

import (
	"context"
	"fmt"
	"log/slog"
	"time"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modshared"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/logz"
	otelmetric "go.opentelemetry.io/otel/metric"
)

type asyncTaskType byte

const (
	_ asyncTaskType = iota
	asyncTaskRegister
	asyncTaskUnregister
	asyncTaskGC
	asyncTaskRefresh
)

const (
	queueSize                              = 4 * 1024
	maxBatchSize                           = 128
	asyncTaskSubmissionDurationMetricName  = "registry_async_submission_duration"
	asyncTaskSubmissionCounterMetricName   = "registry_async_submission"
	asyncTaskSubmissionBatchSizeMetricName = "registry_async_submission_batch_size"
)

type asyncTask struct {
	agentKeys []api.AgentKey
	taskType  asyncTaskType
}

// asyncTracker performs all tunnel tracker (un)registration/refresh/GC I/O for the registry.
// By using a single chanel to submit work we ensure that no races is possible and e.g. a tunnel registration and
// unregistration cannot be reordered.
// Code in Registry that uses asyncTracker must always hold a mutex while both checking state and using
// asyncTracker's methods. This ensures there are no races in the client code i.e. requests to perform I/O
// are submitted to the asyncTracker without any reordering due to concurrency.
type asyncTracker struct {
	log                 *slog.Logger
	api                 modshared.API
	tunnelTracker       Registerer
	ttl                 time.Duration
	taskQueue           chan asyncTask
	submissionDuration  otelmetric.Float64Histogram
	submissionCounter   otelmetric.Int64Counter
	submissionBatchSize otelmetric.Int64Histogram

	batch RegistrationBuilder
	// batchedAgentKeys counts for how many agent keys batch has outstanding operations.
	// It's a coarse way to limit the size of a batch.
	// We don't use number of operations (i.e. register/unregister/refresh) because number of agent keys per operation
	// can be as little as 1 or potentially very huge (thousands for refresh).
	batchedAgentKeys int
}

func newAsyncTracker(log *slog.Logger, api modshared.API, meter otelmetric.Meter, tunnelTracker Registerer, ttl time.Duration) (*asyncTracker, error) {
	submissionDuration, err := meter.Float64Histogram(
		asyncTaskSubmissionDurationMetricName,
		otelmetric.WithUnit("s"),
		otelmetric.WithDescription("The time it takes to submit a tunnel tracking task in seconds"),
		otelmetric.WithExplicitBucketBoundaries(0.0001, 0.001, 0.004, 0.016, 0.064, 0.256, 1.024, 4.096, 16.384),
	)
	if err != nil {
		return nil, err
	}
	submissionCounter, err := meter.Int64Counter(
		asyncTaskSubmissionCounterMetricName,
		otelmetric.WithDescription("The total number of submitted tunnel tracking tasks"),
	)
	if err != nil {
		return nil, err
	}
	submissionBatchSize, err := meter.Int64Histogram(
		asyncTaskSubmissionBatchSizeMetricName,
		otelmetric.WithDescription("The number of batched operations"),
		otelmetric.WithExplicitBucketBoundaries(0, 1, 4, 16, 64, 128),
	)
	if err != nil {
		return nil, err
	}
	return &asyncTracker{
		log:                 log,
		api:                 api,
		tunnelTracker:       tunnelTracker,
		ttl:                 ttl,
		taskQueue:           make(chan asyncTask, queueSize),
		submissionDuration:  submissionDuration,
		submissionCounter:   submissionCounter,
		submissionBatchSize: submissionBatchSize,
		batch:               tunnelTracker.RegistrationBuilder(),
	}, nil
}

func (t *asyncTracker) done() {
	close(t.taskQueue)
}

func (t *asyncTracker) run() {
	for task := range t.taskQueue {
		t.handleTask(task)
	batching:
		for {
			var ok bool
			select {
			case task, ok = <-t.taskQueue:
				if !ok {
					break batching
				}
				t.handleTask(task)
			default:
				// Nothing else in the queue, time to submit what we've accumulated.
				break batching
			}
		}
		t.submitBatch()
	}
}

func (t *asyncTracker) internalRefresh(agentKeys ...api.AgentKey) {
	t.batch.Refresh(t.ttl, agentKeys...)
}

func (t *asyncTracker) internalRegister(agentKeys ...api.AgentKey) {
	t.batch.Register(t.ttl, agentKeys...)
}

func (t *asyncTracker) handleTask(task asyncTask) {
	var op func(...api.AgentKey)
	switch task.taskType {
	case asyncTaskRegister:
		op = t.internalRegister
	case asyncTaskUnregister:
		op = t.batch.Unregister
	case asyncTaskRefresh:
		op = t.internalRefresh
	case asyncTaskGC:
		t.submitBatch()
		deletedKeys, err := t.tunnelTracker.GC(context.Background(), task.agentKeys)
		if err != nil {
			t.api.HandleProcessingError(context.Background(), t.log, "Failed to GC data", err)
		}
		if deletedKeys > 0 {
			t.log.Info("Deleted expired agent tunnel records", logz.RemovedHashKeys(deletedKeys))
		}
		return
	default:
		panic(fmt.Sprintf("unknown task type: %d", task.taskType))
	}

	for len(task.agentKeys) > 0 {
		batch := task.agentKeys[:min(len(task.agentKeys), maxBatchSize-t.batchedAgentKeys)]
		op(batch...)
		t.batchedAgentKeys += len(batch)
		t.maybeSubmitBatch()
		task.agentKeys = task.agentKeys[len(batch):]
	}
}

func (t *asyncTracker) maybeSubmitBatch() {
	if t.batchedAgentKeys >= maxBatchSize {
		t.submitBatch()
	}
}

func (t *asyncTracker) submitBatch() {
	t.submissionBatchSize.Record(context.Background(), int64(t.batchedAgentKeys))
	t.batchedAgentKeys = 0
	err := t.batch.Do(context.Background())
	if err != nil {
		t.api.HandleProcessingError(context.Background(), t.log, "Failed to (un)register tunnels", err)
	}
}

// register registers tunnel with the tracker.
func (t *asyncTracker) register(agentKey api.AgentKey) {
	t.submitTask(asyncTask{
		agentKeys: []api.AgentKey{agentKey},
		taskType:  asyncTaskRegister,
	})
}

// unregister unregisters tunnels with the tracker.
func (t *asyncTracker) unregister(agentKeys ...api.AgentKey) {
	if len(agentKeys) == 0 {
		return
	}
	t.submitTask(asyncTask{
		agentKeys: agentKeys,
		taskType:  asyncTaskUnregister,
	})
}

// GC deletes expired tunnels from the underlying storage.
func (t *asyncTracker) gc(agentKeys []api.AgentKey) {
	if len(agentKeys) == 0 {
		return
	}
	t.submitTask(asyncTask{
		agentKeys: agentKeys,
		taskType:  asyncTaskGC,
	})
}

// Refresh refreshes registered tunnels in the underlying storage.
func (t *asyncTracker) refresh(agentKeys []api.AgentKey) {
	if len(agentKeys) == 0 {
		return
	}
	t.submitTask(asyncTask{
		agentKeys: agentKeys,
		taskType:  asyncTaskRefresh,
	})
}

func (t *asyncTracker) submitTask(task asyncTask) {
	start := time.Now()
	t.taskQueue <- task
	elapsed := float64(time.Since(start)) / float64(time.Second)
	t.submissionDuration.Record(context.Background(), elapsed)
	// Increment the counter after:
	// - submitting - to reduce latency.
	// - calculating the elapsed time - to increase precision.
	t.submissionCounter.Add(context.Background(), 1)
}
