import enum
import functools
import inspect
import time

from prometheus_client import Counter
from prometheus_client import Gauge

_worker_activity = Counter(
    'worker_activity_time_seconds_total',
    'Total seconds spent in each state and operation',
    ['state', 'operation'],
)
_worker_current_state = Gauge(
    'worker_current_state',
    'Current state indicator: 0=Idle, 1=IO, 2=CPU, 3=GPU, 4=Warmup, 5=Startup, 6=Shutdown',
)


class AppState(enum.Enum):
    IDLE = 'idle'
    IO = 'io'
    CPU = 'cpu'
    GPU = 'gpu'
    WARMUP = 'warmup'
    STARTUP = 'startup'
    SHUTDOWN = 'shutdown'


STATE_VALUE_MAP = {
    AppState.IDLE: 0,
    AppState.IO: 1,
    AppState.CPU: 2,
    AppState.GPU: 3,
    AppState.WARMUP: 4,
    AppState.STARTUP: 5,
    AppState.SHUTDOWN: 6,
}


class StateTracker:
    """Ensures mutual exclusivity. The app can only be in ONE state at a time."""

    def __init__(self, enabled: bool = False):
        self._enabled = enabled
        self._current_state = AppState.STARTUP
        self._current_operation = None
        self._last_switch_time = time.perf_counter()
        self.update_gauge()

    def set_state(self, new_state: AppState, operation: str = None):
        if not self._enabled:
            return

        if new_state == self._current_state and operation == self._current_operation:
            return

        now = time.perf_counter()
        duration = now - self._last_switch_time

        _worker_activity.labels(
            state=self._current_state.value,
            operation=self._current_operation if self._current_operation else '',
        ).inc(duration)

        self._current_state = new_state
        self._current_operation = operation
        self._last_switch_time = now
        self.update_gauge()

    def update_gauge(self, state: AppState = None):
        if not self._enabled:
            return
        if state is None:
            state = self._current_state
        _worker_current_state.set(STATE_VALUE_MAP[state])

    def scope(self, state: AppState, operation: str = None):
        return StateScope(self, state, operation)

    def finalize(self):
        """Record the final state duration before shutdown."""
        if not self._enabled:
            return
        now = time.perf_counter()
        duration = now - self._last_switch_time
        _worker_activity.labels(
            state=self._current_state.value,
            operation=self._current_operation if self._current_operation else '',
        ).inc(duration)


class StateScope:
    """Helper to automatically revert previous state and operation"""

    def __init__(self, tracker, state, operation=None):
        self._tracker = tracker
        self._target_state = state
        self._target_operation = operation
        self._previous_state = tracker._current_state
        self._previous_operation = tracker._current_operation

    def __enter__(self):
        self._tracker.set_state(self._target_state, self._target_operation)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._tracker.set_state(self._previous_state, self._previous_operation)
        return False


class _GlobalTracker:
    def __init__(self):
        self._tracker = StateTracker(enabled=False)

    def enable(self, rank: int = 0):
        self._tracker = StateTracker(enabled=(rank == 0))

    def __getattr__(self, name):
        return getattr(self._tracker, name)


state_tracker = _GlobalTracker()


def enable_tracker(rank: int = 0):
    """Enable activity tracking for rank 0 only."""
    state_tracker.enable(rank)


def track_state(state: AppState, operation: str = None):
    def decorator(func):
        op_name = operation or func.__qualname__

        if inspect.iscoroutinefunction(func):

            @functools.wraps(func)
            async def async_wrapper(*args, **kwargs):
                with state_tracker.scope(state, op_name):
                    return await func(*args, **kwargs)

            return async_wrapper
        else:

            @functools.wraps(func)
            def sync_wrapper(*args, **kwargs):
                with state_tracker.scope(state, op_name):
                    return func(*args, **kwargs)

            return sync_wrapper

    return decorator
