# Copyright (c) Microsoft. All rights reserved.

"""Agent runner orchestration for the harness console.

This module provides the HarnessAgentRunner class, which orchestrates agent
invocations with observer lifecycle management. It handles:
- User input dispatch
- Agent streaming with observer notifications
- Follow-up action collection
- Streaming state management
"""

from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from agent_framework import Agent, AgentSession, MessageInjectionMiddleware
    from agent_framework import Message as FrameworkMessage

    from .app_state import FollowUpAction
    from .observers.base import ConsoleObserver
    from .state_driver import IUXStateDriver


class HarnessAgentRunner:
    """Orchestrates agent invocations driven by user-input events from the UI.

    The component invokes the runner's input handlers (run_turn) directly;
    the runner mutates UI state through the supplied IUXStateDriver.

    When the underlying agent has a ``MessageInjectionMiddleware`` wired in
    (as ``create_harness_agent`` does by default), the runner supports message
    injection: input submitted while a turn is streaming is enqueued via
    ``on_streaming_input`` and drained into the ongoing run by the middleware.
    """

    def __init__(
        self,
        agent: Agent,
        observers: list[ConsoleObserver],
        state_driver: IUXStateDriver,
        *,
        max_context_window_tokens: int | None = None,
        max_output_tokens: int | None = None,
    ) -> None:
        """Initialize the agent runner.

        Args:
            agent: The agent to orchestrate.
            observers: List of console observers for lifecycle events.
            state_driver: The UI state driver for observer updates.
            max_context_window_tokens: Optional max context window size for usage display.
            max_output_tokens: Optional max output tokens for usage display.
        """
        self._agent = agent
        self._observers = observers
        self._ux = state_driver
        self._max_context_window_tokens = max_context_window_tokens
        self._max_output_tokens = max_output_tokens
        self._input_gate = asyncio.Semaphore(1)  # Single turn at a time

        # Resolve the message-injection middleware (if any) so streaming-time
        # input can be enqueued into the ongoing run. Absent => injection no-ops.
        from agent_framework import MessageInjectionMiddleware

        self._message_injector: MessageInjectionMiddleware | None = next(
            (m for m in (agent.middleware or []) if isinstance(m, MessageInjectionMiddleware)),
            None,
        )
        # Snapshot of pending injected messages, used to detect consumption
        # during streaming. Safe as instance state because _input_gate
        # serialises turns.
        self._last_pending_messages: list[FrameworkMessage] = []

    async def run_turn(
        self,
        user_input: str,
        session: AgentSession | None = None,
    ) -> None:
        """Run a single agent turn with the given user input.

        Echoes the input, then delegates to the agent loop.

        Args:
            user_input: The user's input text.
            session: Optional agent session for conversation history.
        """
        async with self._input_gate:
            self._ux.write_user_input_echo(user_input)

            from agent_framework import Message

            messages = [Message(role="user", contents=[user_input])]
            await self._run_agent_loop(messages, session)

    async def start_agent_turn(
        self,
        messages: list,
        session: AgentSession | None = None,
    ) -> None:
        """Resume the agent loop with pre-built messages (from follow-up responses).

        Called by the app after the user finishes answering follow-up questions.
        If messages is empty, just completes the turn.

        Args:
            messages: List of Message objects to send to the agent.
            session: Optional agent session.
        """
        async with self._input_gate:
            if not messages:
                self._complete_turn()
                return
            await self._run_agent_loop(messages, session)

    def on_streaming_input(
        self,
        text: str,
        session: AgentSession | None = None,
    ) -> None:
        """Handle user input submitted while an agent turn is streaming.

        The text is enqueued via the ``MessageInjectionMiddleware`` so the agent
        can pick it up on its next opportunity within the ongoing run. No-op if
        the agent has no injection middleware or there is no active session.

        Args:
            text: The user's input text.
            session: The active agent session.
        """
        if self._message_injector is None or session is None:
            return

        text = text.strip()
        if not text:
            return

        from agent_framework import Message

        self._message_injector.enqueue_messages(session, Message(role="user", contents=[text]))
        pending = self._message_injector.get_pending_messages(session)
        self._ux.set_queued_messages([m.text for m in pending])

    def _sync_queued_message_display(self, session: AgentSession | None) -> None:
        """Sync the queued-items display with the injector's pending messages.

        Messages that have been consumed (drained by the middleware) since the
        last sync are echoed to the output area as regular user-input entries.
        No-op if there is no injection middleware or active session.

        Args:
            session: The active agent session.
        """
        if self._message_injector is None or session is None:
            return

        pending = self._message_injector.get_pending_messages(session)

        # The injection middleware drains the whole queue at once, so a message
        # is consumed when it is no longer present in the pending list. Compare
        # by object identity (snapshots share the same Message objects until the
        # queue is cleared) so consumed messages are echoed correctly even if a
        # drain is followed by a new enqueue before the next sync.
        current_ids = {id(m) for m in pending}
        for msg in self._last_pending_messages:
            if id(msg) not in current_ids:
                self._ux.write_user_input_echo(msg.text or "")

        self._last_pending_messages = pending
        self._ux.set_queued_messages([m.text for m in pending])

    async def _run_agent_loop(
        self,
        messages: list,
        session: AgentSession | None,
    ) -> None:
        """Run the agent loop, re-invoking as needed for follow-up messages.

        Loops while there are messages to send. After each stream:
        - Collects follow-up actions from observers
        - If questions exist → queue them and return (UI will collect answers)
        - If only direct messages → loop with those messages
        - If nothing → complete the turn

        Args:
            messages: Initial messages to send.
            session: Optional agent session.
        """
        next_messages = messages

        # Seed the pending-message snapshot so consumed injected messages can be
        # detected and echoed during streaming.
        self._last_pending_messages = (
            self._message_injector.get_pending_messages(session)
            if self._message_injector is not None and session is not None
            else []
        )

        while next_messages:
            # Configure run options
            options = self._configure_run_options(session)

            # Begin streaming
            self._ux.begin_streaming()
            self._ux.begin_streaming_output()
            self._ux.set_show_spinner(True)

            try:
                await self._stream_response_messages(next_messages, session, options)
            except Exception as ex:
                self._ux.append_info_line(
                    f"❌ Stream error: {ex.__class__.__name__}:\n{ex}",
                    color="red",
                )

            # Final sync after streaming (echo any messages consumed on the last update).
            self._sync_queued_message_display(session)

            # Stop spinner and end streaming output
            self._ux.set_show_spinner(False)

            # Collect follow-up actions from observers
            follow_up_actions = await self._collect_follow_up_actions(session)

            # Separate direct messages from questions
            has_follow_ups = len(follow_up_actions) > 0

            # Write no-text warning if applicable
            await self._ux.write_no_text_warning(has_follow_ups)

            # Enqueue all follow-up actions
            for action in follow_up_actions:
                self._ux.enqueue_follow_up_action(action)

            # Check if there are pending questions (UI needs user input)
            if self._ux.has_pending_questions():
                # Pause — the UI will collect answers and call start_agent_turn
                return

            # No questions — drain any accumulated direct messages and loop
            drained = self._ux.take_follow_up_responses()
            next_messages = drained if drained else None

        self._complete_turn()

    def _complete_turn(self) -> None:
        """Complete the current turn (end streaming)."""
        self._ux.end_streaming()

    def _configure_run_options(
        self,
        session: AgentSession | None,
    ) -> dict:
        """Configure run options via observers.

        Each observer can modify the options dict to influence agent behavior.

        Args:
            session: Optional agent session.

        Returns:
            Options dict for agent.run().
        """
        options = {}
        for observer in self._observers:
            observer.configure_run_options(options, self._agent, session)
        return options

    async def _stream_response(
        self,
        user_input: str,
        session: AgentSession | None,
        options: dict,
    ) -> None:
        """Stream agent response from a text input and dispatch to observers.

        Args:
            user_input: The user's input text.
            session: Optional agent session.
            options: Run options configured by observers.
        """
        # Stream response using agent.run(stream=True)
        stream = self._agent.run(
            user_input,
            stream=True,
            session=session,
            options=options,
        )

        # Process each update chunk
        async for update in stream:
            await self._dispatch_update(update, session)

        # Extract usage from the final response
        self._extract_usage(stream)

    async def _stream_response_messages(
        self,
        messages: list,
        session: AgentSession | None,
        options: dict,
    ) -> None:
        """Stream agent response from Message objects and dispatch to observers.

        Args:
            messages: List of Message objects to send.
            session: Optional agent session.
            options: Run options configured by observers.
        """
        stream = self._agent.run(
            messages,
            stream=True,
            session=session,
            options=options,
        )

        async for update in stream:
            await self._dispatch_update(update, session)

        self._extract_usage(stream)

    def _extract_usage(self, stream) -> None:
        """Extract token usage from a completed stream."""
        try:
            get_final = getattr(stream, "get_final_response", None)
            if not get_final:
                return

            import inspect

            if inspect.iscoroutinefunction(get_final):
                return

            final_response = get_final()
            if final_response is None:
                return

            usage = getattr(final_response, "usage_details", None)
            if not isinstance(usage, dict):
                return

            input_tokens = usage.get("input_token_count", 0) or 0
            output_tokens = usage.get("output_token_count", 0) or 0
            if input_tokens or output_tokens:
                self._ux.set_usage_text(self._format_usage(input_tokens, output_tokens))
        except (AttributeError, TypeError):
            pass

    async def _dispatch_update(
        self,
        update,  # AgentResponseUpdate
        session: AgentSession | None,
    ) -> None:
        """Dispatch a single update to all observers.

        Calls observer lifecycle methods in order:
        1. on_response_update (once per update)
        2. on_content (for each content item)
        3. on_text (if text is present)

        Args:
            update: The agent response update.
            session: Optional agent session.
        """
        # on_response_update
        for observer in self._observers:
            await observer.on_response_update(self._ux, update, self._agent, session)

        # on_content for each content item
        if hasattr(update, "contents") and update.contents:
            for content in update.contents:
                for observer in self._observers:
                    await observer.on_content(self._ux, content, self._agent, session)

        # on_text for text chunks
        if hasattr(update, "text") and update.text:
            for observer in self._observers:
                await observer.on_text(self._ux, update.text, self._agent, session)

        # Echo any injected messages consumed by the agent on this update.
        self._sync_queued_message_display(session)

    async def _collect_follow_up_actions(
        self,
        session: AgentSession | None,
    ) -> list[FollowUpAction]:
        """Collect follow-up actions from all observers.

        Called after streaming completes to gather any follow-up questions
        or messages from observers.

        Args:
            session: Optional agent session.

        Returns:
            List of follow-up actions from all observers.
        """
        actions: list[FollowUpAction] = []
        for observer in self._observers:
            observer_actions = await observer.on_stream_complete(self._ux, self._agent, session)
            if observer_actions:
                actions.extend(observer_actions)
        return actions

    def _format_usage(self, input_tokens: int, output_tokens: int) -> str:
        """Format token counts matching C# harness style: 📊 Tokens — input: X | output: Y | total: Z."""
        total_tokens = input_tokens + output_tokens

        input_budget = None
        if self._max_context_window_tokens and self._max_output_tokens:
            input_budget = self._max_context_window_tokens - self._max_output_tokens

        return (
            f"📊 Tokens — input: {self._format_token_count(input_tokens, input_budget)}"
            f" | output: {self._format_token_count(output_tokens, self._max_output_tokens)}"
            f" | total: {self._format_token_count(total_tokens, self._max_context_window_tokens)}"
        )

    @staticmethod
    def _format_token_count(count: int, budget: int | None) -> str:
        """Format a token count, optionally showing budget percentage."""
        if budget and budget > 0:
            pct = count / budget * 100
            return f"{count:,}/{budget:,} ({pct:.1f}%)"
        return f"{count:,}"