// Copyright (c) 2026 Sidero Labs, Inc.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.

package omni_test

import (
	"context"
	"testing"
	"time"

	"github.com/cosi-project/runtime/pkg/controller"
	"github.com/cosi-project/runtime/pkg/controller/generic/qtransform"
	"github.com/cosi-project/runtime/pkg/resource"
	"github.com/cosi-project/runtime/pkg/resource/rtestutils"
	"github.com/cosi-project/runtime/pkg/safe"
	"github.com/cosi-project/runtime/pkg/state"
	"github.com/google/uuid"
	"github.com/siderolabs/go-retry/retry"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"

	"github.com/siderolabs/omni/client/api/omni/specs"
	"github.com/siderolabs/omni/client/pkg/omni/resources/infra"
	"github.com/siderolabs/omni/client/pkg/omni/resources/omni"
	"github.com/siderolabs/omni/client/pkg/omni/resources/system"
	"github.com/siderolabs/omni/internal/backend/runtime/omni/controllers/helpers"
	omnictrl "github.com/siderolabs/omni/internal/backend/runtime/omni/controllers/omni"
	"github.com/siderolabs/omni/internal/backend/runtime/omni/controllers/testutils"
	"github.com/siderolabs/omni/internal/backend/runtime/omni/controllers/testutils/rmock"
	"github.com/siderolabs/omni/internal/backend/runtime/omni/controllers/testutils/rmock/options"
)

type testInfraProvider = qtransform.QController[*infra.MachineRequest, *infra.MachineRequestStatus]

func newTestInfraProvider() *testInfraProvider {
	return qtransform.NewQController(
		qtransform.Settings[*infra.MachineRequest, *infra.MachineRequestStatus]{
			Name: "testInfraProvider",
			MapMetadataFunc: func(request *infra.MachineRequest) *infra.MachineRequestStatus {
				return infra.NewMachineRequestStatus(request.Metadata().ID())
			},
			UnmapMetadataFunc: func(status *infra.MachineRequestStatus) *infra.MachineRequest {
				return infra.NewMachineRequest(status.Metadata().ID())
			},
			TransformExtraOutputFunc: func(_ context.Context, _ controller.ReaderWriter, _ *zap.Logger, machineRequest *infra.MachineRequest,
				machineRequestStatus *infra.MachineRequestStatus,
			) error {
				if machineRequestStatus.TypedSpec().Value.Id == "" {
					machineRequestStatus.TypedSpec().Value.Id = uuid.New().String()
				}

				machineRequestStatus.TypedSpec().Value.Stage = specs.MachineRequestStatusSpec_PROVISIONED

				helpers.CopyAllLabels(machineRequest, machineRequestStatus)

				return nil
			},
		},
	)
}

func TestMachineRequestSetStatusReconcile(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(t.Context(), 20*time.Second)
	t.Cleanup(cancel)

	testutils.WithRuntime(ctx, t, testutils.TestOptions{},
		func(_ context.Context, tc testutils.TestContext) {
			require.NoError(t, tc.Runtime.RegisterQController(omnictrl.NewMachineRequestSetStatusController()))
			require.NoError(t, tc.Runtime.RegisterQController(newTestInfraProvider()))
		},
		func(ctx context.Context, tc testutils.TestContext) {
			testMachineRequestSetStatusReconcile(ctx, t, tc.State)
		},
	)
}

func testMachineRequestSetStatusReconcile(ctx context.Context, t *testing.T, st state.State) {
	var eg errgroup.Group

	reconcilerCtx, reconcilerCancel := context.WithCancel(ctx)

	defer func() {
		reconcilerCancel()

		require.NoError(t, eg.Wait())
	}()

	watchReady := make(chan struct{})
	errCh := make(chan error, 1)

	eg.Go(func() error {
		err := reconcileLabels(reconcilerCtx, st, watchReady)
		if err != nil {
			t.Logf("machine status labels reconciler crashed with error %s", err)

			select {
			case errCh <- err:
			default:
			}
		}

		return err
	})

	select {
	case <-watchReady:
	case err := <-errCh:
		require.FailNow(t, "reconcileLabels failed before watch was established", err)
	case <-ctx.Done():
		require.FailNow(t, "context canceled waiting for reconcileLabels watch to be established")
	}

	machineRequestSet := rmock.Mock[*omni.MachineRequestSet](ctx, t, st,
		options.WithID("test"),
		options.Modify(func(r *omni.MachineRequestSet) error {
			r.TypedSpec().Value.ProviderId = "test"
			r.TypedSpec().Value.TalosVersion = "v1.7.5"

			return nil
		}),
	)

	var err error

	// scale up
	machineRequestSet, err = safe.StateUpdateWithConflicts(ctx, st, machineRequestSet.Metadata(), func(r *omni.MachineRequestSet) error {
		r.TypedSpec().Value.MachineCount = 4

		return nil
	})

	require.NoError(t, err)

	var ids []resource.ID

	err = retry.Constant(time.Second*5).RetryWithContext(ctx, func(ctx context.Context) error {
		ids = rtestutils.ResourceIDs[*infra.MachineRequest](ctx, t, st, state.WithLabelQuery(resource.LabelEqual(omni.LabelMachineRequestSet, machineRequestSet.Metadata().ID())))

		if len(ids) != int(machineRequestSet.TypedSpec().Value.MachineCount) {
			return retry.ExpectedErrorf("expected %d requests got %d", machineRequestSet.TypedSpec().Value.MachineCount, len(ids))
		}

		return nil
	})

	require.NoError(t, err)

	rtestutils.AssertResources(ctx, t, st, ids, func(r *infra.MachineRequest, assert *assert.Assertions) {
		l, ok := r.Metadata().Labels().Get(omni.LabelMachineRequestSet)

		assert.True(ok)
		assert.Equal(l, machineRequestSet.Metadata().ID())

		l, ok = r.Metadata().Labels().Get(omni.LabelInfraProviderID)

		assert.True(ok)
		assert.Equal(l, machineRequestSet.TypedSpec().Value.ProviderId)

		assert.Equal(machineRequestSet.TypedSpec().Value.TalosVersion, r.TypedSpec().Value.TalosVersion)
	})

	rtestutils.AssertResources(ctx, t, st, ids, func(*infra.MachineRequestStatus, *assert.Assertions) {})

	requestStatuses, err := safe.ReaderListAll[*infra.MachineRequestStatus](ctx, st,
		state.WithLabelQuery(resource.LabelEqual(omni.LabelMachineRequestSet, machineRequestSet.Metadata().ID())),
	)

	require.NoError(t, err)

	machineIDs := []string{requestStatuses.Get(0).TypedSpec().Value.Id}

	// remove the first request link
	rtestutils.AssertResources(ctx, t, st, machineIDs, func(r *system.ResourceLabels[*omni.MachineStatus], assertion *assert.Assertions) {
		assertion.True(r.Metadata().Finalizers().Has(omnictrl.MachineRequestSetStatusControllerName), r.Metadata().ID())
	})

	rtestutils.Destroy[*system.ResourceLabels[*omni.MachineStatus]](ctx, t, st, machineIDs)

	rtestutils.AssertNoResource[*infra.MachineRequest](ctx, t, st, ids[0])

	err = retry.Constant(time.Second*5).RetryWithContext(ctx, func(ctx context.Context) error {
		ids = rtestutils.ResourceIDs[*infra.MachineRequest](ctx, t, st, state.WithLabelQuery(resource.LabelEqual(omni.LabelMachineRequestSet, machineRequestSet.Metadata().ID())))

		if len(ids) != int(machineRequestSet.TypedSpec().Value.MachineCount) {
			return retry.ExpectedErrorf("expected %d requests got %d", machineRequestSet.TypedSpec().Value.MachineCount, len(ids))
		}

		return nil
	})

	require.NoError(t, err)

	rtestutils.AssertResources(ctx, t, st, ids, func(r *infra.MachineRequest, assert *assert.Assertions) {
		l, ok := r.Metadata().Labels().Get(omni.LabelMachineRequestSet)

		assert.True(ok)
		assert.Equal(l, machineRequestSet.Metadata().ID())

		l, ok = r.Metadata().Labels().Get(omni.LabelInfraProviderID)

		assert.True(ok)
		assert.Equal(l, machineRequestSet.TypedSpec().Value.ProviderId)

		assert.Equal(machineRequestSet.TypedSpec().Value.TalosVersion, r.TypedSpec().Value.TalosVersion)
	})

	// scale down
	machineRequestSet, err = safe.StateUpdateWithConflicts(ctx, st, machineRequestSet.Metadata(), func(r *omni.MachineRequestSet) error {
		r.TypedSpec().Value.MachineCount = 2

		return nil
	})

	require.NoError(t, err)

	err = retry.Constant(time.Second*5).RetryWithContext(ctx, func(ctx context.Context) error {
		ids = rtestutils.ResourceIDs[*infra.MachineRequest](ctx, t, st, state.WithLabelQuery(resource.LabelEqual(omni.LabelMachineRequestSet, machineRequestSet.Metadata().ID())))

		if len(ids) != int(machineRequestSet.TypedSpec().Value.MachineCount) {
			return retry.ExpectedErrorf("expected %d requests got %d", machineRequestSet.TypedSpec().Value.MachineCount, len(ids))
		}

		return nil
	})

	require.NoError(t, err)

	// remove the machine request set
	rtestutils.DestroyAll[*omni.MachineRequestSet](ctx, t, st)

	requests, err := safe.ReaderListAll[*infra.MachineRequest](ctx, st)
	require.NoError(t, err)

	require.True(t, requests.Len() == 0)
}

//nolint:gocognit
func reconcileLabels(ctx context.Context, st state.State, ready chan<- struct{}) error {
	ch := make(chan state.Event, 64)

	err := st.WatchKind(ctx, infra.NewMachineRequestStatus("").Metadata(), ch)
	if err != nil {
		return err
	}

	close(ready)

	deleteLabels := func(id string) error {
		res := system.NewResourceLabels[*omni.MachineStatus](id)

		deleteCtx, cancel := context.WithTimeout(ctx, time.Second*10)
		defer cancel()

		err = st.TeardownAndDestroy(deleteCtx, res.Metadata())
		if err != nil && !state.IsNotFoundError(err) {
			return err
		}

		return nil
	}

	for {
		select {
		case <-ctx.Done():
			return nil
		case event := <-ch:
			switch event.Type {
			case state.Bootstrapped, state.Noop:
			case state.Errored:
				return event.Error
			case state.Destroyed:
				status := event.Resource.(*infra.MachineRequestStatus) //nolint:errcheck,forcetypeassert

				if err = deleteLabels(status.TypedSpec().Value.Id); err != nil {
					return err
				}
			case state.Created, state.Updated:
				status := event.Resource.(*infra.MachineRequestStatus) //nolint:errcheck,forcetypeassert

				err = safe.StateModify(ctx, st, system.NewResourceLabels[*omni.MachineStatus](status.TypedSpec().Value.Id), func(r *system.ResourceLabels[*omni.MachineStatus]) error {
					if r.Metadata().Phase() == resource.PhaseTearingDown {
						return nil
					}

					r.Metadata().Labels().Set(omni.LabelMachineRequest, status.Metadata().ID())

					helpers.CopyAllLabels(status, r)

					return nil
				})
				if err != nil && !state.IsPhaseConflictError(err) {
					return err
				}
			}
		}
	}
}

func TestMachineRequestForceDestroy(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(t.Context(), 20*time.Second)
	t.Cleanup(cancel)

	testutils.WithRuntime(ctx, t, testutils.TestOptions{},
		func(_ context.Context, tc testutils.TestContext) {
			require.NoError(t, tc.Runtime.RegisterQController(omnictrl.NewMachineRequestSetStatusController()))
			require.NoError(t, tc.Runtime.RegisterQController(newTestInfraProvider()))
		},
		func(ctx context.Context, tc testutils.TestContext) {
			testMachineRequestForceDestroy(ctx, t, tc.State)
		},
	)
}

func testMachineRequestForceDestroy(ctx context.Context, t *testing.T, st state.State) {
	var eg errgroup.Group

	reconcilerCtx, reconcilerCancel := context.WithCancel(ctx)

	defer func() {
		reconcilerCancel()

		require.NoError(t, eg.Wait())
	}()

	watchReady := make(chan struct{})
	errCh := make(chan error, 1)

	eg.Go(func() error {
		err := reconcileLabels(reconcilerCtx, st, watchReady)
		if err != nil {
			t.Logf("machine status labels reconciler crashed with error %s", err)

			select {
			case errCh <- err:
			default:
			}
		}

		return err
	})

	select {
	case <-watchReady:
	case err := <-errCh:
		require.FailNow(t, "reconcileLabels failed before watch was established", err)
	case <-ctx.Done():
		require.FailNow(t, "context canceled waiting for reconcileLabels watch to be established")
	}

	machineRequestSet := rmock.Mock[*omni.MachineRequestSet](ctx, t, st,
		options.WithID("test-force-destroy"),
		options.Modify(func(r *omni.MachineRequestSet) error {
			r.TypedSpec().Value.ProviderId = "test"
			r.TypedSpec().Value.TalosVersion = "v1.7.5"
			r.TypedSpec().Value.MachineCount = 2

			return nil
		}),
	)

	// wait for the requests to be created
	var ids []resource.ID

	err := retry.Constant(time.Second*5).RetryWithContext(ctx, func(ctx context.Context) error {
		ids = rtestutils.ResourceIDs[*infra.MachineRequest](ctx, t, st, state.WithLabelQuery(resource.LabelEqual(omni.LabelMachineRequestSet, machineRequestSet.Metadata().ID())))

		if len(ids) != int(machineRequestSet.TypedSpec().Value.MachineCount) {
			return retry.ExpectedErrorf("expected %d requests got %d", machineRequestSet.TypedSpec().Value.MachineCount, len(ids))
		}

		return nil
	})

	require.NoError(t, err)

	targetRequestID := ids[0]

	// force destroy the ownerless machine request by tearing it down directly
	_, err = st.Teardown(ctx, infra.NewMachineRequest(targetRequestID).Metadata())

	require.NoError(t, err)

	// the original request should be destroyed
	rtestutils.AssertNoResource[*infra.MachineRequest](ctx, t, st, targetRequestID)

	// a replacement request should be created, so total count is still 2
	err = retry.Constant(time.Second*5).RetryWithContext(ctx, func(ctx context.Context) error {
		ids = rtestutils.ResourceIDs[*infra.MachineRequest](ctx, t, st, state.WithLabelQuery(resource.LabelEqual(omni.LabelMachineRequestSet, machineRequestSet.Metadata().ID())))

		if len(ids) != int(machineRequestSet.TypedSpec().Value.MachineCount) {
			return retry.ExpectedErrorf("expected %d requests got %d", machineRequestSet.TypedSpec().Value.MachineCount, len(ids))
		}

		return nil
	})

	require.NoError(t, err)

	// the replacement should be a new request, not the original
	assert.NotContains(t, ids, targetRequestID)
}

func TestInfraProviderDeletion(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(t.Context(), time.Second*10)
	t.Cleanup(cancel)

	providerID := "infra-provider"

	testutils.WithRuntime(ctx, t, testutils.TestOptions{}, func(ctx context.Context, testContext testutils.TestContext) {
		st := testContext.State

		infraProvider := infra.NewProvider(providerID)
		require.NoError(t, st.Create(ctx, infraProvider))

		require.NoError(t, testContext.Runtime.RegisterQController(omnictrl.NewMachineRequestSetStatusController()))
	}, func(ctx context.Context, testContext testutils.TestContext) {
		st := testContext.State

		machineRequestSet := omni.NewMachineRequestSet("test-machine-request-set")
		machineRequestSet.Metadata().Labels().Set(omni.LabelInfraProviderID, providerID)
		machineRequestSet.TypedSpec().Value.MachineCount = 1
		machineRequestSet.TypedSpec().Value.ProviderId = providerID
		machineRequestSet.TypedSpec().Value.TalosVersion = "v1.11.4"

		require.NoError(t, st.Create(ctx, machineRequestSet))

		var ids []resource.ID

		for range 10 {
			ids = rtestutils.ResourceIDs[*infra.MachineRequest](ctx, t, st, state.WithLabelQuery(resource.LabelEqual(omni.LabelMachineRequestSet, machineRequestSet.Metadata().ID())))

			if len(ids) == int(machineRequestSet.TypedSpec().Value.MachineCount) {
				break
			}

			time.Sleep(time.Millisecond * 500)
		}

		require.Equal(t, int(machineRequestSet.TypedSpec().Value.MachineCount), len(ids))

		rtestutils.AssertResources(ctx, t, st, ids, func(r *infra.MachineRequest, assert *assert.Assertions) {
			l, ok := r.Metadata().Labels().Get(omni.LabelMachineRequestSet)

			assert.True(ok)
			assert.Equal(l, machineRequestSet.Metadata().ID())

			l, ok = r.Metadata().Labels().Get(omni.LabelInfraProviderID)

			assert.True(ok)
			assert.Equal(l, machineRequestSet.TypedSpec().Value.ProviderId)

			assert.Equal(machineRequestSet.TypedSpec().Value.TalosVersion, r.TypedSpec().Value.TalosVersion)
		})

		// simulate infra provider operations
		machineRequests, err := safe.ReaderListAll[*infra.MachineRequest](ctx, st, state.WithLabelQuery(resource.LabelEqual(omni.LabelMachineRequestSet, machineRequestSet.Metadata().ID())))
		require.NoError(t, err)
		require.Equal(t, machineRequests.Len(), int(machineRequestSet.TypedSpec().Value.MachineCount))

		machineRequest := machineRequests.Get(0)
		machineRequest.Metadata().Finalizers().Add(providerID)
		require.NoError(t, st.Update(ctx, machineRequest))

		rtestutils.AssertResources(ctx, t, st, ids, func(r *infra.MachineRequest, assert *assert.Assertions) {
			require.True(t, r.Metadata().Finalizers().Has(providerID))
		})

		// tear down machine request set
		_, err = st.Teardown(ctx, machineRequestSet.Metadata())
		require.NoError(t, err)

		// assert that all machine requests are tearing down but still has the infra provider finalizer
		rtestutils.AssertResources[*infra.MachineRequest](ctx, t, st, ids, func(request *infra.MachineRequest, assert *assert.Assertions) {
			assert.Equal(resource.PhaseTearingDown, request.Metadata().Phase())
			assert.True(request.Metadata().Finalizers().Has(providerID))
		})

		// delete infra provider
		rtestutils.Destroy[*infra.Provider](ctx, t, st, []string{providerID})

		for _, id := range ids {
			rtestutils.AssertNoResource[*infra.MachineRequest](ctx, t, st, id)
		}

		require.NoError(t, st.TeardownAndDestroy(ctx, machineRequestSet.Metadata()))
	})
}
