package unit_test

import (
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/coder/coder/v2/agent/unit"
)

const (
	unitA unit.ID = "serviceA"
	unitB unit.ID = "serviceB"
	unitC unit.ID = "serviceC"
	unitD unit.ID = "serviceD"
)

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

	t.Run("Empty Unit Name", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		err := manager.Register("")
		require.ErrorIs(t, err, unit.ErrUnitIDRequired)
		err = manager.AddDependency("", unitA, unit.StatusStarted)
		require.ErrorIs(t, err, unit.ErrUnitIDRequired)
		err = manager.AddDependency(unitA, "", unit.StatusStarted)
		require.ErrorIs(t, err, unit.ErrUnitIDRequired)
		dependencies, err := manager.GetAllDependencies("")
		require.ErrorIs(t, err, unit.ErrUnitIDRequired)
		require.Len(t, dependencies, 0)
		unmetDependencies, err := manager.GetUnmetDependencies("")
		require.ErrorIs(t, err, unit.ErrUnitIDRequired)
		require.Len(t, unmetDependencies, 0)
		err = manager.UpdateStatus("", unit.StatusStarted)
		require.ErrorIs(t, err, unit.ErrUnitIDRequired)
		isReady, err := manager.IsReady("")
		require.ErrorIs(t, err, unit.ErrUnitIDRequired)
		require.False(t, isReady)
		u, err := manager.Unit("")
		require.ErrorIs(t, err, unit.ErrUnitIDRequired)
		assert.Equal(t, unit.Unit{}, u)
	})
}

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

	t.Run("RegisterNewUnit", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Given: a unit is registered
		err := manager.Register(unitA)
		require.NoError(t, err)

		// Then: the unit should be ready (no dependencies)
		u, err := manager.Unit(unitA)
		require.NoError(t, err)
		assert.Equal(t, unitA, u.ID())
		assert.Equal(t, unit.StatusPending, u.Status())
		isReady, err := manager.IsReady(unitA)
		require.NoError(t, err)
		assert.True(t, isReady)
	})

	t.Run("RegisterDuplicateUnit", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Given: a unit is registered
		err := manager.Register(unitA)
		require.NoError(t, err)

		// Newly registered units have StatusPending. We update the unit status to StatusStarted,
		// so we can later assert that it is not overwritten back to StatusPending by the second
		// register call
		manager.UpdateStatus(unitA, unit.StatusStarted)

		// When: the unit is registered again
		err = manager.Register(unitA)

		// Then: a descriptive error should be returned
		require.ErrorIs(t, err, unit.ErrUnitAlreadyRegistered)

		// Then: the unit status should not be overwritten
		u, err := manager.Unit(unitA)
		require.NoError(t, err)
		assert.Equal(t, unit.StatusStarted, u.Status())
		isReady, err := manager.IsReady(unitA)
		require.NoError(t, err)
		assert.True(t, isReady)
	})

	t.Run("RegisterMultipleUnits", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Given: multiple units are registered
		unitIDs := []unit.ID{unitA, unitB, unitC}
		for _, unit := range unitIDs {
			err := manager.Register(unit)
			require.NoError(t, err)
		}

		// Then: all units should be ready initially
		for _, unitID := range unitIDs {
			u, err := manager.Unit(unitID)
			require.NoError(t, err)
			assert.Equal(t, unit.StatusPending, u.Status())
			isReady, err := manager.IsReady(unitID)
			require.NoError(t, err)
			assert.True(t, isReady)
		}
	})
}

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

	t.Run("AddDependencyBetweenRegisteredUnits", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Given: units A and B are registered
		err := manager.Register(unitA)
		require.NoError(t, err)
		err = manager.Register(unitB)
		require.NoError(t, err)

		// Given: Unit A depends on Unit B being unit.StatusStarted
		err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
		require.NoError(t, err)

		// Then: Unit A should not be ready (depends on B)
		u, err := manager.Unit(unitA)
		require.NoError(t, err)
		assert.Equal(t, unit.StatusPending, u.Status())
		isReady, err := manager.IsReady(unitA)
		require.NoError(t, err)
		assert.False(t, isReady)

		// Then: Unit B should still be ready (no dependencies)
		u, err = manager.Unit(unitB)
		require.NoError(t, err)
		assert.Equal(t, unit.StatusPending, u.Status())
		isReady, err = manager.IsReady(unitB)
		require.NoError(t, err)
		assert.True(t, isReady)

		// When: Unit B is started
		err = manager.UpdateStatus(unitB, unit.StatusStarted)
		require.NoError(t, err)

		// Then: Unit A should be ready, because its dependency is now in the desired state.
		isReady, err = manager.IsReady(unitA)
		require.NoError(t, err)
		assert.True(t, isReady)

		// When: Unit B is stopped
		err = manager.UpdateStatus(unitB, unit.StatusPending)
		require.NoError(t, err)

		// Then: Unit A should no longer be ready, because its dependency is not in the desired state.
		isReady, err = manager.IsReady(unitA)
		require.NoError(t, err)
		assert.False(t, isReady)
	})

	t.Run("AddDependencyByAnUnregisteredDependentUnit", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Given Unit B is registered
		err := manager.Register(unitB)
		require.NoError(t, err)

		// Given Unit A depends on Unit B being started
		err = manager.AddDependency(unitA, unitB, unit.StatusStarted)

		// Then: a descriptive error communicates that the dependency cannot be added
		// because the dependent unit must be registered first.
		require.ErrorIs(t, err, unit.ErrUnitNotFound)
	})

	t.Run("AddDependencyOnAnUnregisteredUnit", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Given unit A is registered
		err := manager.Register(unitA)
		require.NoError(t, err)

		// Given Unit B is not yet registered
		// And Unit A depends on Unit B being started
		err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
		require.NoError(t, err)

		// Then: The dependency should be visible in Unit A's status
		dependencies, err := manager.GetAllDependencies(unitA)
		require.NoError(t, err)
		require.Len(t, dependencies, 1)
		assert.Equal(t, unitB, dependencies[0].DependsOn)
		assert.Equal(t, unit.StatusStarted, dependencies[0].RequiredStatus)
		assert.False(t, dependencies[0].IsSatisfied)

		u, err := manager.Unit(unitB)
		require.NoError(t, err)
		assert.Equal(t, unit.StatusNotRegistered, u.Status())

		// Then: Unit A should not be ready, because it depends on Unit B
		isReady, err := manager.IsReady(unitA)
		require.NoError(t, err)
		assert.False(t, isReady)

		// When: Unit B is registered
		err = manager.Register(unitB)
		require.NoError(t, err)

		// Then: Unit A should still not be ready.
		// Unit B is not registered, but it has not been started as required by the dependency.
		isReady, err = manager.IsReady(unitA)
		require.NoError(t, err)
		assert.False(t, isReady)

		// When: Unit B is started
		err = manager.UpdateStatus(unitB, unit.StatusStarted)
		require.NoError(t, err)

		// Then: Unit A should be ready, because its dependency is now in the desired state.
		isReady, err = manager.IsReady(unitA)
		require.NoError(t, err)
		assert.True(t, isReady)
	})

	t.Run("AddDependencyCreatesACyclicDependency", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Register units
		err := manager.Register(unitA)
		require.NoError(t, err)
		err = manager.Register(unitB)
		require.NoError(t, err)
		err = manager.Register(unitC)
		require.NoError(t, err)
		err = manager.Register(unitD)
		require.NoError(t, err)

		// A depends on B
		err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
		require.NoError(t, err)
		// B depends on C
		err = manager.AddDependency(unitB, unitC, unit.StatusStarted)
		require.NoError(t, err)

		// C depends on D
		err = manager.AddDependency(unitC, unitD, unit.StatusStarted)
		require.NoError(t, err)

		// Try to make D depend on A (creates indirect cycle)
		err = manager.AddDependency(unitD, unitA, unit.StatusStarted)
		require.ErrorIs(t, err, unit.ErrCycleDetected)
	})

	t.Run("UpdatingADependency", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Given units A and B are registered
		err := manager.Register(unitA)
		require.NoError(t, err)
		err = manager.Register(unitB)
		require.NoError(t, err)

		// Given Unit A depends on Unit B being unit.StatusStarted
		err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
		require.NoError(t, err)

		// When: The dependency is updated to unit.StatusComplete
		err = manager.AddDependency(unitA, unitB, unit.StatusComplete)
		require.NoError(t, err)

		// Then: Unit A should only have one dependency, and it should be unit.StatusComplete
		dependencies, err := manager.GetAllDependencies(unitA)
		require.NoError(t, err)
		require.Len(t, dependencies, 1)
		assert.Equal(t, unit.StatusComplete, dependencies[0].RequiredStatus)
	})
}

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

	t.Run("UpdateStatusTriggersReadinessRecalculation", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Given units A and B are registered
		err := manager.Register(unitA)
		require.NoError(t, err)
		err = manager.Register(unitB)
		require.NoError(t, err)

		// Given Unit A depends on Unit B being unit.StatusStarted
		err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
		require.NoError(t, err)

		// Then: Unit A should not be ready (depends on B)
		u, err := manager.Unit(unitA)
		require.NoError(t, err)
		assert.Equal(t, unit.StatusPending, u.Status())
		isReady, err := manager.IsReady(unitA)
		require.NoError(t, err)
		assert.False(t, isReady)

		// When: Unit B is started
		err = manager.UpdateStatus(unitB, unit.StatusStarted)
		require.NoError(t, err)

		// Then: Unit A should be ready, because its dependency is now in the desired state.
		u, err = manager.Unit(unitA)
		require.NoError(t, err)
		assert.Equal(t, unit.StatusPending, u.Status())
		isReady, err = manager.IsReady(unitA)
		require.NoError(t, err)
		assert.True(t, isReady)
	})

	t.Run("UpdateStatusWithUnregisteredUnit", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Given Unit A is not registered
		// When: Unit A is updated to unit.StatusStarted
		err := manager.UpdateStatus(unitA, unit.StatusStarted)

		// Then: a descriptive error communicates that the unit must be registered first.
		require.ErrorIs(t, err, unit.ErrUnitNotFound)
	})

	t.Run("LinearChainDependencies", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Given units A, B, and C are registered
		err := manager.Register(unitA)
		require.NoError(t, err)
		err = manager.Register(unitB)
		require.NoError(t, err)
		err = manager.Register(unitC)
		require.NoError(t, err)

		// Create chain: A depends on B being "started", B depends on C being "completed"
		err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
		require.NoError(t, err)
		err = manager.AddDependency(unitB, unitC, unit.StatusComplete)
		require.NoError(t, err)

		// Then: only Unit C should be ready (no dependencies)
		u, err := manager.Unit(unitC)
		require.NoError(t, err)
		assert.Equal(t, unit.StatusPending, u.Status())
		isReady, err := manager.IsReady(unitC)
		require.NoError(t, err)
		assert.True(t, isReady)

		u, err = manager.Unit(unitB)
		require.NoError(t, err)
		assert.Equal(t, unit.StatusPending, u.Status())
		isReady, err = manager.IsReady(unitB)
		require.NoError(t, err)
		assert.False(t, isReady)

		u, err = manager.Unit(unitA)
		require.NoError(t, err)
		assert.Equal(t, unit.StatusPending, u.Status())
		isReady, err = manager.IsReady(unitA)
		require.NoError(t, err)
		assert.False(t, isReady)

		// When: Unit C is completed
		err = manager.UpdateStatus(unitC, unit.StatusComplete)
		require.NoError(t, err)

		// Then: Unit B should be ready, because its dependency is now in the desired state.
		u, err = manager.Unit(unitB)
		require.NoError(t, err)
		assert.Equal(t, unit.StatusPending, u.Status())
		isReady, err = manager.IsReady(unitB)
		require.NoError(t, err)
		assert.True(t, isReady)

		u, err = manager.Unit(unitA)
		require.NoError(t, err)
		assert.Equal(t, unit.StatusPending, u.Status())
		isReady, err = manager.IsReady(unitA)
		require.NoError(t, err)
		assert.False(t, isReady)

		u, err = manager.Unit(unitB)
		require.NoError(t, err)
		assert.Equal(t, unit.StatusPending, u.Status())
		isReady, err = manager.IsReady(unitB)
		require.NoError(t, err)
		assert.True(t, isReady)

		// When: Unit B is started
		err = manager.UpdateStatus(unitB, unit.StatusStarted)
		require.NoError(t, err)

		// Then: Unit A should be ready, because its dependency is now in the desired state.
		u, err = manager.Unit(unitA)
		require.NoError(t, err)
		assert.Equal(t, unit.StatusPending, u.Status())
		isReady, err = manager.IsReady(unitA)
		require.NoError(t, err)
		assert.True(t, isReady)
	})
}

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

	t.Run("GetUnmetDependenciesForUnitWithNoDependencies", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Given: Unit A is registered
		err := manager.Register(unitA)
		require.NoError(t, err)

		// Given: Unit A has no dependencies
		// Then: Unit A should have no unmet dependencies
		unmet, err := manager.GetUnmetDependencies(unitA)
		require.NoError(t, err)
		assert.Empty(t, unmet)
	})

	t.Run("GetUnmetDependenciesForUnitWithUnsatisfiedDependencies", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()
		err := manager.Register(unitA)
		require.NoError(t, err)
		err = manager.Register(unitB)
		require.NoError(t, err)

		// Given: Unit A depends on Unit B being unit.StatusStarted
		err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
		require.NoError(t, err)

		unmet, err := manager.GetUnmetDependencies(unitA)
		require.NoError(t, err)
		require.Len(t, unmet, 1)

		assert.Equal(t, unitA, unmet[0].Unit)
		assert.Equal(t, unitB, unmet[0].DependsOn)
		assert.Equal(t, unit.StatusStarted, unmet[0].RequiredStatus)
		assert.False(t, unmet[0].IsSatisfied)
	})

	t.Run("GetUnmetDependenciesForUnitWithSatisfiedDependencies", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Given: Unit A and Unit B are registered
		err := manager.Register(unitA)
		require.NoError(t, err)
		err = manager.Register(unitB)
		require.NoError(t, err)

		// Given: Unit A depends on Unit B being unit.StatusStarted
		err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
		require.NoError(t, err)

		// When: Unit B is started
		err = manager.UpdateStatus(unitB, unit.StatusStarted)
		require.NoError(t, err)

		// Then: Unit A should have no unmet dependencies
		unmet, err := manager.GetUnmetDependencies(unitA)
		require.NoError(t, err)
		assert.Empty(t, unmet)
	})

	t.Run("GetUnmetDependenciesForUnregisteredUnit", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// When: Unit A is requested
		unmet, err := manager.GetUnmetDependencies(unitA)

		// Then: a descriptive error communicates that the unit must be registered first.
		require.ErrorIs(t, err, unit.ErrUnitNotFound)
		assert.Nil(t, unmet)
	})
}

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

	t.Run("UnitWithMultipleDependencies", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Register all units
		units := []unit.ID{unitA, unitB, unitC, unitD}
		for _, unit := range units {
			err := manager.Register(unit)
			require.NoError(t, err)
		}

		// A depends on B being unit.StatusStarted AND C being "started"
		err := manager.AddDependency(unitA, unitB, unit.StatusStarted)
		require.NoError(t, err)
		err = manager.AddDependency(unitA, unitC, unit.StatusStarted)
		require.NoError(t, err)

		// A should not be ready (depends on both B and C)
		isReady, err := manager.IsReady(unitA)
		require.NoError(t, err)
		assert.False(t, isReady)

		// Update B to unit.StatusStarted - A should still not be ready (needs C too)
		err = manager.UpdateStatus(unitB, unit.StatusStarted)
		require.NoError(t, err)

		isReady, err = manager.IsReady(unitA)
		require.NoError(t, err)
		assert.False(t, isReady)

		// Update C to "started" - A should now be ready
		err = manager.UpdateStatus(unitC, unit.StatusStarted)
		require.NoError(t, err)

		isReady, err = manager.IsReady(unitA)
		require.NoError(t, err)
		assert.True(t, isReady)
	})

	t.Run("ComplexDependencyChain", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Register all units
		units := []unit.ID{unitA, unitB, unitC, unitD}
		for _, unit := range units {
			err := manager.Register(unit)
			require.NoError(t, err)
		}

		// Create complex dependency graph:
		// A depends on B being unit.StatusStarted AND C being "started"
		err := manager.AddDependency(unitA, unitB, unit.StatusStarted)
		require.NoError(t, err)
		err = manager.AddDependency(unitA, unitC, unit.StatusStarted)
		require.NoError(t, err)
		// B depends on D being "completed"
		err = manager.AddDependency(unitB, unitD, unit.StatusComplete)
		require.NoError(t, err)
		// C depends on D being "completed"
		err = manager.AddDependency(unitC, unitD, unit.StatusComplete)
		require.NoError(t, err)

		// Initially only D is ready
		isReady, err := manager.IsReady(unitD)
		require.NoError(t, err)
		assert.True(t, isReady)
		isReady, err = manager.IsReady(unitB)
		require.NoError(t, err)
		assert.False(t, isReady)
		isReady, err = manager.IsReady(unitC)
		require.NoError(t, err)
		assert.False(t, isReady)
		isReady, err = manager.IsReady(unitA)
		require.NoError(t, err)
		assert.False(t, isReady)

		// Update D to "completed" - B and C should become ready
		err = manager.UpdateStatus(unitD, unit.StatusComplete)
		require.NoError(t, err)

		isReady, err = manager.IsReady(unitB)
		require.NoError(t, err)
		assert.True(t, isReady)
		isReady, err = manager.IsReady(unitC)
		require.NoError(t, err)
		assert.True(t, isReady)
		isReady, err = manager.IsReady(unitA)
		require.NoError(t, err)
		assert.False(t, isReady)

		// Update B to unit.StatusStarted - A should still not be ready (needs C)
		err = manager.UpdateStatus(unitB, unit.StatusStarted)
		require.NoError(t, err)

		isReady, err = manager.IsReady(unitA)
		require.NoError(t, err)
		assert.False(t, isReady)

		// Update C to "started" - A should now be ready
		err = manager.UpdateStatus(unitC, unit.StatusStarted)
		require.NoError(t, err)

		isReady, err = manager.IsReady(unitA)
		require.NoError(t, err)
		assert.True(t, isReady)
	})

	t.Run("DifferentStatusTypes", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Register units
		err := manager.Register(unitA)
		require.NoError(t, err)
		err = manager.Register(unitB)
		require.NoError(t, err)
		err = manager.Register(unitC)
		require.NoError(t, err)

		// Given: Unit A depends on Unit B being unit.StatusStarted
		err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
		require.NoError(t, err)
		// Given: Unit A depends on Unit C being "completed"
		err = manager.AddDependency(unitA, unitC, unit.StatusComplete)
		require.NoError(t, err)

		// When: Unit B is started
		err = manager.UpdateStatus(unitB, unit.StatusStarted)
		require.NoError(t, err)

		// Then: Unit A should not be ready, because only one of its dependencies is in the desired state.
		// It still requires Unit C to be completed.
		isReady, err := manager.IsReady(unitA)
		require.NoError(t, err)
		assert.False(t, isReady)

		// When: Unit C is completed
		err = manager.UpdateStatus(unitC, unit.StatusComplete)
		require.NoError(t, err)

		// Then: Unit A should be ready, because both of its dependencies are in the desired state.
		isReady, err = manager.IsReady(unitA)
		require.NoError(t, err)
		assert.True(t, isReady)
	})
}

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

	t.Run("IsReadyWithUnregisteredUnit", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Given: a unit is not registered
		u, err := manager.Unit(unitA)
		require.NoError(t, err)
		assert.Equal(t, unit.StatusNotRegistered, u.Status())
		// Then: the unit is not ready
		isReady, err := manager.IsReady(unitA)
		require.NoError(t, err)
		assert.True(t, isReady)
	})
}

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

	t.Run("ExportSimpleGraph", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Register units
		err := manager.Register(unitA)
		require.NoError(t, err)
		err = manager.Register(unitB)
		require.NoError(t, err)

		// Add dependency
		err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
		require.NoError(t, err)

		dot, err := manager.ExportDOT("test")
		require.NoError(t, err)
		assert.NotEmpty(t, dot)
		assert.Contains(t, dot, "digraph")
	})

	t.Run("ExportComplexGraph", func(t *testing.T) {
		t.Parallel()

		manager := unit.NewManager()

		// Register all units
		units := []unit.ID{unitA, unitB, unitC, unitD}
		for _, unit := range units {
			err := manager.Register(unit)
			require.NoError(t, err)
		}

		// Create complex dependency graph
		// A depends on B and C, B depends on D, C depends on D
		err := manager.AddDependency(unitA, unitB, unit.StatusStarted)
		require.NoError(t, err)
		err = manager.AddDependency(unitA, unitC, unit.StatusStarted)
		require.NoError(t, err)
		err = manager.AddDependency(unitB, unitD, unit.StatusComplete)
		require.NoError(t, err)
		err = manager.AddDependency(unitC, unitD, unit.StatusComplete)
		require.NoError(t, err)

		dot, err := manager.ExportDOT("complex")
		require.NoError(t, err)
		assert.NotEmpty(t, dot)
		assert.Contains(t, dot, "digraph")
	})
}