package resourcelimiter

import (
	"context"
	"fmt"
	"sync"

	"github.com/pingcap/log"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/metrics"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/scheduler/nodemanager"
	rm "github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/resourcemanager/types"
	"go.uber.org/zap"
	"k8s.io/apimachinery/pkg/api/resource"
)

// ResourceLimiter sets limit for the resource usage at keyspace/cluster level.
type ResourceLimiter interface {
	// Acquire resources for the given keyspace according to pool specification.
	// Will return error if:
	// Given pool not found from pool manager, or
	// Given task already acquired for the keyspace, or
	// Keyspace resource usage would exceed limit if granted, or
	// Total resource usage would exceed limit if granted.
	Acquire(keyspaceName string, taskIDs []string, pool rm.WorkerPool) error
	// Release resources previously acquired.
	// Will return error if target task not find.
	Release(keyspaceName, taskID string) error
	// UpdateLimit sets the new limit for the resource limiter.
	// This will not retroactively affect resources already acquired.
	// i.e. if you set limit to 0 after some acquisition, only new acquisitions will be
	// blocked, the already acquired resource will not be auto released.
	UpdateLimit(memTotal, cpuTotal, memKeyspace, cpuKeyspace string) error
	// Initialize updates usage records for all active requests according to node manager.
	// This should be called when resource limiter first start.
	Initialize(ctx context.Context, nodeManager nodemanager.Manager) error
}

// taskUsageRecord represent the resource usage of a single task.
type taskUsageRecord struct {
	mem  resource.Quantity
	cpu  resource.Quantity
	pool rm.WorkerPool
}

// keyspaceUsageRecord represent the resource usage of a particular keyspace.
type keyspaceUsageRecord struct {
	sync.Mutex
	tasks map[string]*taskUsageRecord
	mem   resource.Quantity
	cpu   resource.Quantity
}

type limit struct {
	memTotal    resource.Quantity
	memKeyspace resource.Quantity
	cpuTotal    resource.Quantity
	cpuKeyspace resource.Quantity
}

func NewResourceLimiter(
	poolManager rm.PoolManager,
) ResourceLimiter {
	return &resourceLimiter{
		poolManager:  poolManager,
		usageRecords: make(map[string]*keyspaceUsageRecord),
	}
}

type resourceLimiter struct {
	poolManager rm.PoolManager

	sync.Mutex
	usageRecords map[string]*keyspaceUsageRecord
	mem          resource.Quantity
	cpu          resource.Quantity
	limit        limit
}

func (r *resourceLimiter) Acquire(keyspaceName string, taskIDs []string, pool rm.WorkerPool) error {
	pools, err := r.poolManager.GetPools()
	if err != nil {
		return err
	}
	poolSpec, ok := pools[pool]
	if !ok {
		return fmt.Errorf("pool %s not found", pool)
	}
	taskMem := poolSpec.WorkerPodSpec.Containers[0].Resources.Limits.Memory()
	if taskMem == nil {
		return fmt.Errorf("task mem limit undefined")
	}
	taskCPU := poolSpec.WorkerPodSpec.Containers[0].Resources.Limits.Cpu()
	if taskCPU == nil {
		return fmt.Errorf("task cpu limit undefined")
	}
	r.Lock()
	defer r.Unlock()
	keyspaceUsage, ok := r.usageRecords[keyspaceName]
	if !ok {
		keyspaceUsage = &keyspaceUsageRecord{
			tasks: make(map[string]*taskUsageRecord),
		}
		r.usageRecords[keyspaceName] = keyspaceUsage
	}
	keyspaceUsage.Lock()
	defer keyspaceUsage.Unlock()
	for _, taskID := range taskIDs {
		if keyspaceUsage.tasks[taskID] != nil {
			return fmt.Errorf("task %s already running", taskID)
		}
	}
	taskCount := int64(len(taskIDs))
	// Check per keyspace memory limit.
	if keyspaceUsage.mem.Value()+taskMem.Value()*taskCount > r.limit.memKeyspace.Value() {
		return fmt.Errorf("task memory usage exceeded keyspace limit,"+
			" current usage %s, request count %d, per request usage %s, limit %s",
			keyspaceUsage.mem.String(), taskCount, taskMem.String(), r.limit.memKeyspace.String())
	}
	// Check total memory limit.
	if r.mem.Value()+taskMem.Value()*taskCount > r.limit.memTotal.Value() {
		return fmt.Errorf("task would exceeded total memory limit,"+
			" current usage %s, request count %d, per request usage %s, limit %s",
			r.mem.String(), taskCount, taskMem.String(), r.limit.memTotal.String())
	}
	// Check per keyspace cpu limit.
	if keyspaceUsage.cpu.Value()+taskCPU.Value()*taskCount > r.limit.cpuKeyspace.Value() {
		return fmt.Errorf("task cpu usage exceeded keyspace limit,"+
			" current usage %s, request count %d, per request usage %s, limit %s",
			keyspaceUsage.cpu.String(), taskCount, taskCPU.String(), r.limit.cpuKeyspace.String())
	}
	// Check total cpu limit.
	if r.cpu.Value()+taskCPU.Value()*taskCount > r.limit.cpuTotal.Value() {
		return fmt.Errorf("task would exceeded total memory limit,"+
			" current usage %s, request count %d, per request usage %s, limit %s",
			r.cpu.String(), taskCount, taskCPU.String(), r.limit.cpuTotal.String())
	}
	for _, taskID := range taskIDs {
		keyspaceUsage.tasks[taskID] = &taskUsageRecord{
			mem:  *taskMem,
			cpu:  *taskCPU,
			pool: pool,
		}
		keyspaceUsage.mem.Add(*taskMem)
		r.mem.Add(*taskMem)
		keyspaceUsage.cpu.Add(*taskCPU)
		r.cpu.Add(*taskCPU)
		log.Info("resource acquired",
			zap.String("keyspace-name", keyspaceName),
			zap.String("task-id", taskID),
			zap.String("acquired mem", taskMem.String()),
			zap.String("acquired cpu", taskCPU.String()),
			zap.String("updated keyspace mem usage", keyspaceUsage.mem.String()),
			zap.String("updated keyspace cpu usage", keyspaceUsage.cpu.String()),
			zap.String("updated total mem usage", r.mem.String()),
			zap.String("updated total cpu usage", r.cpu.String()),
		)
		metrics.ResourceUsage.With(prometheus.Labels{
			metrics.ResourceTypeLabel: metrics.ResourceTypeMem,
			metrics.KeyspaceLabel:     keyspaceName,
			metrics.TaskIDLabel:       taskID,
			metrics.PoolLabel:         string(pool),
		}).Add(taskMem.AsApproximateFloat64())
		metrics.ResourceUsage.With(prometheus.Labels{
			metrics.ResourceTypeLabel: metrics.ResourceTypeCPU,
			metrics.KeyspaceLabel:     keyspaceName,
			metrics.TaskIDLabel:       taskID,
			metrics.PoolLabel:         string(pool),
		}).Add(taskCPU.AsApproximateFloat64())
	}
	return nil
}

func (r *resourceLimiter) Release(keyspaceName, taskID string) error {
	r.Lock()
	defer r.Unlock()
	keyspaceUsage, ok := r.usageRecords[keyspaceName]
	if !ok {
		return fmt.Errorf("usage for keyspace %s not found", keyspaceName)
	}
	keyspaceUsage.Lock()
	defer keyspaceUsage.Unlock()
	taskUsage, ok := keyspaceUsage.tasks[taskID]
	if !ok {
		return fmt.Errorf("task %s not found for keyspace %s", taskID, keyspaceName)
	}
	delete(keyspaceUsage.tasks, taskID)
	// Log warning and reset to zero on resource usage underflow during release.
	// Prevents potential tracking inaccuracies due to pool spec changes or race conditions.
	keyspaceUsage.mem.Sub(taskUsage.mem)
	if keyspaceUsage.mem.Sign() < 0 {
		log.Warn("keyspace mem usage underflow, reset to zero",
			zap.String("keyspace-name", keyspaceName),
			zap.String("task-id", taskID),
			zap.String("current-mem-usage", keyspaceUsage.mem.String()),
			zap.String("task-mem-usage", keyspaceUsage.mem.String()),
		)
		keyspaceUsage.mem = resource.Quantity{}
	}
	keyspaceUsage.cpu.Sub(taskUsage.cpu)
	if keyspaceUsage.cpu.Sign() < 0 {
		log.Warn("keyspace cpu usage underflow, reset to zero",
			zap.String("keyspace-name", keyspaceName),
			zap.String("task-id", taskID),
			zap.String("current-cpu-usage", keyspaceUsage.cpu.String()),
			zap.String("task-cpu-usage", keyspaceUsage.cpu.String()),
		)
		keyspaceUsage.cpu = resource.Quantity{}
	}
	r.mem.Sub(taskUsage.mem)
	if r.mem.Sign() < 0 {
		log.Warn("total mem usage underflow, reset to zero",
			zap.String("keyspace-name", keyspaceName),
			zap.String("task-id", taskID),
			zap.String("current-mem-usage", keyspaceUsage.mem.String()),
			zap.String("task-mem-usage", keyspaceUsage.mem.String()),
		)
		r.mem = resource.Quantity{}
	}
	r.cpu.Sub(taskUsage.cpu)
	if r.cpu.Sign() < 0 {
		log.Warn("total cpu usage underflow, reset to zero",
			zap.String("keyspace-name", keyspaceName),
			zap.String("task-id", taskID),
			zap.String("current-cpu-usage", keyspaceUsage.cpu.String()),
			zap.String("task-cpu-usage", keyspaceUsage.cpu.String()),
		)
		r.cpu = resource.Quantity{}
	}
	log.Info("resource released",
		zap.String("keyspace-name", keyspaceName),
		zap.String("task-id", taskID),
		zap.String("released mem", taskUsage.mem.String()),
		zap.String("released cpu", taskUsage.cpu.String()),
		zap.String("updated keyspace mem usage", keyspaceUsage.mem.String()),
		zap.String("updated keyspace cpu usage", keyspaceUsage.cpu.String()),
		zap.String("updated total mem usage", r.mem.String()),
		zap.String("updated total cpu usage", r.cpu.String()),
	)
	metrics.ResourceUsage.With(prometheus.Labels{
		metrics.ResourceTypeLabel: metrics.ResourceTypeMem,
		metrics.KeyspaceLabel:     keyspaceName,
		metrics.TaskIDLabel:       taskID,
		metrics.PoolLabel:         string(taskUsage.pool),
	}).Sub(taskUsage.mem.AsApproximateFloat64())
	metrics.ResourceUsage.With(prometheus.Labels{
		metrics.ResourceTypeLabel: metrics.ResourceTypeCPU,
		metrics.KeyspaceLabel:     keyspaceName,
		metrics.TaskIDLabel:       taskID,
		metrics.PoolLabel:         string(taskUsage.pool),
	}).Sub(taskUsage.cpu.AsApproximateFloat64())
	return nil
}

func (r *resourceLimiter) UpdateLimit(memTotal, cpuTotal, memKeyspace, cpuKeyspace string) error {
	mem, err := resource.ParseQuantity(memTotal)
	if err != nil {
		return err
	}
	cpu, err := resource.ParseQuantity(cpuTotal)
	if err != nil {
		return err
	}
	memK, err := resource.ParseQuantity(memKeyspace)
	if err != nil {
		return err
	}
	cpuK, err := resource.ParseQuantity(cpuKeyspace)
	if err != nil {
		return err
	}
	r.Lock()
	defer r.Unlock()
	r.limit = limit{
		memTotal:    mem,
		cpuTotal:    cpu,
		memKeyspace: memK,
		cpuKeyspace: cpuK,
	}
	log.Info("resource limit updated",
		zap.String("total memory limit", mem.String()),
		zap.String("total cpu limit", cpu.String()),
		zap.String("per keyspace memory limit", memK.String()),
		zap.String("per keyspace cpu limit", cpuK.String()))

	metrics.ResourceLimit.With(prometheus.Labels{
		metrics.ResourceTypeLabel: metrics.ResourceTypeMem,
		metrics.ScopeLabel:        metrics.ScopeTotal,
	}).Set(mem.AsApproximateFloat64())
	metrics.ResourceLimit.With(prometheus.Labels{
		metrics.ResourceTypeLabel: metrics.ResourceTypeCPU,
		metrics.ScopeLabel:        metrics.ScopeTotal,
	}).Set(cpu.AsApproximateFloat64())
	metrics.ResourceLimit.With(prometheus.Labels{
		metrics.ResourceTypeLabel: metrics.ResourceTypeMem,
		metrics.ScopeLabel:        metrics.ScopeKeyspace,
	}).Set(memK.AsApproximateFloat64())
	metrics.ResourceLimit.With(prometheus.Labels{
		metrics.ResourceTypeLabel: metrics.ResourceTypeCPU,
		metrics.ScopeLabel:        metrics.ScopeKeyspace,
	}).Set(cpuK.AsApproximateFloat64())
	return nil
}

func (r *resourceLimiter) Initialize(ctx context.Context, nodeManager nodemanager.Manager) error {
	// Parse pool spec
	poolSpec, err := r.poolManager.GetPools()
	if err != nil {
		return err
	}
	poolMemSpec := map[rm.WorkerPool]resource.Quantity{}
	poolCPUSpec := map[rm.WorkerPool]resource.Quantity{}
	for pool, spec := range poolSpec {
		poolMemLimit := spec.WorkerPodSpec.Containers[0].Resources.Limits.Memory()
		if poolMemLimit == nil {
			return fmt.Errorf("pool %s mem limit undefined", pool)
		}
		poolCPULimit := spec.WorkerPodSpec.Containers[0].Resources.Limits.Cpu()
		if poolCPULimit == nil {
			return fmt.Errorf("pool %s cpu limit undefined", pool)
		}
		poolMemSpec[pool] = *poolMemLimit
		poolCPUSpec[pool] = *poolCPULimit
	}
	// Get active request from node manager and calculate current resource usage.
	bindNodes, err := nodeManager.GetActiveNodes(ctx)
	if err != nil {
		return err
	}
	for _, bindNode := range bindNodes {
		// In the case where pool specs for the active request no longer exists due to
		// pool spec change during task execution, we skip counting that task towards usage.
		mem, ok := poolMemSpec[bindNode.GetWorkerPool()]
		if !ok {
			log.Info("failed to find mem spec for the target pool",
				zap.String("target pool", bindNode.PoolName),
				zap.Any("pool mem spec", poolMemSpec),
			)
			continue
		}
		cpu, ok := poolCPUSpec[bindNode.GetWorkerPool()]
		if !ok {
			log.Info("failed to find cpu spec for the target pool",
				zap.String("target pool", bindNode.PoolName),
				zap.Any("pool cpu spec", poolCPUSpec),
			)
		}
		keyspaceUsage, ok := r.usageRecords[bindNode.KeyspaceName]
		if !ok {
			keyspaceUsage = &keyspaceUsageRecord{
				tasks: make(map[string]*taskUsageRecord),
			}
			r.usageRecords[bindNode.KeyspaceName] = keyspaceUsage
		}
		if keyspaceUsage.tasks[bindNode.TaskID] != nil {
			log.Info("task already counted, will skip",
				zap.String("task-id", bindNode.TaskID),
				zap.String("keyspace-name", bindNode.KeyspaceName),
			)
			continue
		}
		keyspaceUsage.tasks[bindNode.TaskID] = &taskUsageRecord{
			mem: mem,
			cpu: cpu,
		}
		keyspaceUsage.mem.Add(mem)
		keyspaceUsage.cpu.Add(cpu)
		r.mem.Add(mem)
		r.cpu.Add(cpu)
	}
	log.Info("resource limiter initialized",
		zap.String("total memory usage", r.mem.String()),
		zap.String("total cpu usage", r.cpu.String()))
	return nil
}