package taskscaler

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

	"gitlab.com/gitlab-org/fleeting/fleeting"
	flmetrics "gitlab.com/gitlab-org/fleeting/fleeting/metrics"
	"gitlab.com/gitlab-org/fleeting/fleeting/provider"

	"gitlab.com/gitlab-org/fleeting/taskscaler/internal"
	"gitlab.com/gitlab-org/fleeting/taskscaler/internal/capacity"
	"gitlab.com/gitlab-org/fleeting/taskscaler/internal/ratelimit"
	"gitlab.com/gitlab-org/fleeting/taskscaler/metrics"
)

//go:generate mockery --name Taskscaler --with-expecter=true
type Taskscaler interface {
	Shutdown(ctx context.Context)
	Schedule() Schedule
	Reserve(key string) error
	Unreserve(key string)
	Acquire(ctx context.Context, key string) (Acquisition, error)
	Get(key string) Acquisition
	Release(key string)
	Capacity() (available int, potential int)
	ConfigureSchedule(schedules ...Schedule) error
	MetricsCollector() metrics.Collector
	FleetingMetricsCollector() flmetrics.Collector
}

type taskscaler struct {
	opts        options
	group       provider.InstanceGroup
	provisioner *fleeting.Provisioner
	instances   internal.List

	mu           sync.Mutex
	acquisitions map[string]Acquisition
	reservations map[string]struct{}
	pending      int
	lastRequired int
	schedules    schedules
	active       Schedule
	shutdown     func()

	ratelimit *ratelimit.Limiter

	mc metrics.Collector
}

var _ Taskscaler = &taskscaler{}

var (
	ErrNoReservations      = errors.New("taskscaler hasn't been configured to take reservations")
	ErrReservationNotFound = errors.New("reservation not found")
	ErrNoCapacity          = errors.New("no capacity")

	ErrAcquisitionWasRelinquished = errors.New("acquisition was relinquished")
)

var _ Acquisition = &acquisition{}

//go:generate mockery --name Acquisition --with-expecter=true
type Acquisition interface {
	InstanceID() string
	InstanceConnectInfo(ctx context.Context) (provider.ConnectInfo, error)
	Slot() int
	HealthSuccess()
	HealthFailure()
	WithContext(ctx context.Context) (context.Context, context.CancelFunc)
}

type acquisition struct {
	id                      string
	slot                    int
	info                    func(context.Context) (provider.ConnectInfo, error)
	relinquishCallback      func()
	healthSuccessCallback   func()
	healthFailureCallback   func()
	instanceContextCallback func(context.Context) (context.Context, context.CancelFunc)

	groupCtx *internal.CancellableGroupContext
}

func (a *acquisition) InstanceID() string {
	return a.id
}

func (a *acquisition) InstanceConnectInfo(ctx context.Context) (provider.ConnectInfo, error) {
	ctx, cancel := a.WithContext(ctx)
	defer cancel()

	info, err := a.info(ctx)
	if cause := context.Cause(ctx); cause != nil {
		return info, cause
	}

	return info, err
}

func (a *acquisition) Slot() int {
	return a.slot
}

func (a *acquisition) relinquish() {
	a.groupCtx.Cancel(ErrAcquisitionWasRelinquished)

	a.relinquishCallback()
}

func (a *acquisition) HealthSuccess() {
	a.healthSuccessCallback()
}

func (a *acquisition) HealthFailure() {
	a.healthFailureCallback()
}

func (a *acquisition) WithContext(parent context.Context) (context.Context, context.CancelFunc) {
	instanceCtx, instanceCtxCancel := a.instanceContextCallback(parent)
	ctx, cancel := a.groupCtx.WithContext(instanceCtx)

	return ctx, func() {
		cancel()
		instanceCtxCancel()
	}
}

// New returns and initializes a new Taskscaler. The context passed here is
// used only for initialization.
// It is recommended to pass a timeout context, as this call will block until
// the reconcilation process between fleeting and the instance group has
// complete.
// Shutdown() can be called to stop both the taskscaler and fleeting.
func New(ctx context.Context, group provider.InstanceGroup, options ...Option) (Taskscaler, error) {
	opts, err := newOptions(options)
	if err != nil {
		return nil, fmt.Errorf("client option: %w", err)
	}

	ts := &taskscaler{
		opts:         opts,
		group:        group,
		acquisitions: make(map[string]Acquisition),
		reservations: make(map[string]struct{}),
		active:       defaultSchedule,
		mc:           metrics.Init(opts.metricsCollector, opts.maxUseCount, opts.capacityPerInstance),
		ratelimit:    ratelimit.New(opts.scaleThrottleLimit, opts.scaleThrottleBurst),
	}

	limit, burst := ts.ratelimit.LimitBurst()
	opts.logger.Info("instance rate limiting", "limit", limit, "burst", burst)

	return ts, ts.start(ctx)
}

func (ts *taskscaler) start(initCtx context.Context) error {
	mainLoopCtx, cancel := context.WithCancel(context.Background())
	ts.shutdown = cancel

	opts := []fleeting.Option{
		fleeting.WithMaxSize(ts.opts.maxInstances),
		fleeting.WithInstanceGroupSettings(ts.opts.settings),
		fleeting.WithSubscriber(func(instances []fleeting.Instance) {
			ts.updates(mainLoopCtx, instances)
		}),
		fleeting.WithMetricsCollector(ts.opts.fleetingMetricsCollector),
		fleeting.WithUpdateInterval(ts.opts.updateInterval),
		fleeting.WithUpdateIntervalWhenExpecting(ts.opts.updateIntervalWhenExpecting),
		fleeting.WithDeletionRetryInterval(ts.opts.deletionRetryInterval),
		fleeting.WithShutdownDeletionRetries(ts.opts.shutdownDeletionRetries),
		fleeting.WithShutdownDeletionInterval(ts.opts.shutdownDeletionInterval),
	}

	provisioner, err := fleeting.Init(initCtx, ts.opts.logger, ts.group, opts...)
	if err != nil {
		return fmt.Errorf("initializing provisioner: %w", err)
	}

	ts.provisioner = provisioner
	go func() {
		for {
			if mainLoopCtx.Err() != nil {
				return
			}

			time.Sleep(time.Second)
			ts.scale(ts.desired(mainLoopCtx))
		}
	}()

	return nil
}

func (ts *taskscaler) Shutdown(ctx context.Context) {
	ts.opts.logger.Info("received shutdown signal; stopping taskscaler's loop and triggering provisioner shutdown")

	ts.opts.logger.Info("stopping main loop")
	ts.shutdown()

	ts.waitForInstanceReleases(ctx)
	ts.removeAllInstances()

	ts.opts.logger.Info("shutting down provisioner")
	ts.provisioner.Shutdown(ctx)

	ts.opts.logger.Info("taskscaler's shutdown completed")
}

func (ts *taskscaler) waitForInstanceReleases(ctx context.Context) {
	ts.opts.logger.Info("waiting for acquired instances to be released")

	for {
		select {
		case <-ctx.Done():
			return
		default:
			if len(ts.acquisitions) < 1 {
				return
			}

			time.Sleep(100 * time.Millisecond)
		}
	}
}

func (ts *taskscaler) removeAllInstances() {
	if !ts.opts.deleteInstancesOnShutdown {
		return
	}

	ts.opts.logger.Info("removing up all instances")

	for _, instance := range ts.instances.List() {
		if instance.RemoveIfExpired(time.Duration(0)) {
			ts.opts.logger.Info("instance removed", "instance", instance.ID())
		} else {
			ts.opts.logger.Warn("instance couldn't be removed", "instance", instance.ID())
		}
	}
}

func (ts *taskscaler) Schedule() Schedule {
	ts.mu.Lock()
	defer ts.mu.Unlock()

	return ts.active
}

// Acquire blocks until capacity is found and then returns an Acquisition.
//
// Blocking can be canceled by canceling the context. If a reservations are
// enabled, you must first have a reservation with the same key before calling
// Acquire(). With a reservation, Acquire() should return instantly.
//
// It is the callers responsibility to call Release() with the acquisition key
// in all circumstances when done, even if the underlying instance is removed.
//
// Acquisition.WithContext() can be used to return a context that is canceled
// should the instance be removed.
func (ts *taskscaler) Acquire(ctx context.Context, key string) (Acquisition, error) {
	var err error

	ts.mu.Lock()
	if ts.opts.reservations {
		if _, ok := ts.reservations[key]; !ok {
			err = ErrReservationNotFound
		}
	}
	if err == nil {
		ts.pending++
	}
	ts.mu.Unlock()

	if err != nil {
		return nil, err
	}

	defer func() {
		ts.mu.Lock()
		ts.pending--
		ts.mu.Unlock()
	}()

	for {
		if ctx.Err() != nil {
			ts.onMC(func(m metrics.Collector) {
				m.TaskOperationInc(metrics.TaskOperationAcquireTimeout)
			})

			return nil, ctx.Err()
		}

		acq := ts.acquire(key)
		if acq == nil {
			time.Sleep(ts.opts.acquireDelay)
			continue
		}

		ts.onMC(func(m metrics.Collector) {
			m.TaskOperationInc(metrics.TaskOperationAcquire)
		})

		return acq, nil
	}
}

func (ts *taskscaler) acquire(key string) Acquisition {
	ts.mu.Lock()
	defer ts.mu.Unlock()

	for _, instance := range ts.instances.List() {
		slot, ok := instance.Acquire(key)
		if !ok {
			continue
		}

		ts.acquisitions[key] = &acquisition{
			id:   instance.ID(),
			slot: slot,
			info: instance.ConnectInfo,
			relinquishCallback: func() {
				instance.Relinquish(slot)
			},
			healthSuccessCallback:   instance.HealthSuccess,
			healthFailureCallback:   instance.HealthFailure,
			instanceContextCallback: instance.WithContext,
			groupCtx:                internal.NewCancellableGroupContext(),
		}
		delete(ts.reservations, key)

		return ts.acquisitions[key]
	}

	return nil
}

// Get return an Acquisition by its key;.
func (ts *taskscaler) Get(key string) Acquisition {
	ts.mu.Lock()
	defer ts.mu.Unlock()

	return ts.acquisitions[key]
}

// Release releases an acquisition by its key.
func (ts *taskscaler) Release(key string) {
	ts.mu.Lock()
	defer ts.mu.Unlock()

	acq, ok := ts.acquisitions[key]
	if !ok {
		return
	}

	if acq, ok := acq.(*acquisition); ok {
		acq.relinquish()
	}
	delete(ts.acquisitions, key)

	ts.onMC(func(m metrics.Collector) {
		m.TaskOperationInc(metrics.TaskOperationRelease)
	})
}

// Reserve makes a reservation for capacity and provides some guarantees
// that a call to Acquire(), using the same key, will succeed.
//
// Reserve only works when Taskscaler has been configured 'WithReservations'.
//
// When Preemptive Mode is enabled on the active schedule, the guarantee
// that Acquire() will succeed is even greater, as reservations can only
// be made when there's a known immediate capacity to cover it.
//
// Unlike calls to Acquire(), Reserve() doesn't actively cause any scaling
// event to accommodate demand. However, reservations are taken into account
// when scaling down and capacity won't be removed if there's the potential
// it can be used for reserved capacity.
//
// Reserve only returns an error if there's no capacity left for a
// reservation (ErrNoCapacity). When Preemptive Mode is enabled on the active
// schedule, ErrNoCapacity is returned if there's no immediately available
// capacity.
func (ts *taskscaler) Reserve(key string) error {
	if !ts.opts.reservations {
		return ErrNoReservations
	}

	ts.mu.Lock()
	defer ts.mu.Unlock()

	_, ok := ts.reservations[key]
	if ok {
		return nil
	}

	available, potential := ts.capacity()

	// cannot reserve if there's no immediately available or potential capacity
	if potential <= 0 && available <= 0 {
		ts.onMC(func(m metrics.Collector) {
			m.TaskOperationInc(metrics.TaskOperationReserveCapacityFailure)
		})

		return fmt.Errorf("%w: no immediately available or potential capacity", ErrNoCapacity)
	}

	if ts.active.IdleCount > 0 && ts.active.PreemptiveMode && available <= 0 {
		ts.onMC(func(m metrics.Collector) {
			m.TaskOperationInc(metrics.TaskOperationReserveAvailableCapacityFailure)
		})

		return fmt.Errorf("%w: no immediately available capacity", ErrNoCapacity)
	}

	ts.reservations[key] = struct{}{}

	ts.onMC(func(m metrics.Collector) {
		m.TaskOperationInc(metrics.TaskOperationReserve)
	})

	return nil
}

// Unreserve cancels a reservation that was previously made with Reserve()
func (ts *taskscaler) Unreserve(key string) {
	ts.mu.Lock()
	defer ts.mu.Unlock()

	delete(ts.reservations, key)

	ts.onMC(func(m metrics.Collector) {
		m.TaskOperationInc(metrics.TaskOperationUnreserve)
	})
}

// Capacity returns available (immediate) capacity and potential (on-demand)
// capacity.
func (ts *taskscaler) Capacity() (available int, potential int) {
	ts.mu.Lock()
	defer ts.mu.Unlock()

	return ts.capacity()
}

func (ts *taskscaler) capacity() (available int, potential int) {
	capacity := ts.provisioner.Capacity()
	idle, acquired, unavailable := ts.instances.Capacity()
	reservations := len(ts.reservations)

	available = idle - reservations

	if available < 0 {
		// some of the reservations will be chipped away by available capacity,
		// so we adjust how many are left and take them away from potential
		// capacity below
		reservations = -available

		// it's possible in some configurations that we allow reserving
		// more than is immediately available
		available = 0
	}

	// we treat acquired and reservations also unavailable
	unavailable += acquired + reservations

	// if there's no instance cap, we make the total potential capacity
	// effectively unlimited
	total := math.MaxInt
	if capacity.Max > 0 {
		// sometimes we get unexpected instances and those that are running
		// can already exceed our maximum capacity, but if we have them, we might
		// as well use them
		total = max(capacity.Max, capacity.Running) * ts.opts.capacityPerInstance
	}

	// potential capacity is the total capacity we have minus unavailable capacity
	potential = total - unavailable

	if potential < 0 {
		inflight := (capacity.Running + capacity.Deleting) * ts.opts.capacityPerInstance
		if inflight > total && inflight-unavailable < 0 {
			// known edge case: capacity falling below zero can occur if we scaled beyond the
			// max limit (out-of-band), had slots acquired, but now some instances with
			// those slots are now being deleted
			return available, 0
		}

		ts.opts.logger.Error("capacity potential below zero", "total", total, "acquired", acquired, "unavailable", unavailable, "reservations", len(ts.reservations))

		return available, 0
	}

	return available, potential
}

func (ts *taskscaler) updates(ctx context.Context, instances []fleeting.Instance) {
	for _, instance := range instances {
		switch instance.State() {
		case provider.StateCreating, provider.StateDeleting:
			ts.mu.Lock()
			inst := ts.instances.Get(instance.ID())
			ts.mu.Unlock()

			// if creating/deleting then unready the instance
			// this allows for instances to be re-prepared should they come back
			if inst != nil {
				inst.Unready()
			}

		case provider.StateRunning:
			ts.mu.Lock()
			inst := ts.instances.Get(instance.ID())
			if inst == nil {
				inst = internal.New(instance, ts.opts.logger, ts.opts.capacityPerInstance, ts.opts.maxUseCount, ts.opts.failureThreshold)
				ts.instances.Add(inst)
			}
			ts.mu.Unlock()

			inst.Prepare(ctx, instance.Cause(), internal.UpFunc(ts.opts.upFn), func(duration time.Duration, err error) {
				if err != nil {
					ts.ratelimit.Failure()
					return
				}

				ts.ratelimit.Success()

				ts.onMC(func(m metrics.Collector) {
					m.TaskInstanceReadinessTimeObserve(duration)
				})
			})

		case provider.StateDeleted, provider.StateTimeout:
			inst := ts.instances.Get(instance.ID())
			if inst == nil {
				continue
			}

			inst.RemovedExternally()
			ts.instances.Delete(instance.ID())
		}
	}
}

func (ts *taskscaler) desired(ctx context.Context) int {
	c := ts.provisioner.Capacity()
	idle, acquired, unavailable := ts.instances.Capacity()

	ts.mu.Lock()
	pending := ts.pending
	reserved := len(ts.reservations)
	ts.active = ts.schedules.active(time.Now())
	active := ts.active
	ts.mu.Unlock()

	ts.onMC(func(m metrics.Collector) {
		m.MaxUseCountSet(ts.opts.maxUseCount)
		m.MaxTasksPerInstanceSet(ts.opts.capacityPerInstance)

		m.TasksCountSet(metrics.TasksCountIdle, idle)
		m.TasksCountSet(metrics.TasksCountPending, pending)
		m.TasksCountSet(metrics.TasksCountAcquired, acquired)
		m.TasksCountSet(metrics.TasksCountReserved, reserved)
		m.TasksCountSet(metrics.TasksCountUnavailable, unavailable)
	})

	capacityInfo := capacity.CapacityInfo{
		InstanceCount:       c.Creating + c.Running + c.Requested,
		MaxInstanceCount:    c.Max,
		Acquired:            acquired,
		UnavailableCapacity: unavailable,
		Pending:             pending,
		Reserved:            reserved,
		IdleCount:           active.IdleCount,
		ScaleFactor:         active.ScaleFactor,
		ScaleFactorLimit:    active.ScaleFactorLimit,
		CapacityPerInstance: ts.opts.capacityPerInstance,
	}
	required := capacity.RequiredInstances(capacityInfo)
	if required != ts.lastRequired {
		ts.opts.logger.Info("required scaling change", "required", required, "capacity-info", capacityInfo.String())
	}
	ts.lastRequired = required
	return required
}

func (ts *taskscaler) scale(n int) {
	ts.onMC(func(m metrics.Collector) {
		m.DesiredInstancesSet(n)
	})

	if n == 0 {
		return
	}

	if n > 0 {
		rateLimitedN := ts.ratelimit.N(n)
		ts.opts.logger.Debug("scale up", "n", n, "rate-limit-n", rateLimitedN)

		if rateLimitedN == 0 {
			delay, newDelay := ts.ratelimit.Delay()
			if delay > 0 && newDelay {
				ts.opts.logger.Warn("backing off scale up due to ready up failures", "backoff", delay)
			}

			return
		}

		if rateLimitedN != n {
			ts.opts.logger.Info("rate limited scale up", "request", n, "limit", rateLimitedN)
		}

		ts.provisioner.Request(rateLimitedN)

		ts.onMC(func(m metrics.Collector) {
			m.ScaleOperationInc(metrics.ScaleOperationUp)
		})

		return
	}

	ts.opts.logger.Debug("scale down", "n", n)

	ts.onMC(func(m metrics.Collector) {
		m.ScaleOperationInc(metrics.ScaleOperationDown)
	})

	ts.mu.Lock()
	active := ts.active
	ts.mu.Unlock()

	for _, instance := range ts.instances.List() {
		if n == 0 {
			break
		}

		if instance.RemoveIfExpired(active.IdleTime) {
			n++
		}
	}
}

func (ts *taskscaler) onMC(mFn func(m metrics.Collector)) {
	if ts.mc != nil {
		mFn(ts.mc)
	}
}

func (ts *taskscaler) MetricsCollector() metrics.Collector {
	return ts.mc
}

func (ts *taskscaler) FleetingMetricsCollector() flmetrics.Collector {
	return ts.opts.fleetingMetricsCollector
}
