from contextvars import ContextVar
from typing import Optional

from yera.agents.dataclasses import AgentMetadata

_current_agent_context: ContextVar[Optional["AgentContext"]] = ContextVar(
    "agent_context", default=None
)


class AgentContext:
    def __init__(self, metadata: AgentMetadata):
        self.metadata = metadata
        self.instance_number = 0
        self._token = None
        self._graph = None
        self._is_top_of_agent_stack = None

    def __enter__(self):
        self._is_top_of_agent_stack = _current_agent_context.get() is None
        self._token = _current_agent_context.set(self)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        _current_agent_context.reset(self._token)

    def top_level(self) -> bool:
        if self._is_top_of_agent_stack is None:
            raise RuntimeError("This context has not been initialised")
        return self._is_top_of_agent_stack


def no_active_agent_context():
    return _current_agent_context.get() is None


def get_agent_context() -> "AgentContext":
    agent = _current_agent_context.get()
    if agent is None:
        raise RuntimeError("No agent context")
    return agent
