package main

import (
	"sync"
	"time"

	"go.uber.org/zap"
)

type drainController struct {
	readiness      *readiness
	readinessDelay time.Duration
	beginDrain     func() int
	log            *zap.Logger

	once    sync.Once
	beginCh chan struct{}
}

func newDrainController(rdy *readiness, readinessDelay time.Duration, beginDrain func() int, log *zap.Logger) *drainController {
	if log == nil {
		log = zap.NewNop()
	}
	return &drainController{
		readiness:      rdy,
		readinessDelay: readinessDelay,
		beginDrain:     beginDrain,
		log:            log,
		beginCh:        make(chan struct{}),
	}
}

func (d *drainController) Enter(reason string) {
	if d == nil {
		return
	}
	d.once.Do(func() {
		if d.readiness != nil {
			d.readiness.MarkDraining()
		}
		d.log.Info("Proxy marked not ready for drain",
			zap.String("reason", reason),
			zap.Duration("readiness_delay", d.readinessDelay))
		go d.begin(reason)
	})
}

func (d *drainController) WaitForBeginDrain() {
	if d == nil {
		return
	}
	<-d.beginCh
}

func (d *drainController) begin(reason string) {
	defer close(d.beginCh)

	if d.readinessDelay > 0 {
		timer := time.NewTimer(d.readinessDelay)
		defer timer.Stop()
		<-timer.C
	}

	active := 0
	if d.beginDrain != nil {
		active = d.beginDrain()
	}

	d.log.Info("Proxy entered drain mode",
		zap.String("reason", reason),
		zap.Duration("readiness_delay", d.readinessDelay),
		zap.Int("active_connections", active))
}