package impl

import (
	"github.com/elliotchance/pie/v2"
	"github.com/tidbcloud/dataflow-service-ng/app/internal/status_matrix"
)

type statusMatrixImpl struct {
	// status -> action -> next status
	statusMatrix map[string]map[string]string
	// action -> action -> strategy
	friendlyActionMap map[string]map[string]status_matrix.ActionStrategy

	steadyStateMatrix map[string]string
}

func NewStatusMatrix(
	statusMatrix map[string]map[string]string,
	friendlyActionMap map[string]map[string]status_matrix.ActionStrategy,
	steadyStateMatrix map[string]string,
) status_matrix.StatusMatrix {
	return &statusMatrixImpl{
		statusMatrix:      statusMatrix,
		friendlyActionMap: friendlyActionMap,
		steadyStateMatrix: steadyStateMatrix,
	}
}

func (m *statusMatrixImpl) IsActionFriendly2Status(status status_matrix.CurrentStatus, action string) bool {
	nextStatus := m.CalculateNextStatusForAction(status, action)
	return nextStatus != nil
}

func (m *statusMatrixImpl) CalculateNextStatusForAction(
	currentStatus status_matrix.CurrentStatus, action string,
) *status_matrix.NextStatus {
	matrix, ok := m.statusMatrix[currentStatus.Status]
	if !ok {
		return nil
	}
	nextStatus, ok := matrix[action]
	if !ok {
		return nil
	}

	if len(m.friendlyActionMap) == 0 {
		return &status_matrix.NextStatus{
			Status: nextStatus,
		}
	}

	if len(currentStatus.Actions) == 0 {
		return &status_matrix.NextStatus{
			Status:  nextStatus,
			Actions: []string{action},
		}
	}

	return m.isActionFriendly(currentStatus, action, nextStatus)
}

func (m *statusMatrixImpl) isActionFriendly(
	currentStatus status_matrix.CurrentStatus, action string, nextStatus string,
) *status_matrix.NextStatus {
	nextActions := make(map[string]struct{})
	for _, act := range currentStatus.Actions {
		friendlyActions, ok := m.friendlyActionMap[act]
		if !ok {
			continue
		}
		strategy, ok := friendlyActions[action]
		if !ok {
			return nil
		}
		if strategy == status_matrix.Grab {
			nextActions[action] = struct{}{}
			continue
		}
		if strategy == status_matrix.Coexist {
			nextActions[act] = struct{}{}
			nextActions[action] = struct{}{}
		}
	}

	if len(nextActions) == 0 {
		return nil
	}

	return &status_matrix.NextStatus{
		Status:  nextStatus,
		Actions: pie.Keys(nextActions),
	}
}

func (m *statusMatrixImpl) CalculateNextStatusIfActionDone(
	currentStatus status_matrix.CurrentStatus, action string,
) *status_matrix.NextStatus {
	steadyStatus, ok := m.steadyStateMatrix[currentStatus.Status]
	if !ok {
		return nil
	}
	if len(currentStatus.Actions) == 0 {
		return &status_matrix.NextStatus{
			Status: steadyStatus,
		}
	}

	actions := pie.Of(currentStatus.Actions).
		Filter(func(act string) bool {
			return act != action
		}).Result
	if len(actions) == 0 {
		return &status_matrix.NextStatus{
			Status: steadyStatus,
		}
	}
	return &status_matrix.NextStatus{
		Status:  currentStatus.Status,
		Actions: actions,
	}
}

func (m *statusMatrixImpl) QueryAllowActions(currentStatus status_matrix.CurrentStatus) []string {
	matrix, ok := m.statusMatrix[currentStatus.Status]
	if !ok {
		return nil
	}
	return pie.Keys(matrix)
}