"""Protocol and helper for state machines driven by StateRunner.

This module defines the :class:`RunnableMachine` protocol that any
object must satisfy to be driven by :class:`StateRunner`, and a
convenience base class :class:`RunnableMachineHelper` that provides
default implementations of the threading utilities (pause, cancel,
progress, status).

Users subclass :class:`RunnableMachineHelper`, override
:meth:`get_machine`, and implement their long-running work inside
standard ``pytransitions`` callbacks (``on_enter_<state>``,
``before``, ``after``, etc.).

Example::

    from transitions import Machine
    from ox_ui.tkdantic.runnable import RunnableMachineHelper

    class MyPipeline(RunnableMachineHelper):
        def __init__(self):
            super().__init__()
            self._machine = Machine(
                model=self,
                states=['idle', 'running', 'done'],
                transitions=[
                    {'trigger': 'start', 'source': 'idle',
                     'dest': 'running'},
                ],
                initial='idle',
            )

        def get_machine(self):
            return self._machine

        def on_enter_running(self):
            for i in range(10):
                self.wait_if_paused()
                if self.is_cancelled():
                    self.set_status('Cancelled.')
                    return
                self.set_progress((i + 1) / 10)
                self.set_status(f'Step {i + 1}/10')
"""

import logging
import threading
from typing import Any, Optional, Protocol, runtime_checkable

from pydantic import BaseModel
from transitions import Machine

from ox_ui.tkdantic.callbacks import SimpleCallback

LOGGER = logging.getLogger(__name__)


# -------------------------------------------------------------------
# Protocol
# -------------------------------------------------------------------

@runtime_checkable
class RunnableMachine(Protocol):
    """Protocol for objects driven by :class:`StateRunner`.

    Implementors must provide access to a ``pytransitions``
    :class:`Machine`, a model object on which triggers are called,
    cooperative pause/resume/cancel methods, and progress/status
    reporting.
    """

    def get_machine(self) -> Machine:
        """Return the ``pytransitions`` Machine instance."""
        ...

    def get_model(self) -> Any:
        """Return the model object that owns trigger methods.

        In most cases this is ``self`` (the default in
        :class:`RunnableMachineHelper`).
        """
        ...

    def pause(self) -> None:
        """Request the current transition to pause."""
        ...

    def resume(self) -> None:
        """Resume a paused transition."""
        ...

    def cancel(self) -> None:
        """Request cooperative cancellation."""
        ...

    def get_progress(self) -> Optional[float]:
        """Return progress as a float in [0, 1] or None."""
        ...

    def get_status_text(self) -> str:
        """Return current human-readable status string."""
        ...

    def get_state_variables(self) -> Optional[BaseModel]:
        """Return a snapshot of the current state variables.

        Returns a pydantic model instance representing the
        current configurable state, or ``None`` if this machine
        has no user-configurable variables.

        The returned instance is a **snapshot**; callers should
        not mutate it.
        """
        ...

    def set_state_variables(self, variables: BaseModel) -> None:
        """Replace the current state variables wholesale.

        Callers should only invoke this when the machine is
        idle or paused.

        :param variables: validated pydantic model instance.
        """
        ...


# -------------------------------------------------------------------
# Helper base class
# -------------------------------------------------------------------

class RunnableMachineHelper:
    """Convenience base class satisfying :class:`RunnableMachine`.

    Provides thread-safe pause, resume, cancel, progress, and
    status utilities.  Subclasses must override :meth:`get_machine`.

    The threading helpers are designed to be called from within
    ``pytransitions`` callbacks (``on_enter_<state>``, ``before``,
    ``after``, etc.) which run on the worker thread spawned by
    :class:`StateRunner`.
    """

    def __init__(self) -> None:
        """Initialise threading primitives and state."""
        # Set means "running" (not paused).
        self._pause_event = threading.Event()
        self._pause_event.set()

        # Set means "cancel requested".
        self._cancel_event = threading.Event()

        # Guards reads/writes of _progress and _status_text.
        self._state_lock = threading.Lock()
        self._progress: Optional[float] = None
        self._status_text: str = ''
        self._status_headers: dict[str, str] = {}
        self._status_before_pause: str = ''

        # State variables: user-configurable pydantic model.
        self._state_vars: Optional[BaseModel] = None

        # Callbacks fired when state variables change.
        self._callbacks_lock = threading.Lock()
        self._state_changed_callbacks: dict[
            str, SimpleCallback] = {}

    # --- abstract (must override) --------------------------------

    def get_machine(self) -> Machine:
        """Return the pytransitions Machine.

        Subclasses **must** override this.
        """
        raise NotImplementedError(
            'Subclasses must implement get_machine().'
        )

    # --- defaults ------------------------------------------------

    def get_model(self) -> Any:
        """Return the model that owns trigger methods.

        Default returns ``self``, which is correct when the
        Machine was created with ``model=self``.
        """
        return self

    # --- pause / resume / cancel ---------------------------------

    def pause(self) -> None:
        """Request the worker to pause at the next safe point."""
        self._pause_event.clear()
        LOGGER.debug('Pause requested.')

    def resume(self) -> None:
        """Resume a paused worker."""
        self._pause_event.set()
        LOGGER.debug('Resume requested.')

    def cancel(self) -> None:
        """Request cooperative cancellation."""
        # Also resume in case we are paused, so the worker
        # thread unblocks and can check is_cancelled().
        self._cancel_event.set()
        self._pause_event.set()
        LOGGER.debug('Cancel requested.')

    # --- progress / status (thread-safe) -------------------------

    def get_progress(self) -> Optional[float]:
        """Return progress in [0, 1] or None if indeterminate."""
        with self._state_lock:
            return self._progress

    def get_status_text(self) -> str:
        """Return headers (if any) followed by the current status.

        The returned string is for display only; headers and the
        raw status text are stored separately.
        """
        with self._state_lock:
            parts = list(self._status_headers.values())
            if self._status_text:
                parts.append(self._status_text)
            return '\n'.join(parts)

    def set_progress(self, value: Optional[float]) -> None:
        """Set progress.  Call from pytransitions callbacks.

        :param value: float in [0, 1] or None for indeterminate.
        """
        with self._state_lock:
            self._progress = value

    def set_status(self, text: str) -> None:
        """Set the raw status text.  Call from pytransitions callbacks.

        This replaces only the raw portion; status headers are
        unaffected.

        :param text: human-readable status message.
        """
        with self._state_lock:
            self._status_text = text

    def append_status(self, text: str) -> None:
        """Append to the raw status text.

        Useful when a callback wants to add detail without
        overwriting what is already there.

        :param text: text to append (a newline is prepended
            automatically if the raw status is non-empty).
        """
        with self._state_lock:
            if self._status_text:
                self._status_text += '\n' + text
            else:
                self._status_text = text

    def add_status_header(self, name: str, text: str) -> None:
        """Add or replace a named status header.

        Headers are displayed above the raw status text by
        :meth:`get_status_text`.

        :param name: unique key for this header.
        :param text: header text to display.
        """
        with self._state_lock:
            self._status_headers[name] = text

    def remove_status_header(self, name: str) -> None:
        """Remove a status header by name.

        Silently ignores unknown names.

        :param name: key passed to :meth:`add_status_header`.
        """
        with self._state_lock:
            self._status_headers.pop(name, None)

    # --- worker-thread helpers -----------------------------------

    def wait_if_paused(self) -> None:
        """Block until resumed.  Call inside long-running loops.

        Updates the status text to indicate a paused state and
        restores the previous raw status on resume.  Status
        headers are not affected.
        """
        if self._pause_event.is_set():
            return
        with self._state_lock:
            previous = self._status_text
        self.set_status('Paused.')
        LOGGER.debug('Worker paused.')
        self._pause_event.wait()
        LOGGER.debug('Worker resumed.')
        self.set_status(previous)

    def is_cancelled(self) -> bool:
        """Return True if cancellation has been requested."""
        return self._cancel_event.is_set()

    def reset_run_state(self) -> None:
        """Reset pause/cancel/progress for a new transition.

        Called by :class:`StateRunner` before spawning the
        worker thread.
        """
        self._pause_event.set()
        self._cancel_event.clear()
        with self._state_lock:
            self._progress = None
            self._status_text = ''
            self._status_headers.clear()

    # --- state variables -----------------------------------------

    def get_state_variables(self) -> Optional[BaseModel]:
        """Return a deep copy of the current state variables.

        Returns ``None`` if this machine has no user-configurable
        variables (the default).  The returned instance is a
        **snapshot**; callers should not mutate it.
        """
        if self._state_vars is None:
            return None
        return self._state_vars.model_copy(deep=True)

    def set_state_variables(
        self, variables: BaseModel,
    ) -> None:
        """Replace the current state variables wholesale.

        A deep copy of *variables* is stored so the caller
        retains no shared reference.  After storing, all
        registered state-changed callbacks are executed.

        Callers should only invoke this when the machine is
        idle or paused.

        :param variables: validated pydantic model instance.
        """
        self._state_vars = variables.model_copy(deep=True)
        self.execute_state_changed_callbacks()

    # --- state-changed callbacks ---------------------------------

    def add_state_changed_callback(
        self, name: str, callback: SimpleCallback,
    ) -> None:
        """Register a named callback for state-variable changes.

        If *name* already exists it is replaced.

        :param name: unique key for this callback.
        :param callback: :class:`SimpleCallback` to invoke.
        """
        with self._callbacks_lock:
            self._state_changed_callbacks[name] = callback

    def remove_state_changed_callback(self, name: str) -> None:
        """Remove a previously registered callback by name.

        Silently ignores unknown names.

        :param name: the key passed to
            :meth:`add_state_changed_callback`.
        """
        with self._callbacks_lock:
            self._state_changed_callbacks.pop(name, None)

    def execute_state_changed_callbacks(self) -> None:
        """Run all registered state-changed callbacks.

        The callback dict is snapshotted under the lock and
        then released **before** any callback executes.  This
        prevents deadlocks when a callback calls
        :meth:`get_state_variables` or registers / removes
        other callbacks.

        Exceptions in individual callbacks are logged and
        do not prevent remaining callbacks from running.
        """
        with self._callbacks_lock:
            snapshot = list(
                self._state_changed_callbacks.values(),
            )
        for cb in snapshot:
            try:
                cb.function(
                    *(cb.args or []),
                    **(cb.kwargs or {}),
                )
            except Exception:
                LOGGER.exception(
                    'Error in state-changed callback.',
                )
