"""Agent loop — core LLM interaction cycle with tool execution."""

from __future__ import annotations

import asyncio
import time
from typing import Any, Callable

from pi_agents.ai.types import (
    AssistantMessage,
    Context,
    Cost,
    Message,
    Model,
    SimpleStreamOptions,
    TextContent,
    ToolCall,
    ToolResultMessage,
    Usage,
    UserMessage,
)
from pi_agents.ai.utils.event_stream import AssistantMessageEventStream, EventStream

from pi_agents.agent.types import (
    AfterToolCallContext,
    AfterToolCallResult,
    AgentEvent,
    AgentMessage,
    AgentState,
    AgentTool,
    AgentToolResult,
    BeforeToolCallContext,
    BeforeToolCallResult,
)


async def run_agent_loop(
    prompts: list[AgentMessage],
    context: dict[str, Any],
    config: dict[str, Any],
    emit: Callable[[AgentEvent], Any],
    signal: Any = None,
    stream_fn: Callable[..., Any] | None = None,
) -> list[AgentMessage]:
    """Run agent loop emitting events to an async sink function.

    Core algorithm:
    1. Add prompts to context
    2. Build LLM context (system prompt + converted messages + tools)
    3. Call LLM (stream)
    4. If LLM returns tool calls → execute tools → get results → goto 3
    5. If no tool calls → check steering/follow-up queues → done
    """
    model: Model = config.get("model")
    convert_to_llm = config.get("convert_to_llm", lambda msgs: msgs)
    transform_context = config.get("transform_context")
    get_api_key = config.get("get_api_key")
    get_steering_messages = config.get("get_steering_messages")
    get_follow_up_messages = config.get("get_follow_up_messages")
    tool_execution = config.get("tool_execution", "parallel")
    before_tool_call = config.get("before_tool_call")
    after_tool_call = config.get("after_tool_call")
    tools: list[AgentTool] = context.get("tools", [])

    system_prompt = context.get("system_prompt", "")
    messages: list[AgentMessage] = list(context.get("messages", []))

    if stream_fn is None:
        from pi_agents.ai.stream import stream_simple
        stream_fn = stream_simple

    all_messages = list(prompts)

    # Emit agent_start
    await _maybe_await(emit({"type": "agent_start"}))

    # Add prompts to messages
    for prompt in prompts:
        messages.append(prompt)
        await _maybe_await(emit({"type": "message_start", "message": prompt}))
        await _maybe_await(emit({"type": "message_end", "message": prompt}))

    # Outer loop: handles follow-up messages
    while True:
        # Inner loop: handles tool calls and steering
        while True:
            await _maybe_await(emit({"type": "turn_start"}))

            # Build LLM context
            llm_messages = convert_to_llm(messages)

            if transform_context:
                llm_messages = await transform_context(llm_messages, signal)

            llm_context = Context(
                system_prompt=system_prompt,
                messages=llm_messages,
                tools=[AgentTool(name=t.name, description=t.description, parameters=t.parameters) for t in tools] if tools else None,
            )

            # Build stream options
            stream_options = SimpleStreamOptions()
            thinking_level = config.get("thinking_level", "off")
            if thinking_level and thinking_level != "off":
                stream_options.reasoning = thinking_level

            if get_api_key:
                stream_options.api_key = get_api_key(model.provider)

            # Stream LLM response
            event_stream = stream_fn(model, llm_context, stream_options)

            assistant_msg = None
            async for event in event_stream:
                event_type = event.get("type", "")
                if event_type == "start":
                    partial = event.get("partial")
                    await _maybe_await(emit({"type": "message_start", "message": partial}))
                elif event_type in ("text_delta", "thinking_delta", "toolcall_delta"):
                    partial = event.get("partial")
                    await _maybe_await(emit({"type": "message_update", "message": partial, "assistantMessageEvent": event}))
                elif event_type == "done":
                    assistant_msg = event.get("message")
                    await _maybe_await(emit({"type": "message_end", "message": assistant_msg}))
                elif event_type == "error":
                    assistant_msg = event.get("error")
                    await _maybe_await(emit({"type": "message_end", "message": assistant_msg}))

            if assistant_msg is None:
                break

            messages.append(assistant_msg)
            all_messages.append(assistant_msg)

            # Check for errors
            stop_reason = assistant_msg.stop_reason if hasattr(assistant_msg, "stop_reason") else (assistant_msg.get("stopReason") if isinstance(assistant_msg, dict) else "stop")
            if stop_reason in ("error", "aborted"):
                await _maybe_await(emit({"type": "turn_end", "message": assistant_msg, "toolResults": []}))
                await _maybe_await(emit({"type": "agent_end", "messages": all_messages}))
                return all_messages

            # Extract tool calls
            content = assistant_msg.content if hasattr(assistant_msg, "content") else (assistant_msg.get("content", []) if isinstance(assistant_msg, dict) else [])
            tool_calls = [c for c in content if (hasattr(c, "type") and c.type == "toolCall") or (isinstance(c, dict) and c.get("type") == "toolCall")]

            if not tool_calls:
                await _maybe_await(emit({"type": "turn_end", "message": assistant_msg, "toolResults": []}))
                break

            # Execute tool calls
            tool_results = []
            if tool_execution == "parallel":
                tasks = []
                for tc in tool_calls:
                    tasks.append(_execute_tool(tc, tools, assistant_msg, messages, config, emit, signal))
                results = await asyncio.gather(*tasks)
                tool_results.extend(results)
            else:
                for tc in tool_calls:
                    result = await _execute_tool(tc, tools, assistant_msg, messages, config, emit, signal)
                    tool_results.append(result)

            # Add tool results to messages
            for tr in tool_results:
                messages.append(tr)
                all_messages.append(tr)

            await _maybe_await(emit({"type": "turn_end", "message": assistant_msg, "toolResults": tool_results}))

            # Check for steering messages
            if get_steering_messages:
                steering = await get_steering_messages()
                if steering:
                    for s in steering:
                        messages.append(s)
                        all_messages.append(s)
                        await _maybe_await(emit({"type": "message_start", "message": s}))
                        await _maybe_await(emit({"type": "message_end", "message": s}))

        # Check for follow-up messages
        if get_follow_up_messages:
            follow_ups = await get_follow_up_messages()
            if follow_ups:
                for fu in follow_ups:
                    messages.append(fu)
                    all_messages.append(fu)
                    await _maybe_await(emit({"type": "message_start", "message": fu}))
                    await _maybe_await(emit({"type": "message_end", "message": fu}))
                continue

        break

    await _maybe_await(emit({"type": "agent_end", "messages": all_messages}))
    return all_messages


async def run_agent_loop_continue(
    context: dict[str, Any],
    config: dict[str, Any],
    emit: Callable[[AgentEvent], Any],
    signal: Any = None,
    stream_fn: Callable[..., Any] | None = None,
) -> list[AgentMessage]:
    """Continue agent loop from existing context without adding new messages."""
    return await run_agent_loop([], context, config, emit, signal, stream_fn)


async def _execute_tool(
    tool_call: Any,
    tools: list[AgentTool],
    assistant_msg: Any,
    messages: list[Any],
    config: dict[str, Any],
    emit: Callable[[AgentEvent], Any],
    signal: Any,
) -> ToolResultMessage:
    """Execute a single tool call and return a ToolResultMessage."""
    tc_id = tool_call.id if hasattr(tool_call, "id") else tool_call.get("id", "")
    tc_name = tool_call.name if hasattr(tool_call, "name") else tool_call.get("name", "")
    tc_args = tool_call.arguments if hasattr(tool_call, "arguments") else tool_call.get("arguments", {})

    before_tool_call = config.get("before_tool_call")
    after_tool_call = config.get("after_tool_call")

    await _maybe_await(emit({
        "type": "tool_execution_start",
        "toolCallId": tc_id,
        "toolName": tc_name,
        "args": tc_args,
    }))

    # Find matching tool
    matching_tool = next((t for t in tools if t.name == tc_name), None)

    is_error = False
    result_content: list[TextContent] = []
    details = None

    if not matching_tool:
        is_error = True
        result_content = [TextContent(type="text", text=f'Tool "{tc_name}" not found')]
    else:
        # Before hook
        if before_tool_call:
            hook_result = before_tool_call(
                BeforeToolCallContext(assistant_message=assistant_msg, tool_call=tool_call, args=tc_args, context=None)
            )
            if hook_result and hook_result.block:
                is_error = True
                result_content = [TextContent(type="text", text=hook_result.reason or "Tool call blocked")]
                await _maybe_await(emit({
                    "type": "tool_execution_end",
                    "toolCallId": tc_id,
                    "toolName": tc_name,
                    "result": AgentToolResult(content=result_content),
                    "isError": True,
                }))
                return ToolResultMessage(
                    role="toolResult",
                    tool_call_id=tc_id,
                    tool_name=tc_name,
                    content=result_content,
                    is_error=True,
                    timestamp=int(time.time() * 1000),
                )

        try:
            tool_result = await matching_tool.execute(tc_id, tc_args, signal)
            result_content = tool_result.content
            details = tool_result.details
        except Exception as e:
            is_error = True
            result_content = [TextContent(type="text", text=str(e))]

        # After hook
        if after_tool_call:
            hook_result = after_tool_call(
                AfterToolCallContext(
                    assistant_message=assistant_msg,
                    tool_call=tool_call,
                    args=tc_args,
                    result=AgentToolResult(content=result_content, details=details),
                    is_error=is_error,
                    context=None,
                )
            )
            if hook_result:
                if hook_result.content is not None:
                    result_content = hook_result.content
                if hook_result.details is not None:
                    details = hook_result.details
                if hook_result.is_error is not None:
                    is_error = hook_result.is_error

    await _maybe_await(emit({
        "type": "tool_execution_end",
        "toolCallId": tc_id,
        "toolName": tc_name,
        "result": AgentToolResult(content=result_content, details=details),
        "isError": is_error,
    }))

    return ToolResultMessage(
        role="toolResult",
        tool_call_id=tc_id,
        tool_name=tc_name,
        content=result_content,
        details=details,
        is_error=is_error,
        timestamp=int(time.time() * 1000),
    )


async def _maybe_await(result: Any) -> Any:
    """Await result if it's a coroutine."""
    if asyncio.iscoroutine(result):
        return await result
    return result
