package kpoll

import (
	"context"
	"testing"
	"time"

	"gotest.tools/v3/assert/cmp"
	"gotest.tools/v3/poll"
	"k8s.io/apimachinery/pkg/api/errors"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"edge-infra.dev/pkg/k8s/runtime/inventory"
	"edge-infra.dev/pkg/k8s/testing/kmp"
)

// KPoll allows asserting K8s Objects against a live cluster's state with
// default polling options and provides access to common assertions for K8s
// objects.
//
// KPoll should be treated like other assertion helpers: instantiated for each
// individual test.
//
// When instantiated as part of a test hook, this ensures that all assertions
// use a consistent and configurable timeout by default.
type KPoll struct {
	c       client.Client
	ctx     context.Context
	timeout time.Duration
	tick    time.Duration
}

func New(ctx context.Context, c client.Client, timeout, tick time.Duration) *KPoll {
	return &KPoll{c, ctx, timeout, tick}
}

// WaitOn wraps [gotest.tools/v3/poll.WaitOn] with default options based on the
// KPoll configuration.
func (k *KPoll) WaitOn(t *testing.T, c poll.Check, opts ...poll.SettingOp) {
	t.Helper()
	poll.WaitOn(t, c, append(k.pollSettings(), opts...)...)
}

// Compare returns a [gotest.tools/v3/poll.Result] by evaluating Object o against
// Komparison c. The Object is re-fetched on each polling loop, allowing a
// standard [edge-infra.dev/pkg/k8s/testing/kmp.Komparison] to be continually
// evaluated against a cluster's state.
//
// The poll.Result return can be used to wrap a Komparison with additional logic
// in an inline poll.Check:
//
//	kwait.On(t, func(t poll.LogT) poll.Result {
//		// Evaluate check precondition
//		// Perform comparison as polling step
//		return kwait.Compare(ctx, obj, myComparison)
//	})
func (k *KPoll) Compare(o client.Object, c kmp.Komparison, options ...CheckOption) poll.Result {
	opts := k.makeOpts(options...)
	return poll.Compare(func() cmp.Result {
		if opts.fetch {
			err := opts.client.Get(opts.ctx, client.ObjectKeyFromObject(o), o)
			if err != nil && !(errors.IsGone(err) || errors.IsNotFound(err)) {
				return cmp.ResultFromError(err)
			}
		}
		return c(o)
	})
}

// CompareAll is [KPoll.Compare] for sets of objects. The comparison results are
// aggregated to produce the overall result.
func (k *KPoll) CompareAll(objs []client.Object, c kmp.Komparison, options ...CheckOption) poll.Result {
	opts := k.makeOpts(options...)
	results := make([]poll.Result, len(objs), 0)
	for i, o := range objs {
		if opts.fetch {
			err := opts.client.Get(opts.ctx, client.ObjectKeyFromObject(o), o)
			if err != nil && !(errors.IsGone(err) || errors.IsNotFound(err)) {
				return poll.Error(err)
			}
		}
		results[i] = poll.Compare(func() cmp.Result { return c(o) })
	}

	return JoinResults(results...)
}

// Check creates a [gotest.tools/v3/poll.Check] which evaluates Object o using
// Komparison c on each interval of a [KPoll.WaitOn] call. It uses [KPoll.Compare]
// to wrap the Komparison into a function that produces a polling result so that
// the check is always executed against the cluster's live state.
func (k *KPoll) Check(o client.Object, c kmp.Komparison, options ...CheckOption) poll.Check {
	return func(_ poll.LogT) poll.Result {
		return k.Compare(o, c, options...)
	}
}

// CheckAll is [KPoll.Check] for sets of objects.
func (k *KPoll) CheckAll(objs []client.Object, c kmp.Komparison, options ...CheckOption) poll.Check {
	return func(_ poll.LogT) poll.Result {
		return k.CompareAll(objs, c, options...)
	}
}

// ObjExists checks that a single object exists on the cluster.
//
// [edge-infra.dev/pkg/k8s/object.IsDeleted] is used to determine object
// existence.
func (k *KPoll) ObjExists(o client.Object, options ...CheckOption) poll.Check {
	opts := k.makeOpts(options...)
	return Exists(opts.ctx, opts.client, o)
}

// ObjsExist checks that all objects exist on the cluster.
//
// [edge-infra.dev/pkg/k8s/object.IsDeleted] is used to determine object
// existence.
func (k *KPoll) ObjsExist(objs []client.Object, options ...CheckOption) poll.Check {
	opts := k.makeOpts(options...)
	return Exists(opts.ctx, opts.client, objs...)
}

// ObjDeleted checks that a single object doesn't exist on the cluster.
//
// [edge-infra.dev/pkg/k8s/object.IsDeleted] is used to determine object
// existence.
func (k *KPoll) ObjDeleted(o client.Object, options ...CheckOption) poll.Check {
	opts := k.makeOpts(options...)
	return Deleted(opts.ctx, opts.client, o)
}

// ObjsDeleted checks that none of the provided objects exist on the cluster.
//
// [edge-infra.dev/pkg/k8s/object.IsDeleted] is used to determine object
// existence.
func (k *KPoll) ObjsDeleted(objs []client.Object, options ...CheckOption) poll.Check {
	opts := k.makeOpts(options...)
	return Deleted(opts.ctx, opts.client, objs...)
}

// InventoryExists checks that all member objects of the expected inventory
// exist on the cluster.
//
// [edge-infra.dev/pkg/k8s/object.IsDeleted] is used to determine object
// existence.
func (k *KPoll) InventoryExists(exp *inventory.ResourceInventory, options ...CheckOption) poll.Check {
	opts := k.makeOpts(options...)
	return InventoryExists(opts.ctx, opts.client, exp)
}

func (k *KPoll) InventoryPruned(curr, old *inventory.ResourceInventory, options ...CheckOption) poll.Check {
	opts := k.makeOpts(options...)
	return InventoryPruned(opts.ctx, opts.client, curr, old)
}

func (k *KPoll) makeOpts(options ...CheckOption) *checkOpts {
	opts := &checkOpts{
		ctx:    k.ctx,
		client: k.c,
		fetch:  true,
	}
	for _, option := range options {
		option(opts)
	}
	return opts
}

// pollSettings returns the default [gotest.tools/v3/poll.WaitOn] options based
// on the object's state.
func (k *KPoll) pollSettings() []poll.SettingOp {
	return []poll.SettingOp{poll.WithDelay(k.tick), poll.WithTimeout(k.timeout)}
}