// Package runtime owns shared client session state. It contains no socket or
// OS calls, so supervisor workers can outlive a replaceable network session.
package runtime

import (
	"errors"
	"math"
	"sync"
	"time"
)

var ErrInvalidTransition = errors.New("invalid client session transition")

type State uint8

const (
	StateBackoff State = iota + 1
	StateConnecting
	StateHello
	StateReconciling
	StateActive
	StateClosing
)

type BackoffOptions struct {
	Initial     time.Duration
	Maximum     time.Duration
	StableReset time.Duration
}

func (options BackoffOptions) Validate() error {
	if options.Initial <= 0 || options.Maximum < options.Initial || options.StableReset <= 0 {
		return errors.New("invalid client reconnect backoff options")
	}
	return nil
}

// Jitter returns a delay in [0, cap]. It is injected so all backoff boundaries
// are testable without sleeps or probabilistic assertions.
type Jitter func(cap time.Duration) time.Duration

type SessionMachine struct {
	mu          sync.Mutex
	options     BackoffOptions
	jitter      Jitter
	state       State
	failures    uint32
	nextConnect time.Duration
	activeSince time.Duration
}

func NewSessionMachine(options BackoffOptions, jitter Jitter) (*SessionMachine, error) {
	if err := options.Validate(); err != nil {
		return nil, err
	}
	if jitter == nil {
		return nil, errors.New("client reconnect jitter is required")
	}
	return &SessionMachine{options: options, jitter: jitter, state: StateBackoff}, nil
}

func (machine *SessionMachine) State() State {
	machine.mu.Lock()
	defer machine.mu.Unlock()
	return machine.state
}

func (machine *SessionMachine) NextConnectAt() time.Duration {
	machine.mu.Lock()
	defer machine.mu.Unlock()
	return machine.nextConnect
}

func (machine *SessionMachine) StartConnecting(now time.Duration) error {
	machine.mu.Lock()
	defer machine.mu.Unlock()
	if machine.state != StateBackoff || now < machine.nextConnect {
		return ErrInvalidTransition
	}
	machine.state = StateConnecting
	return nil
}

func (machine *SessionMachine) TransportConnected() error {
	return machine.transition(StateConnecting, StateHello)
}

func (machine *SessionMachine) Welcomed() error {
	return machine.transition(StateHello, StateReconciling)
}

func (machine *SessionMachine) Reconciled(now time.Duration) error {
	machine.mu.Lock()
	defer machine.mu.Unlock()
	if machine.state != StateReconciling {
		return ErrInvalidTransition
	}
	machine.state = StateActive
	machine.activeSince = now
	return nil
}

func (machine *SessionMachine) BeginClosing() error {
	machine.mu.Lock()
	defer machine.mu.Unlock()
	if machine.state == StateBackoff || machine.state == StateClosing {
		return ErrInvalidTransition
	}
	machine.state = StateClosing
	return nil
}

// Failed moves every in-flight session state to backoff. Only a continuous
// active period reaching StableReset clears exponential history.
func (machine *SessionMachine) Failed(now time.Duration) time.Duration {
	machine.mu.Lock()
	defer machine.mu.Unlock()
	if machine.state == StateActive && now >= machine.activeSince && now-machine.activeSince >= machine.options.StableReset {
		machine.failures = 0
	}
	if machine.failures < math.MaxUint32 {
		machine.failures++
	}
	cap := machine.capLocked()
	delay := machine.jitter(cap)
	if delay < 0 {
		delay = 0
	}
	if delay > cap {
		delay = cap
	}
	machine.state = StateBackoff
	machine.nextConnect = now + delay
	machine.activeSince = 0
	return delay
}

func (machine *SessionMachine) transition(from, to State) error {
	machine.mu.Lock()
	defer machine.mu.Unlock()
	if machine.state != from {
		return ErrInvalidTransition
	}
	machine.state = to
	return nil
}

func (machine *SessionMachine) capLocked() time.Duration {
	cap := machine.options.Initial
	for index := uint32(1); index < machine.failures && cap < machine.options.Maximum; index++ {
		if cap > machine.options.Maximum/2 {
			return machine.options.Maximum
		}
		cap *= 2
	}
	if cap > machine.options.Maximum {
		return machine.options.Maximum
	}
	return cap
}