"""Runner — the zero-boilerplate entry point for AlbusOS.

Simplest possible usage (from the ``albusos`` package)::

    from albusos import agent, run

    a = agent("researcher", instructions="Research topics.", tools=["web.*"])
    result = await run(a, "What is quantum computing?")

For multi-turn conversations use ``Session``::

    from albusos import Session

    session = Session(agent_def)
    r1 = await session.run("Hello")
    r2 = await session.run("Follow up")

The runner wires PathwayService, PathwayVM, Context, ServiceContainer,
and StateStore internally.  Callers never touch those objects.
"""

from __future__ import annotations

import asyncio
import contextlib
import uuid
from typing import TYPE_CHECKING

from core.agent import Agent
from core.builders.pathway import PathwayBuilder
from core.pathways.context import Context
from core.pathways.service import PathwayService
from core.pathways.service_container import ServiceContainer
from core.pathways.vm.executor import PathwayVM
from core.state import FileStateStore
from core.tools import get_tool_registry
from core.types.execution import ExecutionMode

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from core.pathways.pathway import Pathway
    from core.protocols.services import StateStoreLike
    from core.types.agents import AgentDefinition
    from core.types.execution import ExecutionBudget, ExecutionResult
    from core.types.stream import StreamEvent

# Sentinel pathway ID used when the agent doesn't specify a custom pathway.
DEFAULT_PATHWAY_ID = "__default__"


# =============================================================================
# Default pathway — the universal tool-calling agent loop
# =============================================================================


def _build_default_pathway() -> Pathway:
    """Build the default tool-calling pathway.

    This is a single LLMNode whose effective config (model, tools, system
    prompt, max_tool_rounds) comes from RunState — which ``Agent.turn()``
    populates from the ``AgentDefinition``.  So one pathway works for
    every agent.

    Graph::

        input ──▶ llm ("{{input.goal}}") ──▶ output
    """
    return (
        PathwayBuilder("default-agent-loop", pathway_id=DEFAULT_PATHWAY_ID)
        .description("Universal tool-calling agent loop")
        .llm("llm", "{{input.goal}}")
        .output({"response": "{{llm.response}}"})
        .connect("input", "llm")
        .connect("llm", "output")
        .build()
    )


# =============================================================================
# run() — single-shot convenience
# =============================================================================


async def run(
    agent_or_definition: Agent | AgentDefinition,
    message: str,
    *,
    pathway: Pathway | None = None,
    thread_id: str | None = None,
    budget: ExecutionBudget | None = None,
    state_store: StateStoreLike | None = None,
    mode: ExecutionMode = ExecutionMode.CONVERSATIONAL,
) -> ExecutionResult:
    """Execute a single agent turn with zero manual wiring.

    Parameters
    ----------
    agent_or_definition
        An ``AgentDefinition`` (from ``AgentBuilder.build()`` or
        ``agent()``) or an already-constructed ``Agent``.
    message
        The user message for this turn.
    pathway
        A ``Pathway`` to deploy.  When omitted and the agent uses the
        default pathway ID, a standard tool-calling pathway is created
        automatically.
    thread_id
        Conversation thread id.  Auto-generated when omitted.
    budget
        Execution budget (tokens, cost, duration limits).
    state_store
        Custom state store.  Defaults to ``FileStateStore()``.
    mode
        Execution mode — influences streaming, checkpointing, and error
        handling defaults.  Defaults to ``CONVERSATIONAL``.

    Returns
    -------
    ExecutionResult
        The agent's response, metrics, and status.
    """
    session = Session(
        agent_or_definition,
        pathway=pathway,
        state_store=state_store,
        mode=mode,
    )
    return await session.run(message, thread_id=thread_id, budget=budget)


# =============================================================================
# Session — multi-turn convenience
# =============================================================================


class Session:
    """Reusable session that holds wiring across multiple ``run()`` calls.

    Typical usage::

        session = Session(agent_def)
        r1 = await session.run("Hello")
        r2 = await session.run("Tell me more")

    The session creates one ``PathwayService``, ``PathwayVM``,
    ``Context``, and ``StateStore`` internally, so callers never
    touch those objects.
    """

    def __init__(
        self,
        agent_or_definition: Agent | AgentDefinition,
        *,
        pathway: Pathway | None = None,
        state_store: StateStoreLike | None = None,
        mode: ExecutionMode = ExecutionMode.CONVERSATIONAL,
    ) -> None:
        self._mode = mode
        store = state_store or FileStateStore()

        # --- PathwayService + deploy ---
        self._pathway_service = PathwayService()
        if pathway is not None:
            self._pathway_service.deploy(pathway)

        # Auto-deploy the default pathway when the agent references it.
        defn = (
            agent_or_definition.definition
            if isinstance(agent_or_definition, Agent)
            else agent_or_definition
        )
        if (
            defn.pathway == DEFAULT_PATHWAY_ID
            and self._pathway_service.load(DEFAULT_PATHWAY_ID) is None
        ):
            self._pathway_service.deploy(_build_default_pathway())

        # --- Tool handlers from global registry ---
        registry = get_tool_registry()
        tool_handlers = registry.as_handler_dict()

        # --- ServiceContainer ---
        from core.agent import get_agent_repository

        services = ServiceContainer(
            pathway_service=self._pathway_service,
            state_store=store,
            agent_repo=get_agent_repository(),
        )

        # --- Context + PathwayVM ---
        ctx = Context(tools=tool_handlers, services=services)
        self._vm = PathwayVM(ctx)

        # --- Agent ---
        if isinstance(agent_or_definition, Agent):
            self._agent = agent_or_definition
        else:
            self._agent = Agent(
                definition=agent_or_definition,
                state_store=store,
            )

        # Default thread id for the session (stable across calls)
        self._default_thread_id = f"session-{uuid.uuid4().hex[:12]}"

    # -- public API --

    @property
    def agent(self) -> Agent:
        """The ``Agent`` backing this session."""
        return self._agent

    @property
    def mode(self) -> ExecutionMode:
        """The session's execution mode."""
        return self._mode

    @property
    def pathway_service(self) -> PathwayService:
        """The session's ``PathwayService`` (for deploying additional pathways)."""
        return self._pathway_service

    async def run(
        self,
        message: str,
        *,
        thread_id: str | None = None,
        budget: ExecutionBudget | None = None,
        pathway_override: str | None = None,
    ) -> ExecutionResult:
        """Execute one conversational turn.

        Parameters
        ----------
        message
            User message.
        thread_id
            Thread id.  Reuses the session-level default when omitted,
            giving you multi-turn memory for free.
        budget
            Execution budget overrides.
        pathway_override
            Use a different pathway for this turn (must already be
            deployed on the session's ``PathwayService``).

        Returns
        -------
        ExecutionResult
        """
        return await self._agent.turn(
            message,
            thread_id or self._default_thread_id,
            pathway_vm=self._vm,
            pathway_service=self._pathway_service,
            budget=budget,
            pathway_override=pathway_override,
        )

    def deploy(self, pathway: Pathway) -> None:
        """Deploy an additional pathway into this session's service."""
        self._pathway_service.deploy(pathway)

    async def stream(
        self,
        message: str,
        *,
        thread_id: str | None = None,
        budget: ExecutionBudget | None = None,
        pathway_override: str | None = None,
    ) -> AsyncIterator[StreamEvent]:
        """Execute one turn, yielding ``StreamEvent`` as tokens arrive.

        Usage::

            async for event in session.stream("Hello"):
                if event.type == "token":
                    print(event.content, end="", flush=True)
                elif event.type == "done":
                    print()

        Parameters
        ----------
        message
            User message.
        thread_id
            Thread id (defaults to session-level default).
        budget
            Execution budget overrides.
        pathway_override
            Use a different pathway for this turn.

        Yields
        ------
        StreamEvent
            Token deltas, tool call notifications, and final result.
        """
        from core.types.stream import StreamEvent as _StreamEvent
        from core.types.stream import TurnComplete

        queue: asyncio.Queue[_StreamEvent | None] = asyncio.Queue()

        def _callback(event: _StreamEvent) -> None:
            queue.put_nowait(event)

        async def _run_turn() -> ExecutionResult:
            result = await self._agent.turn(
                message,
                thread_id or self._default_thread_id,
                pathway_vm=self._vm,
                pathway_service=self._pathway_service,
                budget=budget,
                pathway_override=pathway_override,
                stream_callback=_callback,
            )
            # Signal completion then end-of-stream
            queue.put_nowait(TurnComplete(result=result))
            queue.put_nowait(None)
            return result

        task = asyncio.create_task(_run_turn())

        try:
            while True:
                event = await queue.get()
                if event is None:
                    break
                yield event
        finally:
            if not task.done():
                task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await task

    async def interactive(
        self,
        *,
        prompt: str = "you> ",
        name: str | None = None,
    ) -> None:
        """Run an interactive REPL session.

        In ``CONVERSATIONAL`` mode, streams tokens to stdout as they
        arrive.  In other modes, prints the full response after
        completion.

        Loops on stdin.  Exits on ``quit`` / ``exit`` / ``q`` /
        Ctrl-C / EOF.

        Parameters
        ----------
        prompt
            Input prompt string shown before each user turn.
        name
            Display name for the greeting banner.  Defaults to the
            agent's display name or ID.
        """
        import sys

        display = name or self._agent.definition.display_name
        print(f"\n  {display} ready.  Type 'quit' to exit.\n")

        while True:
            try:
                raw = await asyncio.to_thread(input, prompt)
                user_input = raw.strip()
            except (EOFError, KeyboardInterrupt):
                print("\nGoodbye.")
                break

            if not user_input:
                continue
            if user_input.lower() in ("quit", "exit", "q"):
                print("Goodbye.")
                break

            if self._mode == ExecutionMode.CONVERSATIONAL:
                # Stream tokens to stdout as they arrive
                print("\nassistant> ", end="", flush=True)
                result = None
                async for event in self.stream(user_input):
                    if event.type == "token":
                        sys.stdout.write(event.content)
                        sys.stdout.flush()
                    elif event.type == "tool_start":
                        sys.stdout.write(f"\n  [calling {event.tool_name}...]")
                        sys.stdout.flush()
                    elif event.type == "tool_done":
                        status = "done" if event.success else "failed"
                        sys.stdout.write(f" {status}\n")
                        sys.stdout.flush()
                    elif event.type == "done":
                        result = event.result
                print()

                if result is not None:
                    meta: list[str] = []
                    if hasattr(result, "steps_taken") and result.steps_taken:
                        meta.append(f"{result.steps_taken} steps")
                    if hasattr(result, "tokens_used") and result.tokens_used:
                        meta.append(f"{result.tokens_used:,} tokens")
                    if meta:
                        print(f"  [{' | '.join(meta)}]\n")
            else:
                # Non-streaming: run to completion then print
                result = await self.run(user_input)

                if result.success:
                    print(f"\nassistant> {result.response}\n")
                    meta_parts: list[str] = []
                    if result.steps_taken:
                        meta_parts.append(f"{result.steps_taken} steps")
                    if result.tokens_used:
                        meta_parts.append(f"{result.tokens_used:,} tokens")
                    if meta_parts:
                        print(f"  [{' | '.join(meta_parts)}]\n")
                else:
                    print(f"\n  error: {result.response}\n")


__all__ = [
    "DEFAULT_PATHWAY_ID",
    "ExecutionMode",
    "Session",
    "run",
]
