package probe

import (
	"context"
	"sync"
	"sync/atomic"
	"time"
)

type (
	ReadinessFunc func() Readiness
	LivenessFunc  func(context.Context) Liveness

	ReadinessAggregateFuncs []ReadinessFunc
)

func (fs ReadinessAggregateFuncs) check() (a ReadinessAggregate) {
	a = make(ReadinessAggregate, 0, len(fs))
	for _, f := range fs {
		a = append(a, f())
	}

	return a
}

type Reporter interface {
	ReportReadiness() Readiness
	ReportLiveness() Liveness

	RegisterReadiness(ReadinessFunc)
	RegisterLiveness(LivenessFunc)
}

func NewReporter(livenessTimeout time.Duration) *reporter {
	return &reporter{
		livenessTimeout: livenessTimeout,
	}
}

type reporter struct {
	readinessComponents ReadinessAggregateFuncs
	livenessComponents  []LivenessFunc

	livenessTimeout time.Duration

	mu     sync.Mutex
	livemu sync.Mutex
	readmu sync.Mutex
}

func (r *reporter) ReportReadiness() Readiness {
	r.readmu.Lock()
	defer r.readmu.Unlock()

	return r.readinessComponents.check().Status()
}

func (r *reporter) ReportLiveness() (out Liveness) {
	r.livemu.Lock()
	defer r.livemu.Unlock()

	ctx, cancel := context.WithTimeout(context.Background(), r.livenessTimeout)
	defer cancel()

	for _, f := range r.livenessComponents {
		status := f(ctx)
		if status == LivenessTimeout {
			return status
		}
	}

	return LivenessOk
}

func (r *reporter) RegisterReadiness(f ReadinessFunc) {
	r.readmu.Lock()
	defer r.readmu.Unlock()

	r.readinessComponents = append(r.readinessComponents, f)
}

func (r *reporter) RegisterLiveness(f LivenessFunc) {
	r.livemu.Lock()
	defer r.livemu.Unlock()

	r.livenessComponents = append(r.livenessComponents, f)
}

func SimpleReadinessSwitcher() (f ReadinessFunc, toggle func(newStatus Readiness)) {
	var status atomic.Int32

	f = func() Readiness {
		return Readiness(status.Load())
	}

	toggle = func(newStatus Readiness) {
		status.Store(int32(newStatus))
	}

	return f, toggle
}

func SimpleLivenessSwitcher() (f LivenessFunc, toggle func()) {
	var liveness atomic.Bool

	f = func(context.Context) Liveness {
		if liveness.Load() {
			return LivenessOk
		}

		return LivenessUnknown
	}

	toggle = func() {
		liveness.Store(true)
	}

	return f, toggle
}