"""SDK-runtime operation tracking — the registry + watch stream.

The operation *model* — :class:`~mgf.vm.core.contracts.OperationHandle`,
:class:`~mgf.vm.core.contracts.OperationStatus`,
:class:`~mgf.vm.core.contracts.OperationKind` — lives in **mgf-vm-core**
(per the locked decision; ``api.md`` §4.2). It is the pure-data type that
round-trips between caller and backend; this module does NOT redefine it,
it imports and re-exports it.

What lives *here* is the SDK-side runtime that turns submit/watch into a
working two-step flow:

* :class:`OperationRegistry` — an in-memory, thread-safe store of in-flight
  and completed operations. Used by transports that own the op lifecycle
  locally (the in-process transport; test fakes). Transition methods
  (:meth:`mark_running`, :meth:`mark_completed`, …) advance a handle and
  **notify waiters** via a :class:`threading.Condition` — watchers block
  on the condition rather than polling.
* :class:`WatchOpStream` — the :class:`Iterator` a transport's
  ``watch_op`` returns. Each ``next()`` blocks on the registry's condition
  until the op's revision advances, then yields the new handle. Exhausts
  once a terminal status is yielded.

Error carriage (``api.md`` §7): a ``FAILED`` handle's
:attr:`~mgf.vm.core.contracts.OperationHandle.error` carries the **original**
core :class:`~mgf.vm.core.exceptions.VmManagerError` — no re-wrapping. The
client's sync-collapse (``submit_op`` → ``watch_op`` → unpack) re-raises that
exact exception so the caller sees the typed cause.
"""

from __future__ import annotations

import threading
from datetime import UTC, datetime
from typing import TYPE_CHECKING

# Re-export the model from core — single source of truth (locked decision).
from mgf.vm.core.contracts import (
    OperationHandle,
    OperationKind,
    OperationStatus,
)

if TYPE_CHECKING:
    from collections.abc import Callable, Iterator

    from mgf.vm.core.exceptions import VmManagerError
    from mgf.vm.core.ids import OperationId

__all__ = [
    "OperationHandle",
    "OperationKind",
    "OperationRegistry",
    "OperationStatus",
    "WatchOpStream",
]


class _Entry:
    """A registry slot: the current handle plus a monotonic revision.

    The revision increments on every transition; watchers compare against a
    last-seen revision to detect "something changed" without re-equality-checking
    the (frozen) handle. ``client_op_id`` is retained for dedup bookkeeping.
    """

    __slots__ = ("client_op_id", "handle", "revision")

    def __init__(self, handle: OperationHandle, client_op_id: str) -> None:
        self.handle = handle
        self.client_op_id = client_op_id
        self.revision = 0


class OperationRegistry:
    """In-memory, thread-safe store of operations with notify-on-transition.

    Concurrency model: a single :class:`threading.RLock` guards all state, and
    a :class:`threading.Condition` bound to that lock provides the wakeup
    channel. Mutators (:meth:`submit`, :meth:`mark_running`,
    :meth:`mark_completed`, :meth:`mark_failed`, :meth:`mark_cancelled`) take
    the lock, mutate, bump the entry revision, and ``notify_all`` — so any
    :class:`WatchOpStream` blocked in :meth:`wait_for_change` wakes exactly when
    a transition lands. **No polling, no sleeps.**

    Dedup (D-G5): :meth:`submit` keys on ``client_op_id``. A second submit with
    the same ``client_op_id`` returns the already-issued handle (the original
    op), never creating a duplicate.
    """

    def __init__(self) -> None:
        self._lock = threading.RLock()
        self._cond = threading.Condition(self._lock)
        self._by_op: dict[OperationId, _Entry] = {}
        self._by_client_op: dict[str, OperationId] = {}

    # -- submission --------------------------------------------------------

    def submit(
        self,
        op_id: OperationId,
        kind: OperationKind,
        *,
        client_op_id: str,
        started_at: datetime | None = None,
    ) -> OperationHandle:
        """Register a freshly-issued ``PENDING`` op, or return the dedup hit.

        If ``client_op_id`` was seen before, the existing handle is returned
        unchanged (idempotent submit; ``api.md`` §4.2). Otherwise a new
        ``PENDING`` handle is created, stored, and returned.
        """
        with self._lock:
            existing_op = self._by_client_op.get(client_op_id)
            if existing_op is not None:
                return self._by_op[existing_op].handle

            handle = OperationHandle.pending(op_id, kind, started_at or datetime.now(UTC))
            self._by_op[op_id] = _Entry(handle, client_op_id)
            self._by_client_op[client_op_id] = op_id
            return handle

    # -- queries -----------------------------------------------------------

    def get(self, op_id: OperationId) -> OperationHandle:
        """Return the current handle for ``op_id``.

        Raises :class:`KeyError` if the op is unknown to this registry.
        """
        with self._lock:
            return self._by_op[op_id].handle

    def list(self) -> list[OperationHandle]:
        """Snapshot of every known handle (order unspecified)."""
        with self._lock:
            return [entry.handle for entry in self._by_op.values()]

    def watch(self, op_id: OperationId) -> WatchOpStream:
        """Return a :class:`WatchOpStream` over ``op_id``'s transitions.

        Raises :class:`KeyError` if the op is unknown.
        """
        with self._lock:
            if op_id not in self._by_op:
                raise KeyError(op_id)
        return WatchOpStream(self, op_id)

    # -- transitions (each notifies waiters) -------------------------------

    def mark_running(self, op_id: OperationId) -> OperationHandle:
        """Transition ``PENDING`` → ``RUNNING`` and notify waiters."""
        return self._transition(op_id, lambda h: h.started())

    def mark_completed(
        self,
        op_id: OperationId,
        result: object,
        *,
        finished_at: datetime | None = None,
    ) -> OperationHandle:
        """Transition to ``SUCCEEDED`` with ``result`` and notify waiters."""
        return self._transition(
            op_id,
            lambda h: h.succeeded(result, finished_at=finished_at or datetime.now(UTC)),
        )

    def mark_failed(
        self,
        op_id: OperationId,
        error: VmManagerError,
        *,
        finished_at: datetime | None = None,
    ) -> OperationHandle:
        """Transition to ``FAILED`` carrying the ORIGINAL core ``error``.

        No re-wrapping (``api.md`` §7): the client's sync-collapse re-raises
        this exact exception.
        """
        return self._transition(
            op_id,
            lambda h: h.failed(error, finished_at=finished_at or datetime.now(UTC)),
        )

    def mark_cancelled(
        self,
        op_id: OperationId,
        *,
        finished_at: datetime | None = None,
    ) -> OperationHandle:
        """Transition to ``CANCELLED`` and notify waiters."""
        return self._transition(
            op_id,
            lambda h: h.cancelled(finished_at=finished_at or datetime.now(UTC)),
        )

    def _transition(
        self,
        op_id: OperationId,
        advance: Callable[[OperationHandle], OperationHandle],
    ) -> OperationHandle:
        with self._cond:
            entry = self._by_op[op_id]
            entry.handle = advance(entry.handle)
            entry.revision += 1
            self._cond.notify_all()
            return entry.handle

    # -- the wait primitive WatchOpStream blocks on ----------------------

    def wait_for_change(
        self,
        op_id: OperationId,
        last_revision: int,
    ) -> tuple[OperationHandle, int]:
        """Block until ``op_id``'s revision exceeds ``last_revision``.

        Returns ``(handle, revision)`` at the new revision. Wakes via the
        :class:`threading.Condition` the moment a transition method calls
        ``notify_all`` — never spins. Raises :class:`KeyError` if the op is
        unknown.
        """
        with self._cond:
            entry = self._by_op[op_id]
            self._cond.wait_for(lambda: entry.revision > last_revision)
            return entry.handle, entry.revision


class WatchOpStream:
    """The :class:`Iterator` ``watch_op`` returns — one yield per transition.

    Blocks on :meth:`OperationRegistry.wait_for_change` (Condition-backed, not
    polling) and yields each new :class:`OperationHandle` in revision order. On
    construction it captures the op's current revision so the *next* transition
    is the first thing yielded; the iterator exhausts after yielding a terminal
    handle (``SUCCEEDED`` / ``FAILED`` / ``CANCELLED``).
    """

    def __init__(self, registry: OperationRegistry, op_id: OperationId) -> None:
        self._registry = registry
        self._op_id = op_id
        # Seed from the current revision so we yield transitions, not the
        # state at subscribe time.
        handle, revision = registry.get(op_id), 0
        with registry._cond:
            revision = registry._by_op[op_id].revision
        self._handle = handle
        self._last_revision = revision
        self._terminal_yielded = False

    def __iter__(self) -> Iterator[OperationHandle]:
        return self

    def __next__(self) -> OperationHandle:
        if self._terminal_yielded:
            raise StopIteration
        handle, revision = self._registry.wait_for_change(self._op_id, self._last_revision)
        self._last_revision = revision
        self._handle = handle
        if handle.status.is_terminal:
            self._terminal_yielded = True
        return handle

    def close(self) -> None:
        """Release the stream. Idempotent; the in-memory stream holds no
        engine resources, so this just marks it exhausted."""
        self._terminal_yielded = True
