"""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.observability.trace import ExecutionTrace
    from core.pathways.checkpoint import DAGCheckpoint
    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,
    inputs: dict[str, object] | None = None,
) -> 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``.
    inputs
        Structured inputs to pass to the pathway.  Merged with the
        standard fields (goal, conversation_history, budget), with
        caller-supplied keys taking precedence.

    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, inputs=inputs)


# =============================================================================
# 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.

    Observability is always on — every execution is traced
    automatically.  Access the most recent trace via
    ``session.last_trace``, or look up by execution ID with
    ``session.get_trace(eid)``.

    Optional integrations (``otel``, ``metrics``) attach additional
    listeners for export to external systems.
    """

    def __init__(
        self,
        agent_or_definition: Agent | AgentDefinition,
        *,
        pathway: Pathway | None = None,
        state_store: StateStoreLike | None = None,
        mode: ExecutionMode = ExecutionMode.CONVERSATIONAL,
        otel: bool = False,
        metrics: bool = False,
    ) -> None:
        self._mode = mode
        self._otel_exporter: object | None = None
        self._metrics_collector: object | None = None
        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)

        # --- Optional observability integrations ---
        if otel:
            from core.observability.otel import OTelSpanExporter

            self._otel_exporter = OTelSpanExporter()
            self._vm.add_listener(self._otel_exporter)

        if metrics:
            from core.observability.metrics import MetricsCollector

            self._metrics_collector = MetricsCollector()
            self._vm.add_listener(self._metrics_collector)

        # --- 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

    @property
    def last_trace(self) -> ExecutionTrace | None:
        """The execution trace from the most recent run.

        Tracing is always on — no flag needed.
        """
        return self._vm.last_trace

    def get_trace(self, execution_id: str) -> ExecutionTrace | None:
        """Look up a completed trace by *execution_id*.

        Parameters
        ----------
        execution_id
            The execution ID to look up.

        Returns
        -------
        ExecutionTrace | None
            The trace, or ``None`` if not found.
        """
        return self._vm.get_trace(execution_id)

    def save_trace(
        self,
        path: str,
        *,
        pathway: Pathway | None = None,
        title: str | None = None,
    ) -> str:
        """Save the last execution trace as a self-contained HTML file.

        Parameters
        ----------
        path
            File path to write (e.g. ``"trace.html"``).
        pathway
            The pathway to use for the DAG visualization.  When omitted,
            the default pathway is used.
        title
            Page title.

        Returns
        -------
        str
            The resolved path of the written file.

        Raises
        ------
        ValueError
            If no trace is available.
        """
        trace = self.last_trace
        if trace is None:
            raise ValueError("No trace available — no execution has completed yet.")

        pw = pathway or self._pathway_service.load(trace.pathway_id or DEFAULT_PATHWAY_ID)
        if pw is None:
            raise ValueError(f"Pathway {trace.pathway_id!r} not found in PathwayService")

        from core.observability.html_viewer import save_trace_html

        return str(save_trace_html(pw, trace, path, title=title))

    async def run(
        self,
        message: str,
        *,
        thread_id: str | None = None,
        budget: ExecutionBudget | None = None,
        pathway_override: str | None = None,
        inputs: dict[str, object] | 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``).
        inputs
            Structured inputs to pass to the pathway.  Merged with the
            standard fields (goal, conversation_history, budget), with
            caller-supplied keys taking precedence.

        Returns
        -------
        ExecutionResult
        """
        effective_thread_id = thread_id or self._default_thread_id

        # --- Set logging context for this turn ---
        from core.observability.logging import (
            reset_log_agent_id,
            reset_log_thread_id,
            set_log_agent_id,
            set_log_thread_id,
        )

        agent_token = set_log_agent_id(self._agent.definition.id)
        thread_token = set_log_thread_id(effective_thread_id)

        try:
            result = await self._agent.turn(
                message,
                effective_thread_id,
                pathway_vm=self._vm,
                pathway_service=self._pathway_service,
                budget=budget,
                pathway_override=pathway_override,
                inputs=inputs,
            )
        finally:
            # Persist trace to state store for later retrieval
            trace = self.last_trace
            if trace is not None:
                with contextlib.suppress(Exception):
                    await self._agent._state_store.save_trace(trace)

            # --- Reset logging context ---
            reset_log_agent_id(agent_token)
            reset_log_thread_id(thread_token)

        return result

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

    async def resume_from_checkpoint(
        self,
        execution_id: str,
    ) -> ExecutionResult:
        """Resume a failed pathway execution from its last checkpoint.

        Loads the ``DAGCheckpoint`` from the session's state store, resolves
        the pathway, and re-executes from the point of failure — skipping
        already-completed nodes.

        Parameters
        ----------
        execution_id
            The ``execution_id`` of the failed run.

        Returns
        -------
        ExecutionResult

        Raises
        ------
        ValueError
            If no checkpoint exists for the given *execution_id*.
        """
        from core.pathways.checkpoint import DAGCheckpoint

        store = self._agent._state_store
        checkpoint = await store.load_checkpoint(execution_id)
        if checkpoint is None or not isinstance(checkpoint, DAGCheckpoint):
            raise ValueError(f"No checkpoint found for execution_id={execution_id!r}")

        pathway = self._pathway_service.load(checkpoint.pathway_id)
        if pathway is None:
            raise ValueError(
                f"Pathway {checkpoint.pathway_id!r} not deployed "
                f"(needed to resume execution {execution_id!r})"
            )

        record = await self._vm.resume(
            pathway,
            checkpoint,
            checkpoint_store=store,
        )

        # Build an ExecutionResult consistent with a normal run
        from core.types.execution import (
            ExecutionResult as _ExecutionResult,
        )

        # Pull tokens from the trace
        trace = self._vm.get_trace(execution_id)
        tokens = trace.total_tokens if trace else 0

        return _ExecutionResult(
            success=record.status.value == "completed",
            response=record.outputs.get("response", ""),
            pathway_record=record,
            tokens_used=tokens,
        )

    async def list_checkpoints(self) -> list[DAGCheckpoint]:
        """List all persisted DAG checkpoints in the session's state store.

        Returns
        -------
        list[DAGCheckpoint]
            Checkpoints sorted by timestamp (most recent first).
        """
        from core.pathways.checkpoint import DAGCheckpoint

        store = self._agent._state_store
        store_root = getattr(store, "_root", None)
        if store_root is None:
            return []

        from pathlib import Path

        checkpoints_dir = Path(store_root) / "checkpoints"
        if not checkpoints_dir.exists():
            return []

        results: list[DAGCheckpoint] = []
        for path in checkpoints_dir.glob("*.json"):
            try:
                text = path.read_text()
                results.append(DAGCheckpoint.model_validate_json(text))
            except Exception:
                continue

        results.sort(key=lambda c: c.timestamp, reverse=True)
        return results

    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")


async def replay(
    checkpoint_id: str,
    *,
    agent_or_definition: Agent | AgentDefinition | None = None,
    state_store: StateStoreLike | None = None,
) -> ExecutionResult:
    """Replay a pathway execution from a checkpoint.

    Convenience wrapper around ``Session.resume_from_checkpoint()``
    for one-shot replay without manually creating a Session.

    Parameters
    ----------
    checkpoint_id
        The ``execution_id`` of the checkpointed run to replay.
    agent_or_definition
        An agent or definition. When omitted, a minimal agent is
        created for replaying the checkpoint's pathway.
    state_store
        Custom state store. Defaults to ``FileStateStore()``.

    Returns
    -------
    ExecutionResult
    """
    store = state_store or FileStateStore()

    # Load checkpoint to discover the pathway
    from core.pathways.checkpoint import DAGCheckpoint

    checkpoint = await store.load_checkpoint(checkpoint_id)
    if checkpoint is None or not isinstance(checkpoint, DAGCheckpoint):
        raise ValueError(f"No checkpoint found for execution_id={checkpoint_id!r}")

    if agent_or_definition is None:
        # Create a minimal agent just for replay
        from core.builders.agent import AgentBuilder

        agent_def = AgentBuilder("replay-agent").build()
        session = Session(agent_def, state_store=store)
    else:
        session = Session(agent_or_definition, state_store=store)

    return await session.resume_from_checkpoint(checkpoint_id)


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