package internal

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

	"github.com/hashicorp/go-hclog"

	"gitlab.com/gitlab-org/fleeting/fleeting"
	"gitlab.com/gitlab-org/fleeting/fleeting/provider"
)

type Instance struct {
	instance fleeting.Instance
	logger   hclog.Logger

	capacityPerInstance int
	maxUseCount         int
	failureThreshold    int

	mu                      sync.Mutex
	info                    provider.ConnectInfo
	acquiredSlots           map[int]string
	used                    int
	preparing               bool
	ready                   bool
	removing                bool
	consecutiveFailureCount int
	lastUsedAt              time.Time

	groupCtx *CancellableGroupContext
}

var (
	ErrInstanceReachedMaxUseCount       = errors.New("instance reached max use count")
	ErrInstanceUnhealthy                = errors.New("instance is unhealthy")
	ErrInstanceExceededMaxIdleTime      = errors.New("instance exceeded max idle time")
	ErrInstancePreparationFailed        = errors.New("instance connection preparation failed")
	ErrInstanceReadyUpPreparationFailed = errors.New("instance ready up preparation failed")
	ErrInstanceExternallyRemoved        = errors.New("instance externally removed")
)

func New(instance fleeting.Instance, logger hclog.Logger, capacityPerInstance, maxUseCount, failureThreshold int) *Instance {
	logger = logger.With("instance", instance.ID())

	return &Instance{
		instance:            instance,
		logger:              logger,
		capacityPerInstance: capacityPerInstance,
		maxUseCount:         maxUseCount,
		failureThreshold:    failureThreshold,
		acquiredSlots:       map[int]string{},
		groupCtx:            NewCancellableGroupContext(),
	}
}

func (inst *Instance) ID() string {
	return inst.instance.ID()
}

func (inst *Instance) Acquire(key string) (slot int, success bool) {
	inst.mu.Lock()
	defer inst.mu.Unlock()

	switch {
	case inst.removing, !inst.ready:
		return 0, false

	case inst.reachedCapacityPerInstance():
		return 0, false

	case inst.reachedMaxUseCount():
		return 0, false
	}

	inst.used++
	for i := 0; i < inst.capacityPerInstance; i++ {
		if _, occupied := inst.acquiredSlots[i]; !occupied {
			inst.acquiredSlots[i] = key
			return i, true
		}
	}

	panic("no remaining slots")
}

func (inst *Instance) reachedMaxUseCount() bool {
	return inst.maxUseCount > 0 && inst.used >= inst.maxUseCount
}

func (inst *Instance) reachedCapacityPerInstance() bool {
	return len(inst.acquiredSlots) >= inst.capacityPerInstance
}

func (inst *Instance) unhealthy() bool {
	return inst.consecutiveFailureCount >= inst.failureThreshold
}

func (inst *Instance) Relinquish(slot int) {
	inst.mu.Lock()
	defer inst.mu.Unlock()
	if _, ok := inst.acquiredSlots[slot]; ok {
		delete(inst.acquiredSlots, slot)
	} else {
		panic("relinquish of slot not acquired")
	}

	inst.lastUsedAt = time.Now()
	if inst.reachedMaxUseCount() && len(inst.acquiredSlots) == 0 {
		inst.remove(ErrInstanceReachedMaxUseCount)
	}
}

func (inst *Instance) WithContext(ctx context.Context) (context.Context, context.CancelFunc) {
	return inst.groupCtx.WithContext(ctx)
}

// Capacity returns the acquired and unavailable capacity of an instance.
//
// The unavailable capacity is typically 0, indicating that all capacity is
// available (other than whatever the acquired capacity is).
//
// The unavailable capacity is above zero when a max used count is configured,
// and the used count is affecting the capacity:
//   - if you have a capacity of 5 and a used count of 1/100, you still have a
//     capacity of 5. The unavailable capacity is therefore 0.
//   - if you have a capacity of 5 and a used count of 96/100, you now only
//     have a capacity of 4. The unavailable capacity is therefore 1.
func (inst *Instance) Capacity() (acquired, unavailable int, deleted bool) {
	inst.mu.Lock()
	defer inst.mu.Unlock()

	if !inst.instance.DeletedAt().IsZero() {
		return 0, 0, true
	}

	acquired = len(inst.acquiredSlots)

	if inst.removing {
		return acquired, inst.capacityPerInstance - acquired, false
	}

	if inst.maxUseCount <= 0 {
		return acquired, 0, false
	}

	if !inst.removing && inst.maxUseCount > 0 {
		// used and acquired are incremented together, and we don't want
		// the unavailable count to include the current acquired count,
		// so we deduct it here.
		remainingUses := inst.maxUseCount - (inst.used - acquired)
		if remainingUses < inst.capacityPerInstance {
			unavailable = inst.capacityPerInstance - remainingUses
		}
	}

	return acquired, unavailable, false
}

func (inst *Instance) IsReady() bool {
	inst.mu.Lock()
	defer inst.mu.Unlock()

	return inst.ready
}

func (inst *Instance) Run(fn func(info provider.ConnectInfo) error) error {
	return fn(inst.info)
}

// ConnectInfo returns the connection information for the instance.
//
// If the connection info has expired, it will be attempted to be refreshed. If
// that attempt fails, the error will be returned along with the expired
// connection info. On success, the new connection info will be cached for
// subsequent calls. If HealthFailure() is called, the connection cache will
// be invalidated.
func (inst *Instance) ConnectInfo(ctx context.Context) (provider.ConnectInfo, error) {
	inst.mu.Lock()
	defer inst.mu.Unlock()

	// the connection cache has expired if expiration is set and we're 10 seconds away
	// from expiration
	expired := inst.info.Expires != nil && inst.info.Expires.Before(time.Now().Add(10*time.Second))

	// refresh if connection has errors or cache has expired
	if inst.consecutiveFailureCount > 0 || expired {
		info, err := inst.instance.ConnectInfo(ctx)
		if err != nil {
			return inst.info, fmt.Errorf("refreshing connect info: %w", err)
		}
		inst.info = info
	}

	return inst.info, nil
}

type UpFunc func(id string, info provider.ConnectInfo, cause fleeting.Cause) (keys []string, used int, err error)

type readyFunc func(time.Duration, error)

func (inst *Instance) Prepare(ctx context.Context, cause fleeting.Cause, upFunc UpFunc, readyFn readyFunc) {
	inst.mu.Lock()
	defer inst.mu.Unlock()

	if inst.removing {
		return
	}

	if inst.preparing || inst.ready {
		return
	}

	inst.preparing = true

	go func() {
		defer func() {
			inst.mu.Lock()
			inst.preparing = false
			inst.mu.Unlock()
		}()

		start := time.Now()
		ready := func(msg string, err error) {
			preparationTime := time.Since(start)
			if err == nil {
				inst.logger.Info(msg, "took", preparationTime)
			} else {
				inst.logger.Error(msg, "took", preparationTime, "err", err)
			}

			if readyFn != nil {
				readyFn(preparationTime, err)
			}
		}

		ctx, cancel := inst.WithContext(ctx)
		defer cancel()

		if !inst.prepareWait(ctx) {
			return
		}

		info, err := inst.prepareConnect(ctx)
		if err != nil {
			ready("connection preparation failed", err)
			inst.Remove(ErrInstancePreparationFailed)
			return
		}
		inst.info = info

		if upFunc != nil {
			keys, used, err := upFunc(inst.ID(), info, cause)
			if err != nil {
				ready("ready up preparation failed", err)
				inst.Remove(ErrInstanceReadyUpPreparationFailed)
				return
			}

			inst.mu.Lock()

			acquiredSlotsByKey := map[string]int{}
			for slot, key := range inst.acquiredSlots {
				acquiredSlotsByKey[key] = slot
			}
			keysWithoutSlot := []string{}
			for _, key := range keys {
				if _, ok := acquiredSlotsByKey[key]; !ok {
					keysWithoutSlot = append(keysWithoutSlot, key)
				}
			}
			if len(keysWithoutSlot) > 0 {
				// If the underlying resource manager
				// (e.g. nesting) has slots for the
				// extra keys, then taskscaler will
				// stomp them when reassigning the
				// slot (good and expected). If not,
				// they will remain as orphan
				// resources (bad).
				inst.logger.Warn(
					"the following keys are present on instance %q but are not assigned slots in taskscaler: %v: %v",
					inst.ID(),
					keysWithoutSlot,
				)
			}

			inst.used = used
			inst.mu.Unlock()
		}

		inst.mu.Lock()
		inst.ready = true
		inst.mu.Unlock()

		ready("instance is ready", nil)
	}()
}

func (inst *Instance) prepareWait(ctx context.Context) bool {
	ctx, cancel := context.WithTimeout(ctx, time.Minute)
	defer cancel()

	return inst.instance.ReadyWait(ctx)
}

func (inst *Instance) prepareConnect(ctx context.Context) (provider.ConnectInfo, error) {
	return inst.instance.ConnectInfo(ctx)
}

func (inst *Instance) Unready() {
	inst.mu.Lock()
	defer inst.mu.Unlock()

	if !inst.ready {
		return
	}

	inst.ready = false
	inst.consecutiveFailureCount = 0
}

// RemoveIfExpired determines if an instances should be deleted according to
// several rules and then deletes the instance.
func (inst *Instance) RemoveIfExpired(minIdleTime time.Duration) bool {
	inst.mu.Lock()
	defer inst.mu.Unlock()

	// don't expire acquired instances
	if len(inst.acquiredSlots) > 0 {
		return false
	}

	if inst.reachedMaxUseCount() {
		inst.remove(ErrInstanceReachedMaxUseCount)
		return true
	}

	if inst.unhealthy() {
		inst.remove(ErrInstanceUnhealthy)
		return true
	}

	// use the last idle time if available, otherwise,
	// the time the instance was provisioned at
	idleTime := inst.instance.ProvisionedAt()
	if inst.lastUsedAt.After(idleTime) {
		idleTime = inst.lastUsedAt
	}

	// keep the instance if it hasn't yet exceeded the minimum idle time
	if time.Since(idleTime) < minIdleTime {
		return false
	}

	inst.remove(ErrInstanceExceededMaxIdleTime)

	return true
}

// Remove marks an instance for removal
func (inst *Instance) Remove(reason error) {
	inst.mu.Lock()
	defer inst.mu.Unlock()
	inst.remove(reason)
}

func (inst *Instance) remove(reason error) {
	if inst.removing {
		return
	}

	inst.logger.Info("instance marked for removal", "reason", reason)

	inst.instance.Delete()
	inst.removing = true
	inst.groupCtx.Cancel(reason)
}

// RemovedExternally is used to signal that this instance no longer exists
// according to the instance group.
func (inst *Instance) RemovedExternally() {
	inst.groupCtx.Cancel(ErrInstanceExternallyRemoved)

	inst.mu.Lock()
	defer inst.mu.Unlock()

	if inst.removing {
		return
	}

	inst.logger.Error("instance unexpectedly removed", "slots", inst.acquiredSlots, "used", inst.used, "max-use-count", inst.maxUseCount)
}

// HealthSuccess reports a successful connection verifying health of
// the instance
func (inst *Instance) HealthSuccess() {
	inst.mu.Lock()
	defer inst.mu.Unlock()

	inst.logger.Info("health success", "failure-count", inst.consecutiveFailureCount)

	inst.consecutiveFailureCount = 0
}

// HealthFailure reports an unsuccessful connection verifying the
// unhealthiness of the instance
func (inst *Instance) HealthFailure() {
	inst.mu.Lock()
	defer inst.mu.Unlock()
	inst.consecutiveFailureCount++

	inst.logger.Warn("health failure", "failure-count", inst.consecutiveFailureCount)
}
