package nodegroup

import (
	"fmt"
	"time"
)

// State is a typed string for the node instance group state machine.
// Values are derived from the canonical string constants in consts.go.
type State string

const (
	StateNormal           State = State(NodeInstanceGroupStateNormal)
	StateShortageDetected State = State(NodeInstanceGroupStateShortageDetected)
	StateFallbackCreating State = State(NodeInstanceGroupStateFallbackCreating)
	StateFallbackReady    State = State(NodeInstanceGroupStateFallbackReady)
	StateRecovering       State = State(NodeInstanceGroupStateRecovering)
	StateDrainingFallback State = State(NodeInstanceGroupStateDrainingFallback)
	StateDeletingFallback State = State(NodeInstanceGroupStateDeletingFallback)
	StateDeprecated       State = State(NodeInstanceGroupStateDeprecated)
)

type Event string

const (
	EventShortageDetected   Event = "shortage_detected"
	EventFallbackCreated    Event = "fallback_created"
	EventFallbackReady      Event = "fallback_ready"
	EventCapacityRestored   Event = "capacity_restored"
	EventDrainComplete      Event = "drain_complete"
	EventFallbackDeleted    Event = "fallback_deleted"
	EventDeprecationStarted Event = "deprecation_started"
	EventCooldownExpired    Event = "cooldown_expired"
)

type NodeInstanceGroupStateMachine struct {
	CurrentState     State
	SharedPoolID     string
	PrimarySpec      string
	LastTransitionAt time.Time
	CooldownDuration time.Duration
	now              func() time.Time
}

type Transition struct {
	From  State
	Event Event
	To    State
}

var validTransitions = []Transition{
	{StateNormal, EventShortageDetected, StateShortageDetected},
	{StateNormal, EventDeprecationStarted, StateDeprecated},

	{StateShortageDetected, EventFallbackCreated, StateFallbackCreating},
	{StateShortageDetected, EventCapacityRestored, StateNormal},

	{StateFallbackCreating, EventFallbackReady, StateFallbackReady},

	{StateFallbackReady, EventCapacityRestored, StateRecovering},

	{StateRecovering, EventCooldownExpired, StateDrainingFallback},
	{StateRecovering, EventShortageDetected, StateFallbackReady},

	{StateDrainingFallback, EventDrainComplete, StateDeletingFallback},
	{StateDrainingFallback, EventShortageDetected, StateFallbackReady},

	{StateDeletingFallback, EventFallbackDeleted, StateNormal},
}

func NewStateMachine(sharedPoolID, primarySpec string, initialState State, cooldown time.Duration) *NodeInstanceGroupStateMachine {
	if cooldown == 0 {
		cooldown = 30 * time.Minute
	}
	return &NodeInstanceGroupStateMachine{
		CurrentState:     initialState,
		SharedPoolID:     sharedPoolID,
		PrimarySpec:      primarySpec,
		LastTransitionAt: time.Now(),
		CooldownDuration: cooldown,
		now:              time.Now,
	}
}

// WithClock overrides the clock function used for time-dependent operations.
// Intended for testing; production code should use the default (time.Now).
func (sm *NodeInstanceGroupStateMachine) WithClock(now func() time.Time) *NodeInstanceGroupStateMachine {
	sm.now = now
	sm.LastTransitionAt = now()
	return sm
}

func (sm *NodeInstanceGroupStateMachine) CanTransition(event Event) bool {
	for _, t := range validTransitions {
		if t.From == sm.CurrentState && t.Event == event {
			return true
		}
	}
	return false
}

func (sm *NodeInstanceGroupStateMachine) Transition(event Event) error {
	for _, t := range validTransitions {
		if t.From == sm.CurrentState && t.Event == event {
			sm.CurrentState = t.To
			sm.LastTransitionAt = sm.now()
			return nil
		}
	}
	return fmt.Errorf("invalid transition: %s + %s", sm.CurrentState, event)
}

func (sm *NodeInstanceGroupStateMachine) NextState(event Event) (State, error) {
	for _, t := range validTransitions {
		if t.From == sm.CurrentState && t.Event == event {
			return t.To, nil
		}
	}
	return "", fmt.Errorf("no transition from %s on %s", sm.CurrentState, event)
}

func (sm *NodeInstanceGroupStateMachine) IsCooldownExpired() bool {
	return sm.now().Sub(sm.LastTransitionAt) >= sm.CooldownDuration
}

func (sm *NodeInstanceGroupStateMachine) TimeUntilCooldownExpires() time.Duration {
	elapsed := sm.now().Sub(sm.LastTransitionAt)
	if elapsed >= sm.CooldownDuration {
		return 0
	}
	return sm.CooldownDuration - elapsed
}

func (sm *NodeInstanceGroupStateMachine) RequiresFallback() bool {
	return sm.CurrentState == StateShortageDetected ||
		sm.CurrentState == StateFallbackCreating ||
		sm.CurrentState == StateFallbackReady
}

func (sm *NodeInstanceGroupStateMachine) HasActiveFallback() bool {
	return sm.CurrentState == StateFallbackReady ||
		sm.CurrentState == StateRecovering ||
		sm.CurrentState == StateDrainingFallback
}

func (sm *NodeInstanceGroupStateMachine) IsStable() bool {
	return sm.CurrentState == StateNormal ||
		sm.CurrentState == StateFallbackReady ||
		sm.CurrentState == StateDeprecated
}

func (sm *NodeInstanceGroupStateMachine) AllowedEvents() []Event {
	var events []Event
	for _, t := range validTransitions {
		if t.From == sm.CurrentState {
			events = append(events, t.Event)
		}
	}
	return events
}

func IsValidState(s State) bool {
	switch s {
	case StateNormal, StateShortageDetected, StateFallbackCreating,
		StateFallbackReady, StateRecovering, StateDrainingFallback,
		StateDeletingFallback, StateDeprecated:
		return true
	}
	return false
}

func StateFromString(s string) (State, error) {
	state := State(s)
	if !IsValidState(state) {
		return "", fmt.Errorf("invalid state: %s", s)
	}
	return state, nil
}