"""Interactive message processing for AgentRunner.

Extracted from agent.py to reduce file size.
Spec source: vectl step ``shrink-agent-py`` in ``gate-review-debt`` phase.
"""

# @invar:allow shell_result: Interact helpers return Result types
# @invar:allow function_size: process_interact_message - conversation continuation
# @shell_orchestration: Interact helpers delegate to MCP server and state I/O

from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING, Any, AsyncIterator, cast

from pydantic_ai import Agent
from pydantic_ai.exceptions import ModelHTTPError, UserError as PAgentUserError
from pydantic_ai.messages import ToolCallPart
from pydantic_ai.settings import ModelSettings
from pydantic_ai.usage import UsageLimits
from returns.result import Failure, Result, Success

from anima.core.prune import assemble_tool_output_history
from anima.core.schemas import (
    AgentState,
    AnimaError,
    CompletionPolicy,
    InternalError,
    MaxTurnsReachedError,
    TokenUsage,
)
from anima.shell.agent_budget import check_budget_exhaustion, update_state
from anima.shell.mcp_toolset_policy import (
    begin_tool_output_budget_scope,
    end_tool_output_budget_scope,
)
from anima.shell.agent_response import check_auto_complete, record_tool_usage
from anima.shell.state import append_log

if TYPE_CHECKING:
    from anima.core.tracker import BudgetTracker
    from anima.shell.mcp_agent import AgentMCPServer


# @invar:allow shell_result: Internal helper returning tool argument preview string
# @shell_complexity: Extracts first string arg from ToolCallPart, truncates to 72 chars
# @shell_orchestration: Used by CallToolsNode stream processing to emit tool_call events
def _extract_tool_arg_preview(tool_call_part: ToolCallPart) -> str:
    """Extract a preview string from tool call arguments.

    Gets the first string-valued argument, truncated to 72 characters
    with ellipsis if longer. Returns empty string if no string args.

    Args:
        tool_call_part: The ToolCallPart containing tool call information.

    Returns:
        Preview string of tool arguments, or empty string if no string args.
    """
    args = tool_call_part.args
    if args is None:
        return ""

    # Get args as dict to find first string value
    try:
        args_dict = tool_call_part.args_as_dict()
    except Exception:
        # If parsing fails, try raw args
        if isinstance(args, str):
            return args[:72] + "..." if len(args) > 72 else args
        return ""

    if not args_dict:
        return ""

    # Find first string-valued argument
    for value in args_dict.values():
        if isinstance(value, str):
            preview = value
            if len(preview) > 72:
                preview = preview[:72] + "..."
            return preview

    # No string arguments found
    return ""


# @invar:allow shell_result: Internal helper returning message or None
# @shell_orchestration: Async I/O via MCP server queue and asyncio.wait_for
# @shell_complexity: Queue polling with timeout and error handling
async def poll_pending_message(
    mcp_server: AgentMCPServer | None,
    timeout_seconds: float | None,
    inject_followup_hint_fn: Any,
) -> str | None:
    """Poll for pending messages from interact() MCP tool.

    Spec source: vectl step
    ``agent-loop-resilience.fix-interact-loop-premature-exit-100ms-timeout-bug``
    requires interactive mode polling to avoid self-termination on brief
    empty-queue windows between talk messages.

    Args:
        mcp_server: The MCP server instance or None.
        timeout_seconds: Maximum wait for a queued message. ``None`` blocks
            until a message arrives.
        inject_followup_hint_fn: Callable to transform message with hint injection.

    Returns:
        The pending message string, or None if no message pending.
    """
    if mcp_server is None:
        await asyncio.sleep(0)
        return None

    queue = mcp_server.get_pending_message()
    if queue is None:
        await asyncio.sleep(0)
        return None

    try:
        if timeout_seconds is None:
            message = await queue.get()
            return inject_followup_hint_fn(message)
        message = await asyncio.wait_for(queue.get(), timeout=timeout_seconds)
        return inject_followup_hint_fn(message)
    except asyncio.TimeoutError:
        return None


# @invar:allow shell_result: Internal helper preserving original AgentRunner method signature
# @shell_orchestration: Delegates to read_state file I/O
def read_current_state(instance_id: str) -> AgentState | None:
    """Read persisted state for interactive-loop exit checks.

    Spec source: vectl step
    ``agent-loop-resilience.fix-interact-loop-premature-exit-100ms-timeout-bug``
    requires daemon interactive loops to exit only after explicit finish.

    Args:
        instance_id: The agent's unique identifier.

    Returns:
        Current AgentState, or None on read failure.
    """
    from anima.shell.state import read_state

    state_result = read_state(instance_id)
    if isinstance(state_result, Failure):
        return None
    return state_result.unwrap().state


# @invar:allow shell_result: Internal helper returning drained message list
# @shell_orchestration: Drains MCP server queue (I/O boundary)
# @shell_complexity: Queue draining with message classification
def drain_phase1_pending_messages(mcp_server: AgentMCPServer | None) -> list[str]:
    """Drain messages queued during initial task execution (phase 1).

    Spec source: vectl step
    ``agent-loop-resilience.drain-pending-messages-after-phase-1-completes``.

    Args:
        mcp_server: The MCP server instance or None.

    Returns:
        Messages queued before interact loop starts, oldest first.
    """
    if mcp_server is None:
        return []

    drain_pending_messages_fn = getattr(mcp_server, "drain_pending_messages", None)
    if callable(drain_pending_messages_fn):
        drained = drain_pending_messages_fn()
        if isinstance(drained, list):
            return [message for message in drained if isinstance(message, str)]

    queue = mcp_server.get_pending_message()
    if queue is None:
        return []

    drained_fallback: list[str] = []
    while True:
        try:
            drained_fallback.append(queue.get_nowait())
        except asyncio.QueueEmpty:
            break
    return drained_fallback


# @invar:allow shell_result: Internal helper returning bool flag
# @shell_orchestration: Reads MCP server event (I/O boundary)
def consume_stop_turn_request(mcp_server: AgentMCPServer | None) -> bool:
    """Consume a stop-turn control request if present.

    Spec source: plan step ``talk-stop-command`` requires a control event
    checked between PydanticAI nodes so /stop cancels only the current turn.

    Args:
        mcp_server: The MCP server instance or None.

    Returns:
        True if a stop request was consumed.
    """
    if mcp_server is None:
        return False

    stop_turn_event = mcp_server.get_stop_turn_event()
    if not stop_turn_event.is_set():
        return False

    stop_turn_event.clear()
    return True


# @invar:allow function_size: Conversation continuation with message handling
# @shell_complexity: Message type dispatch with approval, continuation, and error handling
async def process_interact_message(
    agent: Agent[Any],
    message: str,
    message_history: list[Any],
    *,
    instance_id: str,
    mcp_server: AgentMCPServer | None,
    model_settings: ModelSettings | None,
    turn_count: int,
    max_turns: int | None,
    token_usage: TokenUsage,
    budget_tracker: BudgetTracker,
    no_tool_streak: int,
    completion_policy: CompletionPolicy,
    in_grace_window: bool,
    check_interrupt_fn: Any,
    queue_followup_hint_fn: Any,
    tool_output_per_call_cap_bytes: int,
    tool_output_per_turn_cap_bytes: int,
) -> tuple[Result[list[Any], AnimaError], int, int, bool]:
    """Process an interact message as a continuation of the conversation.

    Args:
        agent: The PydanticAI Agent instance.
        message: The user message to inject.
        message_history: The conversation history to continue from.
        instance_id: The agent's unique identifier.
        mcp_server: MCP server for busy/response signaling.
        model_settings: Optional model settings with timeout.
        turn_count: Current turn count.
        max_turns: Maximum turns allowed.
        token_usage: Mutable token usage accumulator.
        budget_tracker: Budget tracker instance.
        no_tool_streak: Current no-tool streak.
        completion_policy: Agent's completion policy.
        in_grace_window: Whether in grace window.
        check_interrupt_fn: Callable that returns interrupt result or None.
        queue_followup_hint_fn: Callable to queue followup hint after response.
        tool_output_per_call_cap_bytes: Per-tool-result cap for current-turn
            message assembly.
        tool_output_per_turn_cap_bytes: Total cap for cumulative current-turn
            tool-result content.

    Returns:
        Tuple of (result, new_turn_count, new_no_tool_streak, new_in_grace_window).
        result is Success with updated message history, or Failure with error.
    """
    from pydantic_ai._agent_graph import CallToolsNode, ModelRequestNode
    from pydantic_ai.messages import ModelResponse

    # Log user message for talk history replay
    # Spec source: vectl step talk-ux-bugs-v2.fix-history-replay-user-messages
    # User messages must be logged with message_text and turn_number for replay.
    # Turn number is turn_count + 1 (this interaction's turn).
    # This matches the stream session (start_stream_session) and response log.
    append_log(
        instance_id,
        "user_message",
        {"message_text": message, "turn_number": turn_count + 1},
    )

    # Reset turn tracking for this interaction
    turn_count_local = 0
    no_tool_streak_local = no_tool_streak

    response_text: str | None = None
    stop_requested = False

    async def _iter_stream_text_deltas(stream: object) -> AsyncIterator[str]:
        stream_text = getattr(stream, "stream_text", None)
        if not callable(stream_text):
            return
        stream_iter = cast(Any, stream_text(delta=True, debounce_by=0.05))
        async for delta in stream_iter:
            if isinstance(delta, str) and delta:
                yield delta

    if mcp_server is not None:
        # Spec source: vectl step
        # agent-loop-resilience.return-agent-busy-status-instead-of-hard-timeout
        # marks active turn so concurrent interact calls can return AGENT_BUSY.
        mcp_server.set_busy(True)
        # Initialize streaming session for this turn
        # Spec: talk-streaming-transport.md - prepare stream queue before agent iteration
        await mcp_server.start_stream_session(turn_count + turn_count_local + 1)

    try:
        # Find the boundary of the last model exchange: preserve the most
        # recent assistant response + its tool results (the model needs this
        # context to understand what just happened). Degrade everything older.
        # Find the last ModelResponse — everything from it onwards is "current".
        preserve_idx = len(message_history)  # fallback: degrade all
        for i in range(len(message_history) - 1, -1, -1):
            msg = message_history[i]
            # ModelResponse has 'parts' with TextPart/ToolCallPart
            if hasattr(msg, "parts") and any(
                hasattr(p, "tool_name") and hasattr(p, "args")
                for p in getattr(msg, "parts", [])
            ):
                preserve_idx = i
                break

        history_for_interact = (
            assemble_tool_output_history(
                message_history,
                preserve_after_index=preserve_idx,
                per_call_cap_bytes=tool_output_per_call_cap_bytes,
                per_turn_cap_bytes=tool_output_per_turn_cap_bytes,
            )
            if message_history
            else message_history
        )
        turn_budget_token = begin_tool_output_budget_scope(
            per_call_cap_bytes=tool_output_per_call_cap_bytes,
            per_turn_cap_bytes=tool_output_per_turn_cap_bytes,
        )
        prev_input_tokens = 0
        prev_output_tokens = 0
        try:
            async with agent.iter(
                user_prompt=message,
                message_history=history_for_interact,
                model_settings=model_settings,
                usage_limits=UsageLimits(request_limit=None),
            ) as run:
                async for node in run:
                    if consume_stop_turn_request(mcp_server):
                        stop_requested = True
                        append_log(
                            instance_id,
                            "interact_turn_stopped",
                            {"message_length": len(message)},
                        )
                        break

                    # Check for interrupt
                    killed = check_interrupt_fn()
                    if killed:
                        return (
                            Failure(InternalError("Agent interrupted during interact")),
                            turn_count_local,
                            no_tool_streak_local,
                            in_grace_window,
                        )

                    # Track turns: ModelRequestNode indicates a model turn
                    if isinstance(node, ModelRequestNode):
                        turn_count_local += 1
                        update_state(
                            instance_id, turn_count + turn_count_local, token_usage
                        )

                        stream_ctx = getattr(run, "ctx", None)
                        if mcp_server is not None and stream_ctx is not None:
                            async with node.stream(stream_ctx) as agent_stream:
                                async for delta in _iter_stream_text_deltas(
                                    agent_stream
                                ):
                                    await mcp_server.enqueue_stream_chunk(
                                        turn_number=turn_count + turn_count_local,
                                        chunk=delta,
                                    )

                        from anima.shell.agent_budget import handle_model_request_turn

                        exceeded, in_grace_window, _ = handle_model_request_turn(
                            turn_count + turn_count_local,
                            max_turns,
                            in_grace_window,
                            budget_tracker,
                            instance_id,
                        )
                        if exceeded:
                            return (
                                Failure(
                                    MaxTurnsReachedError(
                                        f"Max turns ({max_turns}) reached during interact"
                                    )
                                ),
                                turn_count_local,
                                no_tool_streak_local,
                                in_grace_window,
                            )

                    # Track tool calls for auto-completion
                    if isinstance(node, CallToolsNode):
                        no_tool_streak_local = record_tool_usage(
                            node, no_tool_streak_local, completion_policy
                        )

                        # Emit tool_call stream events for each ToolCallPart
                        # Spec: data-layer.emit-tool-call-from-agent-loop
                        if mcp_server is not None:
                            model_response = node.model_response
                            for part in model_response.parts:
                                if isinstance(part, ToolCallPart):
                                    tool_name = part.tool_name
                                    arg_preview = _extract_tool_arg_preview(part)
                                    await mcp_server.enqueue_stream_tool_call(
                                        turn_number=turn_count + turn_count_local,
                                        tool_name=tool_name,
                                        tool_arg_preview=arg_preview,
                                    )

                    # Track token usage delta from run.
                    # run.usage() returns CUMULATIVE usage within the iter() call.
                    # We must compute the delta since the last check to avoid
                    # double-counting (F1 bug: adding cumulative on every node).
                    usage = run.usage()
                    if usage:
                        curr_input = usage.input_tokens or 0
                        curr_output = usage.output_tokens or 0
                        delta_input = curr_input - prev_input_tokens
                        delta_output = curr_output - prev_output_tokens
                        prev_input_tokens = curr_input
                        prev_output_tokens = curr_output

                        if delta_input > 0 or delta_output > 0:
                            token_usage.prompt_tokens += delta_input
                            token_usage.completion_tokens += delta_output
                            # last_request_prompt_tokens = this request's actual
                            # input size (for compaction decisions)
                            token_usage.last_request_prompt_tokens = curr_input
                            budget_tracker.record_usage(delta_input + delta_output)
                            update_state(
                                instance_id,
                                turn_count + turn_count_local,
                                token_usage,
                            )

                        # Check budget exhaustion after each token usage
                        # Spec source: spec-compliance-v2.budget-error-type
                        budget_result, in_grace_window = check_budget_exhaustion(
                            budget_tracker, in_grace_window, instance_id
                        )
                        if budget_result is not None:
                            # Emit stream error before returning
                            # Spec: talk-streaming-transport.md - terminal event on failure
                            if mcp_server is not None:
                                await mcp_server.enqueue_stream_error(
                                    turn_number=turn_count + turn_count_local + 1,
                                    error="Budget exceeded",
                                    code="BUDGET_EXCEEDED",
                                )
                            return (
                                Failure(budget_result.failure())
                                if isinstance(budget_result, Failure)
                                else budget_result,  # type: ignore[arg-type]
                                turn_count_local,
                                no_tool_streak_local,
                                in_grace_window,
                            )

                    # Check auto-completion
                    completed = check_auto_complete(
                        completion_policy,
                        no_tool_streak_local,
                        instance_id,
                        token_usage,
                    )
                    if completed:
                        break

                # Extract response text from the last model response
                all_messages = run.all_messages()
                for msg in reversed(all_messages):
                    if isinstance(msg, ModelResponse):
                        for part in msg.parts:
                            if hasattr(part, "part_kind") and part.part_kind == "text":
                                response_text = part.content
                                break
                        if response_text:
                            break

                # If no response text found, provide a default
                if response_text is None:
                    if stop_requested:
                        response_text = "[Current turn cancelled by /stop]"
                    else:
                        response_text = "[Agent completed without text response]"

                # Emit stream done event before response signaling
                # Spec: talk-streaming-transport.md - terminal event with usage snapshot
                if mcp_server is not None:
                    final_chars = len(response_text) if response_text else 0
                    stop_reason = "stopped" if stop_requested else "complete"
                    await mcp_server.enqueue_stream_done(
                        turn_number=turn_count + turn_count_local + 1,
                        stop_reason=stop_reason,
                        final_response_chars=final_chars,
                        prompt_tokens=token_usage.prompt_tokens,
                        completion_tokens=token_usage.completion_tokens,
                        total_tokens=token_usage.total_tokens,
                    )

                # Set response via MCP server
                if mcp_server is not None:
                    mcp_server.set_response(
                        response_text, turn_count + turn_count_local
                    )
                    queue_followup_hint_fn()

                # Log the response text for talk replay and logs --responses
                # Spec source: vectl step spawn-result-persistence.log-response-text
                # Event name must match logs --responses filter and replay_talk_history.
                append_log(
                    instance_id,
                    "response",
                    {
                        "response_text": response_text,
                        "turn_number": turn_count + turn_count_local,
                    },
                )

                # Return updated message history for next iteration
                return (
                    Success(all_messages),
                    turn_count_local,
                    no_tool_streak_local,
                    in_grace_window,
                )
        finally:
            end_tool_output_budget_scope(turn_budget_token)

    except (
        ModelHTTPError,
        PAgentUserError,
        RuntimeError,
        OSError,
        ValueError,
        TypeError,
    ) as e:
        append_log(
            instance_id,
            "interact_error",
            {"error": str(e), "error_type": type(e).__name__},
        )
        # Emit stream error event
        # Spec: talk-streaming-transport.md - terminal event on error
        if mcp_server is not None:
            await mcp_server.enqueue_stream_error(
                turn_number=turn_count + turn_count_local + 1,
                error=str(e),
                code="INTERNAL",
            )
        return (
            Failure(InternalError(f"Interact error: {e}")),
            turn_count_local,
            no_tool_streak_local,
            in_grace_window,
        )
    except Exception as e:
        # Catch-all for unhandled exceptions (e.g., openai.APIConnectionError,
        # openai.APITimeoutError, httpx errors) that are not in the explicit
        # exception list above. Without this, the agent process crashes and
        # the talk session sees "peer closed connection".
        append_log(
            instance_id,
            "interact_error",
            {"error": str(e), "error_type": type(e).__name__},
        )
        # Emit stream error event
        # Spec: talk-streaming-transport.md - terminal event on error
        if mcp_server is not None:
            await mcp_server.enqueue_stream_error(
                turn_number=turn_count + turn_count_local + 1,
                error=str(e),
                code="INTERNAL",
            )
        return (
            Failure(InternalError(f"Interact error: {e}")),
            turn_count_local,
            no_tool_streak_local,
            in_grace_window,
        )
    finally:
        if mcp_server is not None:
            mcp_server.set_busy(False)

    # Should not reach here, but satisfy type checker
    return (
        Failure(InternalError("Unexpected exit from interact processing")),
        turn_count_local,
        no_tool_streak_local,
        in_grace_window,
    )


__all__ = [
    "consume_stop_turn_request",
    "drain_phase1_pending_messages",
    "poll_pending_message",
    "process_interact_message",
    "read_current_state",
]
