import multiprocessing as mp
from collections.abc import Iterator
from contextvars import ContextVar
from typing import Optional

from yera.events.models import FailureExitEventData

from .models.in_event import InputEvent
from .models.out_event import OutputEvent

_event_stream_context: ContextVar[Optional["EventStream"]] = ContextVar(
    "event_stream", default=None
)


def push_output(event: OutputEvent):
    stream = EventStream.get_current()
    stream.push_output(event)
    if event.block_type == "exit" and isinstance(event.data, FailureExitEventData):
        stream.failure_exit_event_pushed = True


def await_input(timeout=None) -> InputEvent:
    stream = EventStream.get_current()
    return stream.pop_input(timeout)


def push_input(in_event: InputEvent):
    stream = EventStream.get_current()
    stream.push_input(in_event)


def push_text_input(text: str):
    in_event = InputEvent(identifier="ABC", type="text", data=text)
    push_input(in_event)


def is_restricted_stream() -> bool:
    """Return whether the current event stream is restricted (only certain block types allowed)."""
    stream = EventStream.get_current()
    return stream.is_restricted_stream


def failure_exit_event_already_pushed() -> bool:
    """Return True if a nested agent already pushed a failure exit event."""
    stream = EventStream.get_current()
    return stream.failure_exit_event_pushed


class EventStream:
    def __init__(self, out_event_queue, in_event_queue):
        self.out_event_queue = out_event_queue
        self.in_event_queue = in_event_queue
        self.is_restricted_stream: bool = False
        self.failure_exit_event_pushed: bool = False
        super().__init__()

    def push_output(self, event: OutputEvent) -> None:
        if not isinstance(event, OutputEvent):
            raise TypeError(f"Expected OutputEvent got {type(event).__name__}")
        self.out_event_queue.put(event)

    def push_input(self, in_event: InputEvent) -> None:
        if not isinstance(in_event, InputEvent):
            raise TypeError(f"Expected InputEvent got {type(in_event).__name__}")
        self.in_event_queue.put(in_event)

    def pop_output(self, timeout: float | None = None) -> OutputEvent:
        return self.out_event_queue.get(timeout=timeout)

    def pop_input(self, timeout: float | None = None) -> InputEvent:
        return self.in_event_queue.get(timeout=timeout)

    def has_output(self) -> bool:
        return not self.out_event_queue.empty()

    def has_input(self) -> bool:
        return not self.in_event_queue.empty()

    def iter_output(self) -> Iterator[OutputEvent]:
        while self.has_output():
            yield self.pop_output()

    def set_current(self):
        if _event_stream_context.get() is not None:
            raise RuntimeError("Event stream is already set in context.")
        _event_stream_context.set(self)

    @classmethod
    def get_current(cls):
        reg = _event_stream_context.get()
        if reg is None:
            raise RuntimeError(
                "No event stream has been set in current context. "
                "Call EventStream's set_current method first."
            )
        return reg

    @staticmethod
    def build() -> "EventStream":
        reg = _event_stream_context.get()
        if reg is not None:
            return reg

        ctx = mp.get_context("spawn")
        manager = ctx.Manager()
        stream = EventStream(manager.Queue(), manager.Queue())
        stream.set_current()
        return stream
