/src/free_claude_code/messaging/trees/runtime.py

"""Atomic runtime aggregate for one messaging conversation tree."""

import asyncio
from dataclasses import dataclass
from uuid import uuid4

from loguru import logger

from ..models import MessageScope
from .graph import MessageTreeGraph
from .identity import TreeIdentity
from .node import MessageNode, MessageReferenceKind, MessageState
from .queue import MessageNodeQueue
from .snapshot import TreeSnapshot
from .transitions import (
    AdmissionRejection,
    CompletionResult,
    FailureResult,
    MessageSubtreeRemoval,
    NodeClaim,
    NodeUiTarget,
    NodeView,
    QueueDecision,
    QueueEntry,
    ReplyTarget,
    TreeCancellation,
)


@dataclass(slots=True)
class _ActiveClaim:
    """Runtime execution identity kept separate from the node's UI state."""

    claim: NodeClaim
    cancellation_requested: bool = False


class MessageTree:
    """Own graph, queue, claim identity, and every concurrency invariant."""

    def __init__(
        self,
        root_node: MessageNode,
        *,
        graph: MessageTreeGraph | None = None,
    ) -> None:
        self._graph = graph or MessageTreeGraph(root_node)
        self._queue = MessageNodeQueue()
        self._lock = asyncio.Lock()
        self._active: _ActiveClaim | None = None
        self._restored_snapshot: TreeSnapshot | None = None
        self._restored_stale_targets: tuple[NodeUiTarget, ...] = ()
        logger.debug("Created MessageTree with root {}", self.root_id)

    @property
    def root_id(self) -> str:
        return self._graph.root_id

    @property
    def identity(self) -> TreeIdentity:
        return self._graph.identity

    @property
    def restored_snapshot(self) -> TreeSnapshot | None:
        """Normalized startup snapshot captured before the tree is published."""
        return self._restored_snapshot

    @property
    def restored_stale_targets(self) -> tuple[NodeUiTarget, ...]:
        """UI targets normalized from runnable to interrupted on restore."""
        return self._restored_stale_targets

    def _ui_target(self, node: MessageNode) -> NodeUiTarget:
        if node.status_message_id is None:
            raise ValueError("Runnable node has no status message")
        return NodeUiTarget(
            scope=node.scope,
            node_id=node.node_id,
            status_message_id=node.status_message_id,
        )

    def _queue_entries(self) -> tuple[QueueEntry, ...]:
        entries: list[QueueEntry] = []
        for node_id in self._queue.items():
            node = self._graph.get_node(node_id)
            if node is None or node.state is not MessageState.PENDING:
                continue
            entries.append(
                QueueEntry(node=self._ui_target(node), position=len(entries) + 1)
            )
        return tuple(entries)

    def _claim(self, node: MessageNode) -> NodeClaim:
        node.update_state(MessageState.IN_PROGRESS)
        claim = NodeClaim(
            identity=self.identity,
            claim_id=uuid4().hex,
            node=self._ui_target(node),
            prompt=node.prompt,
            parent_session_id=self._graph.get_parent_session_id(node.node_id),
        )
        self._active = _ActiveClaim(claim=claim)
        return claim

    def _enqueue_or_claim(self, node_id: str) -> QueueDecision:
        node = self._graph.get_node(node_id)
        if node is None or node.state is not MessageState.PENDING:
            return QueueDecision(
                claim=None,
                position=None,
                snapshot=None,
                rejection=AdmissionRejection.DUPLICATE,
            )

        if self._active is None:
            claim = self._claim(node)
            return QueueDecision(
                claim=claim,
                position=None,
                snapshot=self._graph.snapshot(),
            )

        if not self._queue.put(node_id):
            return QueueDecision(
                claim=None,
                position=None,
                snapshot=None,
                rejection=AdmissionRejection.DUPLICATE,
            )

        position = self._queue.qsize()
        logger.info("Queued node {}, position {}", node_id, position)
        return QueueDecision(
            claim=None,
            position=position,
            snapshot=self._graph.snapshot(),
        )

    async def enqueue_or_claim(self, node_id: str) -> QueueDecision:
        """Atomically reject, queue, or exclusively claim an existing node."""
        async with self._lock:
            return self._enqueue_or_claim(node_id)

    async def add_and_enqueue(
        self,
        node_id: str,
        scope: MessageScope,
        prompt: str,
        status_message_id: str,
        parent_id: str,
        parent_reference_id: str,
    ) -> QueueDecision:
        """Atomically add a reply and admit it to this tree."""
        async with self._lock:
            self._graph.add_node(
                node_id=node_id,
                scope=scope,
                prompt=prompt,
                status_message_id=status_message_id,
                parent_id=parent_id,
                parent_reference_id=parent_reference_id,
            )
            return self._enqueue_or_claim(node_id)

    async def finish_and_claim_next(self, claim_id: str) -> CompletionResult:
        """Release only the matching claim and atomically select its successor."""
        async with self._lock:
            if self._active is None or self._active.claim.claim_id != claim_id:
                return CompletionResult(
                    next_claim=None,
                    queue=self._queue_entries(),
                )

            self._active = None
            next_claim: NodeClaim | None = None
            while node_id := self._queue.pop():
                node = self._graph.get_node(node_id)
                if node is not None and node.state is MessageState.PENDING:
                    next_claim = self._claim(node)
                    break

            return CompletionResult(
                next_claim=next_claim,
                queue=self._queue_entries(),
            )

    async def cancel_node(
        self,
        node_id: str,
    ) -> TreeCancellation:
        """Atomically cancel one active, queued, or stale runnable node."""
        async with self._lock:
            node = self._graph.get_node(node_id)
            active_claim = (
                self._active.claim
                if self._active is not None
                and self._active.claim.node.node_id == node_id
                else None
            )
            if active_claim is not None:
                active = self._active
                if active is not None:
                    active.cancellation_requested = True
            if node is None:
                return TreeCancellation(
                    nodes=(),
                    active_claim=active_claim,
                    queue_update=None,
                )

            queue_changed = self._queue.remove(node_id)
            cancelled_nodes: tuple[NodeUiTarget, ...] = ()
            if node.state in (MessageState.PENDING, MessageState.IN_PROGRESS):
                node.mark_error()
                cancelled_nodes = (self._ui_target(node),)
            elif node.state is MessageState.ERROR and active_claim is not None:
                cancelled_nodes = (self._ui_target(node),)
            return TreeCancellation(
                nodes=cancelled_nodes,
                active_claim=active_claim,
                queue_update=self._queue_entries() if queue_changed else None,
            )

    async def cancel_all(
        self,
    ) -> TreeCancellation:
        """Atomically cancel every runnable node present at the transition."""
        async with self._lock:
            cancelled_nodes: list[NodeUiTarget] = []
            seen: set[str] = set()
            active_claim: NodeClaim | None = None

            if self._active is not None:
                active_claim = self._active.claim
                self._active.cancellation_requested = True
                active_node = self._graph.get_node(active_claim.node.node_id)
                if active_node is not None and active_node.state in (
                    MessageState.PENDING,
                    MessageState.IN_PROGRESS,
                ):
                    active_node.mark_error()
                    seen.add(active_node.node_id)
                    cancelled_nodes.append(self._ui_target(active_node))
                elif active_node is not None:
                    seen.add(active_node.node_id)
                    if active_node.state is MessageState.ERROR:
                        cancelled_nodes.append(self._ui_target(active_node))

            queued_ids = self._queue.drain()
            for node_id in queued_ids:
                node = self._graph.get_node(node_id)
                if node is None or node.state not in (
                    MessageState.PENDING,
                    MessageState.IN_PROGRESS,
                ):
                    continue
                node.mark_error()
                seen.add(node_id)
                cancelled_nodes.append(self._ui_target(node))

            for node in self._graph.all_nodes():
                if node.node_id in seen or node.state not in (
                    MessageState.PENDING,
                    MessageState.IN_PROGRESS,
                ):
                    continue
                node.mark_error()
                cancelled_nodes.append(self._ui_target(node))

            return TreeCancellation(
                nodes=tuple(cancelled_nodes),
                active_claim=active_claim,
                queue_update=() if queued_ids else None,
            )

    async def remove_message_subtree(
        self,
        reference_id: str,
    ) -> MessageSubtreeRemoval:
        """Atomically cancel and detach one literal platform reply subtree."""
        async with self._lock:
            resolved = self._graph.resolve_reference(reference_id)
            reference_ids = tuple(self._graph.get_reference_descendants(reference_id))
            if resolved is None or not reference_ids:
                empty = TreeCancellation(
                    nodes=(),
                    active_claim=None,
                    queue_update=None,
                )
                return MessageSubtreeRemoval(
                    cancellation=empty,
                    removed_message_ids=frozenset(),
                    removed_entire_tree=False,
                )

            owner, reference_kind = resolved
            removed_node_ids = {
                candidate
                for candidate in reference_ids
                if self._graph.get_node(candidate) is not None
            }
            affected_node_ids = set(removed_node_ids)
            if reference_kind is MessageReferenceKind.STATUS:
                affected_node_ids.add(owner.node_id)
            active_claim = (
                self._active.claim
                if self._active is not None
                and self._active.claim.node.node_id in affected_node_ids
                else None
            )
            if active_claim is not None:
                active = self._active
                if active is not None:
                    active.cancellation_requested = True
            cancelled_nodes: list[NodeUiTarget] = []
            queue_changed = False

            for node_id in affected_node_ids:
                node = self._graph.get_node(node_id)
                if node is None:
                    continue
                queue_changed = self._queue.remove(node_id) or queue_changed
                if node.state in (MessageState.PENDING, MessageState.IN_PROGRESS):
                    target = self._ui_target(node)
                    node.mark_error()
                    cancelled_nodes.append(target)
                elif (
                    node.state is MessageState.ERROR
                    and active_claim is not None
                    and active_claim.node.node_id == node_id
                ):
                    cancelled_nodes.append(self._ui_target(node))

            if reference_kind is MessageReferenceKind.STATUS:
                self._graph.clear_status(owner.node_id)
            removed_entire_tree = self.root_id in removed_node_ids
            self._graph.remove_nodes(removed_node_ids)
            cancellation = TreeCancellation(
                nodes=tuple(cancelled_nodes),
                active_claim=active_claim,
                queue_update=self._queue_entries() if queue_changed else None,
            )
            return MessageSubtreeRemoval(
                cancellation=cancellation,
                removed_message_ids=frozenset(reference_ids),
                removed_entire_tree=removed_entire_tree,
            )

    async def record_session(
        self, claim_id: str, session_id: str
    ) -> TreeSnapshot | None:
        """Record a real CLI session only for the currently active claim."""
        async with self._lock:
            if (
                self._active is None
                or self._active.claim.claim_id != claim_id
                or self._active.cancellation_requested
            ):
                return None
            node = self._graph.get_node(self._active.claim.node.node_id)
            if node is None or node.state is not MessageState.IN_PROGRESS:
                return None
            node.update_state(MessageState.IN_PROGRESS, session_id=session_id)
            return self._graph.snapshot()

    async def complete_claim(
        self, claim_id: str, session_id: str | None
    ) -> TreeSnapshot | None:
        """Mark the currently active claim complete."""
        async with self._lock:
            if (
                self._active is None
                or self._active.claim.claim_id != claim_id
                or self._active.cancellation_requested
            ):
                return None
            node = self._graph.get_node(self._active.claim.node.node_id)
            if node is None or node.state not in (
                MessageState.IN_PROGRESS,
                MessageState.ERROR,
            ):
                return None
            node.update_state(MessageState.COMPLETED, session_id=session_id)
            return self._graph.snapshot()

    async def fail_claim(
        self,
        claim_id: str,
        *,
        propagate: bool,
    ) -> FailureResult:
        """Atomically fail the active claim and its pending descendants."""
        async with self._lock:
            if (
                self._active is None
                or self._active.claim.claim_id != claim_id
                or self._active.cancellation_requested
            ):
                return FailureResult(affected=(), queue_update=None, snapshot=None)
            node = self._graph.get_node(self._active.claim.node.node_id)
            if node is None:
                return FailureResult(affected=(), queue_update=None, snapshot=None)

            affected: list[NodeUiTarget] = []
            queue_changed = False
            if node.state is not MessageState.COMPLETED:
                if node.state is not MessageState.ERROR:
                    node.mark_error()
                    affected.append(self._ui_target(node))

                if propagate:
                    for descendant_id in self._graph.get_descendants(node.node_id)[1:]:
                        child = self._graph.get_node(descendant_id)
                        if child is None or child.state is not MessageState.PENDING:
                            continue
                        child.mark_error()
                        queue_changed = (
                            self._queue.remove(child.node_id) or queue_changed
                        )
                        affected.append(self._ui_target(child))

            return FailureResult(
                affected=tuple(affected),
                queue_update=self._queue_entries() if queue_changed else None,
                snapshot=self._graph.snapshot(),
            )

    async def resolve_reply(self, reference_id: str) -> ReplyTarget | None:
        """Resolve a node/status reference without exposing the mutable graph."""
        async with self._lock:
            resolved = self._graph.resolve_reference(reference_id)
            if resolved is None:
                return None
            node, reference_kind = resolved
            return ReplyTarget(
                node_id=node.node_id,
                reference_id=reference_id,
                reference_kind=reference_kind,
                queue_position=(self._queue.qsize() + 1)
                if self._active is not None
                else None,
            )

    async def node_view(self, node_id: str) -> NodeView | None:
        """Return a copied node read model."""
        async with self._lock:
            node = self._graph.get_node(node_id)
            if node is None:
                return None
            return NodeView(
                identity=self.identity,
                node_id=node.node_id,
                state=node.state,
                parent_id=node.parent_id,
                session_id=node.session_id,
            )

    async def snapshot(self) -> TreeSnapshot:
        """Capture a detached persistence snapshot under the aggregate lock."""
        async with self._lock:
            return self._graph.snapshot()

    async def message_ids_for_chat(self, platform: str, chat_id: str) -> set[str]:
        """Copy every prompt and FCC status belonging to one platform chat."""
        async with self._lock:
            if self.identity.scope.platform != str(platform) or (
                self.identity.scope.chat_id != str(chat_id)
            ):
                return set()
            return self._graph.all_reference_ids()

    @classmethod
    def from_snapshot(cls, snapshot: TreeSnapshot) -> MessageTree:
        """Restore and reconcile interrupted nodes before publishing the tree."""
        graph = MessageTreeGraph.from_snapshot(snapshot)
        tree = cls(graph.get_root(), graph=graph)
        stale_targets: list[NodeUiTarget] = []
        for node in graph.all_nodes():
            if node.state in (MessageState.PENDING, MessageState.IN_PROGRESS):
                stale_targets.append(tree._ui_target(node))
                node.mark_error()
        tree._restored_stale_targets = tuple(stale_targets)
        tree._restored_snapshot = graph.snapshot()
        return tree


__all__ = ["MessageTree"]