"""``VmLifecycleStateMachine`` — pure-Python lifecycle logic.

The brain that knows which transitions are legal, drives them
serially per VM, and fires :class:`LifecycleObserver` /
:class:`AsyncLifecycleObserver` on every legal transition.

The state machine itself is **engine-neutral**. It doesn't know
libvirt or Firecracker; it only knows *states* and *legal
transitions between them*. The actual work (starting a VM, taking
a snapshot) is the adapter's job, driven by the CP's reconciler.

Per ``contracts.md`` §6:

* Observers are **non-blocking** (<50 ms budget); exceptions are
  caught and logged but never propagated to the transition caller.
* The state machine is **NOT thread-safe** — one instance per VM,
  caller serializes (the CP achieves this by per-``VmId``
  reconciliation work queues).
* Observer order is registration order; observers do not declare
  dependencies on each other.
"""

from __future__ import annotations

import asyncio
import inspect
import logging
import time
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import TYPE_CHECKING

from mgf.vm.core.contracts import (
    AsyncLifecycleObserver,
    LifecycleObserver,
    TransitionCause,
    TransitionEvent,
    VmState,
)

if TYPE_CHECKING:
    from collections.abc import Callable

    from mgf.vm.core.ids import OperationId, VmId

LOG = logging.getLogger(__name__)


# Each transition (from → to) is either legal (in the table below)
# or illegal (raises :class:`InvalidTransition` on attempt).
#
# Empty set for ``DEFINED`` means: from DEFINED you can go to any
# named state (handled below). The table is conservative — only the
# transitions the lifecycle explicitly supports.
_LEGAL_TRANSITIONS: dict[VmState | None, frozenset[VmState]] = {
    None: frozenset({VmState.DEFINED}),  # initial define
    VmState.DEFINED: frozenset({VmState.STARTING, VmState.STOPPED}),
    VmState.STARTING: frozenset({VmState.RUNNING, VmState.FAILED}),
    VmState.RUNNING: frozenset(
        {
            VmState.PAUSING,
            VmState.STOPPING,
            VmState.MIGRATING,
            VmState.FAILED,
        },
    ),
    VmState.PAUSING: frozenset({VmState.PAUSED, VmState.FAILED}),
    VmState.PAUSED: frozenset({VmState.RESUMING, VmState.STOPPING}),
    VmState.RESUMING: frozenset({VmState.RUNNING, VmState.FAILED}),
    VmState.STOPPING: frozenset({VmState.STOPPED, VmState.FAILED}),
    VmState.STOPPED: frozenset({VmState.STARTING, VmState.DEFINED}),
    VmState.MIGRATING: frozenset({VmState.RUNNING, VmState.FAILED}),
    VmState.FAILED: frozenset({VmState.STOPPING, VmState.STOPPED, VmState.DEFINED}),
}


# Observer budget per contracts.md §6 — 50 ms. Exceedances are
# logged but don't fail the transition (the state machine refuses
# to be brittle in the presence of slow hooks).
_OBSERVER_BUDGET_S: float = 0.050


class InvalidTransition(Exception):
    """Raised when a caller tries an illegal state transition.

    Carries enough context to render a useful error message in CP
    logs and surface to users via the SDK.
    """

    def __init__(self, vm_id: VmId, from_state: VmState | None, to_state: VmState) -> None:
        self.vm_id = vm_id
        self.from_state = from_state
        self.to_state = to_state
        super().__init__(
            f"VM {vm_id}: illegal transition {from_state} -> {to_state}",
        )


@dataclass(slots=True)
class _ObserverEntry:
    """One registered observer + classification (sync vs. async)."""

    observer: LifecycleObserver | AsyncLifecycleObserver
    is_async: bool


class VmLifecycleStateMachine:
    """Per-VM lifecycle state machine.

    One instance per VM. The state is held internally; transitions
    fire registered observers. Observer exceptions are caught and
    logged — never propagated to ``transition`` callers.

    Usage::

        sm = VmLifecycleStateMachine(vm_id=VmId("vm-abc"))
        sm.add_observer(audit_logger)
        sm.add_observer(reconciler_drift_detector)
        sm.transition(VmState.DEFINED, cause=TransitionCause.USER_REQUEST)
        sm.transition(VmState.STARTING, cause=TransitionCause.USER_REQUEST)
        # ... adapter does the actual start work ...
        sm.transition(VmState.RUNNING, cause=TransitionCause.ENGINE_EVENT)
    """

    def __init__(
        self,
        vm_id: VmId,
        *,
        clock: Callable[[], datetime] = lambda: datetime.now(UTC),
    ) -> None:
        self._vm_id = vm_id
        self._clock = clock
        self._state: VmState | None = None
        self._observers: list[_ObserverEntry] = []
        self._history: list[TransitionEvent] = []

    @property
    def vm_id(self) -> VmId:
        return self._vm_id

    @property
    def state(self) -> VmState | None:
        """``None`` before the first transition (the VM hasn't been
        defined yet)."""
        return self._state

    @property
    def history(self) -> list[TransitionEvent]:
        """Read-only view of every transition this VM has been through.

        Returned list is a copy; callers may not mutate it.
        """
        return list(self._history)

    # --- observers --------------------------------------------------------

    def add_observer(self, observer: LifecycleObserver | AsyncLifecycleObserver) -> None:
        """Register an observer. Fired on every transition in
        registration order. Idempotent on identity — adding the same
        observer twice is a no-op."""
        for entry in self._observers:
            if entry.observer is observer:
                return
        is_async = inspect.iscoroutinefunction(observer.on_transition)
        self._observers.append(_ObserverEntry(observer=observer, is_async=is_async))

    def remove_observer(self, observer: LifecycleObserver | AsyncLifecycleObserver) -> None:
        """Unregister an observer. Idempotent on absence."""
        self._observers = [e for e in self._observers if e.observer is not observer]

    # --- transitions ------------------------------------------------------

    def transition(
        self,
        to_state: VmState,
        *,
        cause: TransitionCause,
        operation_id: OperationId | None = None,
        metadata: dict[str, str] | None = None,
    ) -> TransitionEvent:
        """Drive the machine to ``to_state``.

        Validates the transition against the legal-transitions table;
        raises :class:`InvalidTransition` if disallowed.

        On success, fires every observer in registration order. Per
        ``contracts.md`` §6, observer exceptions are caught and logged
        — never propagated. Slow observers (>50 ms) are logged as
        warnings but their results still count.

        Returns the :class:`TransitionEvent` that was fired (also
        appended to :attr:`history`).
        """
        legal_targets = _LEGAL_TRANSITIONS.get(self._state, frozenset())
        if to_state not in legal_targets:
            raise InvalidTransition(self._vm_id, self._state, to_state)

        event = TransitionEvent(
            vm_id=self._vm_id,
            from_state=self._state,
            to_state=to_state,
            timestamp=self._clock(),
            cause=cause,
            operation_id=operation_id,
            metadata=dict(metadata) if metadata is not None else {},
        )
        # Apply state THEN fire observers — observers see consistent
        # state when they query the machine.
        self._state = to_state
        self._history.append(event)
        self._fire_observers(event)
        return event

    def force_state(
        self,
        new_state: VmState,
        *,
        cause: TransitionCause = TransitionCause.RECONCILER,
        metadata: dict[str, str] | None = None,
    ) -> TransitionEvent:
        """Bypass the legal-transitions table — reconciler-only.

        Used when the engine reports an out-of-band state change
        (e.g., a VM crashed on the host; the host agent observed it
        and the reconciler must converge to ``FAILED`` even from a
        state that wouldn't normally allow that transition).

        Still fires observers. Callers (the reconciler) take full
        responsibility for the integrity of the transition; the
        machine trusts them.
        """
        event = TransitionEvent(
            vm_id=self._vm_id,
            from_state=self._state,
            to_state=new_state,
            timestamp=self._clock(),
            cause=cause,
            operation_id=None,
            metadata=dict(metadata) if metadata is not None else {"forced": "true"},
        )
        self._state = new_state
        self._history.append(event)
        self._fire_observers(event)
        return event

    # --- observer dispatch ------------------------------------------------

    def _fire_observers(self, event: TransitionEvent) -> None:
        """Fire each registered observer.

        Exceptions are caught + logged; budget exceedances are logged
        as warnings. Async observers are awaited on the current event
        loop if one is running; if not, they're scheduled via
        ``asyncio.run`` (single-shot) — adapters running fully
        synchronously won't typically register async observers.
        """
        for entry in self._observers:
            start = time.monotonic()
            try:
                if entry.is_async:
                    self._dispatch_async(entry.observer, event)
                else:
                    entry.observer.on_transition(event)
            except Exception:
                LOG.exception(
                    "observer %s raised on transition %s -> %s",
                    type(entry.observer).__name__,
                    event.from_state,
                    event.to_state,
                )
            finally:
                duration_s = time.monotonic() - start
                if duration_s > _OBSERVER_BUDGET_S:
                    LOG.warning(
                        "observer %s exceeded %dms budget: took %.1fms on %s -> %s",
                        type(entry.observer).__name__,
                        int(_OBSERVER_BUDGET_S * 1000),
                        duration_s * 1000,
                        event.from_state,
                        event.to_state,
                    )

    @staticmethod
    def _dispatch_async(
        observer: LifecycleObserver | AsyncLifecycleObserver,
        event: TransitionEvent,
    ) -> None:
        """Invoke an async observer.

        If we're inside a running event loop, schedule a task and
        await its completion via ``run_coroutine_threadsafe`` (the
        state machine is not async-aware).

        If no event loop is running, use ``asyncio.run`` for a
        bounded single-shot.
        """
        coro = observer.on_transition(event)
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            asyncio.run(coro)  # type: ignore[arg-type]
            return
        # We're inside an event loop already. Schedule and block.
        # NOTE: this is a deliberate choice — the state machine's
        # transition() is synchronous by design (per contracts.md §6),
        # and we don't want to leak coroutines that never complete.
        # Adapters running inside an async stack should NOT use
        # async observers paired with synchronous .transition() calls;
        # they should call .atransition() (TBD; not in this iteration).
        future = asyncio.run_coroutine_threadsafe(coro, loop)  # type: ignore[arg-type]
        future.result(timeout=_OBSERVER_BUDGET_S)


__all__ = [
    "InvalidTransition",
    "VmLifecycleStateMachine",
]
