from __future__ import annotations

import multiprocessing as mp
import sys
from datetime import datetime
from time import monotonic, sleep
from typing import TYPE_CHECKING

from yera.agents.registry import get_agent_name
from yera.events.models import OutputEvent
from yera.events.stream import EventStream, push_text_input

if TYPE_CHECKING:
    from rich.console import Console


class PyRuntimeExecutor:
    """Execute an agent in a separate process using an EventStream."""

    def __init__(self, agent, args, is_restricted_runtime: bool = False):
        """Initialise the runtime executor.

        Args:
            agent: The agent to execute.
            args: Positional arguments to pass to the agent's ``invoke`` method.
            is_restricted_runtime: Whether the runtime is restricted (built-in stream handler).
        """
        self.process: mp.Process | None = None
        self.agent = agent
        self.args = args
        self.is_restricted_runtime = is_restricted_runtime
        self.stream: EventStream | None = None

    @staticmethod
    def _agent_process(agent, stream: EventStream, *args):
        """Target function for the subprocess.

        Sets the current EventStream, marks the agent as top-level, and invokes it.
        """
        stream.set_current()
        agent.top_level = True
        agent.invoke(*args)

    def start(self) -> None:
        """Start the agent execution in a new process with an internally created EventStream."""
        if self.process is not None and self.process.is_alive():
            # Already running; no-op
            return

        ctx = mp.get_context("spawn")
        self.stream = EventStream.build()
        self.stream.is_restricted_stream = self.is_restricted_runtime
        self.process = ctx.Process(
            target=self._agent_process,
            args=(self.agent, self.stream, *self.args),
        )
        self.process.start()

    def stop(self) -> None:
        """Stop the agent process if it is running and clean up resources."""
        # Capture process locally to avoid race conditions
        process = self.process
        if process is None:
            return

        process.join(timeout=2)
        if process.is_alive():
            process.terminate()
            process.join()

        self.process = None
        self.stream = None

    def is_running(self) -> bool:
        """Check if the executor process is running.

        Returns:
            True if process is alive and stream is available
        """
        return (
            self.process is not None
            and self.process.is_alive()
            and self.stream is not None
        )

    def __enter__(self):
        """Start execution with a fresh EventStream and return self."""
        self.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Ensure the subprocess is terminated when leaving the context."""
        self.stop()
        return False


def stream_handler(
    metadata,
    *,
    console: Console | None = None,
    interactive: bool | None = None,
) -> OutputEvent:
    """Handle an event stream and render it to the console.

    Args:
        metadata: The metadata of the agent to run.
        console: The console to render to.
        interactive: Whether the stdin is interactive.
    """
    # Lazily load Console to defer heavy Rich imports
    from rich.console import Console as _Console

    from yera.console import ConsoleRenderer

    if console is None:
        console = _Console(force_jupyter=False)
    if interactive is None:
        interactive = sys.stdin.isatty()
    renderer = ConsoleRenderer(console, interactive=interactive)
    stream = EventStream.build()
    start_time_mono = monotonic()
    start_time_dt = datetime.now()

    # Use metadata to identify the top-level agent for startup panel
    top_level_agent_id = metadata.identifier
    top_level_agent_name = get_agent_name(top_level_agent_id) or metadata.name

    renderer.print_startup(start_time_dt, top_level_agent_name)

    exit_event = None

    while exit_event is None:
        for event in stream.iter_output():
            if event.block_type == "exit":
                exit_event = event
                break

            renderer.print_agent(event.agent_instance.agent_id)

            user_input = renderer.print_event(event)
            if user_input is not None:
                push_text_input(user_input)

        sleep(0.1)

    run_time_seconds = monotonic() - start_time_mono
    failing_agent_name = get_agent_name(exit_event.agent_instance.agent_id)
    renderer.print_exit(exit_event.data, run_time_seconds, failing_agent_name)
    list(stream.iter_output())
    return exit_event
