"""Agent loop: the core processing engine."""

from __future__ import annotations

import asyncio
import json
import re
from typing import Any, Awaitable, Callable

from loguru import logger

from exoclaw.agent.conversation import Conversation
from exoclaw.agent.tools.protocol import Tool, ToolContext
from exoclaw.agent.tools.registry import ToolRegistry
from exoclaw.bus.events import InboundMessage, OutboundMessage
from exoclaw.bus.protocol import Bus
from exoclaw.providers.protocol import LLMProvider


class AgentLoop:
    """
    The agent loop is the core processing engine.

    It:
    1. Receives messages from the bus
    2. Asks the Conversation for the prompt
    3. Calls the LLM
    4. Executes tool calls
    5. Records the turn and sends the response back
    """

    def __init__(
        self,
        bus: Bus,
        provider: LLMProvider,
        conversation: Conversation,
        model: str | None = None,
        max_iterations: int = 40,
        temperature: float = 0.1,
        max_tokens: int = 4096,
        reasoning_effort: str | None = None,
        tools: list[Tool] | None = None,
        registry: ToolRegistry | None = None,
        # Optional lifecycle callbacks — all default to None (backwards compatible).
        # Plugins pass these in; the loop calls them at the appropriate point.
        on_pre_context: Callable[[str, str, str, str], Awaitable[str]] | None = None,
        on_pre_tool: Callable[[str, dict[str, Any], str], Awaitable[str | None]] | None = None,
        on_post_turn: Callable[[list[dict[str, Any]], str, str, str], Awaitable[None]] | None = None,
        on_max_iterations: Callable[[str, str, str], Awaitable[None]] | None = None,
    ):
        self.bus = bus
        self.provider = provider
        self.model = model or provider.get_default_model()
        self.max_iterations = max_iterations
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.reasoning_effort = reasoning_effort
        self._extra_tools = tools or []
        self._on_pre_context = on_pre_context
        self._on_pre_tool = on_pre_tool
        self._on_post_turn = on_post_turn
        self._on_max_iterations = on_max_iterations

        self.conversation = conversation

        self.tools = registry or ToolRegistry()
        for tool in self._extra_tools:
            self.tools.register(tool)
            if hasattr(tool, "set_bus"):
                tool.set_bus(self.bus)
        self._running = False
        self._active_tasks: dict[str, list[asyncio.Task[None]]] = {}  # session_key -> tasks
        self._processing_lock = asyncio.Lock()
        self._current_ctx: ToolContext | None = None  # set before each _run_agent_loop call

    def _notify_tools_inbound(self, msg: InboundMessage) -> None:
        """Let tools that care about inbound messages update their state."""
        for tool in self.tools._tools.values():
            if hasattr(tool, "on_inbound"):
                tool.on_inbound(msg)

    def _collect_plugin_context(self) -> list[str]:
        """Collect system_context() strings from tools that provide them."""
        ctx = []
        for tool in self.tools._tools.values():
            if hasattr(tool, "system_context"):
                try:
                    result = tool.system_context()
                    if result and isinstance(result, str):
                        ctx.append(result)
                except Exception:
                    logger.exception("Error collecting system_context from tool {}", getattr(tool, "name", "?"))
        return ctx

    @staticmethod
    def _strip_think(text: str | None) -> str | None:
        """Remove <think>…</think> blocks that some models embed in content."""
        if not text:
            return None
        return re.sub(r"<think>[\s\S]*?</think>", "", text).strip() or None

    @staticmethod
    def _tool_hint(tool_calls: list[Any]) -> str:
        """Format tool calls as concise hint, e.g. 'web_search("query")'."""
        def _fmt(tc: Any) -> str:
            args = (tc.arguments[0] if isinstance(tc.arguments, list) else tc.arguments) or {}
            val = next(iter(args.values()), None) if isinstance(args, dict) else None
            if not isinstance(val, str):
                return str(tc.name)
            return f'{tc.name}("{val[:40]}…")' if len(val) > 40 else f'{tc.name}("{val}")'
        return ", ".join(_fmt(tc) for tc in tool_calls)

    async def _run_agent_loop(
        self,
        initial_messages: list[dict[str, Any]],
        on_progress: Callable[..., Awaitable[None]] | None = None,
    ) -> tuple[str | None, list[str], list[dict[str, Any]]]:
        """Run the agent iteration loop. Returns (final_content, tools_used, messages)."""
        messages = initial_messages
        iteration = 0
        final_content = None
        tools_used: list[str] = []

        while iteration < self.max_iterations:
            iteration += 1

            response = await self.provider.chat(
                messages=messages,
                tools=self.tools.get_definitions(),
                model=self.model,
                temperature=self.temperature,
                max_tokens=self.max_tokens,
                reasoning_effort=self.reasoning_effort,
            )

            if response.has_tool_calls:
                if on_progress:
                    thought = self._strip_think(response.content)
                    if thought:
                        await on_progress(thought)
                    await on_progress(self._tool_hint(response.tool_calls), tool_hint=True)

                tool_call_dicts = [
                    {
                        "id": tc.id,
                        "type": "function",
                        "function": {
                            "name": tc.name,
                            "arguments": json.dumps(tc.arguments, ensure_ascii=False)
                        }
                    }
                    for tc in response.tool_calls
                ]
                msg: dict[str, Any] = {"role": "assistant", "content": response.content}
                if tool_call_dicts:
                    msg["tool_calls"] = tool_call_dicts
                if response.reasoning_content is not None:
                    msg["reasoning_content"] = response.reasoning_content
                if response.thinking_blocks:
                    msg["thinking_blocks"] = response.thinking_blocks
                messages = [*messages, msg]

                for tool_call in response.tool_calls:
                    tools_used.append(tool_call.name)
                    args_str = json.dumps(tool_call.arguments, ensure_ascii=False)
                    logger.info("Tool call: {}({})", tool_call.name, args_str[:200])
                    if self._on_pre_tool:
                        sk = self._current_ctx.session_key if self._current_ctx else ""
                        rejection = await self._on_pre_tool(tool_call.name, tool_call.arguments, sk)
                        if rejection:
                            logger.info("on_pre_tool rejected {}: {}", tool_call.name, rejection[:100])
                            result = rejection
                        else:
                            result = await self.tools.execute(tool_call.name, tool_call.arguments, self._current_ctx)
                    else:
                        result = await self.tools.execute(tool_call.name, tool_call.arguments, self._current_ctx)
                    messages = [
                        *messages,
                        {"role": "tool", "tool_call_id": tool_call.id,
                         "name": tool_call.name, "content": result},
                    ]
            else:
                clean = self._strip_think(response.content)
                if response.finish_reason == "error":
                    logger.error("LLM returned error: {}", (clean or "")[:200])
                    final_content = clean or "Sorry, I encountered an error calling the AI model."
                    break
                msg = {"role": "assistant", "content": clean}
                if response.reasoning_content is not None:
                    msg["reasoning_content"] = response.reasoning_content
                if response.thinking_blocks:
                    msg["thinking_blocks"] = response.thinking_blocks
                messages = [*messages, msg]
                final_content = clean
                break

        if final_content is None and iteration >= self.max_iterations:
            logger.warning("Max iterations ({}) reached", self.max_iterations)
            final_content = (
                f"I reached the maximum number of tool call iterations ({self.max_iterations}) "
                "without completing the task. You can try breaking the task into smaller steps."
            )
            if self._on_max_iterations and self._current_ctx:
                ctx = self._current_ctx
                asyncio.ensure_future(self._on_max_iterations(ctx.session_key, ctx.channel, ctx.chat_id))

        return final_content, tools_used, messages

    async def run(self) -> None:
        """Run the agent loop, dispatching messages as tasks to stay responsive to /stop."""
        self._running = True
        logger.info("Agent loop started")

        try:
            while self._running:
                try:
                    msg = await asyncio.wait_for(self.bus.consume_inbound(), timeout=1.0)
                except asyncio.TimeoutError:
                    continue

                if msg.content.strip().lower() == "/stop":
                    await self._handle_stop(msg)
                else:
                    task = asyncio.create_task(self._dispatch(msg))
                    self._active_tasks.setdefault(msg.session_key, []).append(task)

                    def _remove_task(t: asyncio.Task[None], k: str = msg.session_key) -> None:
                        tasks = self._active_tasks.get(k, [])
                        if t in tasks:
                            tasks.remove(t)

                    task.add_done_callback(_remove_task)
        finally:
            all_tasks = [t for ts in self._active_tasks.values() for t in ts if not t.done()]
            for t in all_tasks:
                t.cancel()
            if all_tasks:
                await asyncio.gather(*all_tasks, return_exceptions=True)
            logger.info("Agent loop stopped")

    async def _handle_stop(self, msg: InboundMessage) -> None:
        """Cancel all active tasks and subagents for the session."""
        tasks = self._active_tasks.pop(msg.session_key, [])
        cancelled = sum(1 for t in tasks if not t.done() and t.cancel())
        for t in tasks:
            try:
                await t
            except (asyncio.CancelledError, Exception):
                pass
        sub_cancelled = 0
        for tool in self.tools._tools.values():
            if hasattr(tool, "cancel_by_session"):
                sub_cancelled += await tool.cancel_by_session(msg.session_key)
        total = cancelled + sub_cancelled
        content = f"⏹ Stopped {total} task(s)." if total else "No active task to stop."
        await self.bus.publish_outbound(OutboundMessage(
            channel=msg.channel, chat_id=msg.chat_id, content=content,
        ))

    async def _dispatch(self, msg: InboundMessage) -> None:
        """Process a message under the global lock."""
        async with self._processing_lock:
            try:
                response = await self._process_message(msg)
                if response is not None:
                    await self.bus.publish_outbound(response)
                elif msg.channel == "cli":
                    await self.bus.publish_outbound(OutboundMessage(
                        channel=msg.channel, chat_id=msg.chat_id,
                        content="", metadata=msg.metadata or {},
                    ))
            except asyncio.CancelledError:
                logger.info("Task cancelled for session {}", msg.session_key)
                raise
            except Exception:
                logger.exception("Error processing message for session {}", msg.session_key)
                await self.bus.publish_outbound(OutboundMessage(
                    channel=msg.channel, chat_id=msg.chat_id,
                    content="Sorry, I encountered an error.",
                ))

    def stop(self) -> None:
        """Stop the agent loop."""
        self._running = False
        logger.info("Agent loop stopping")

    async def _process_message(
        self,
        msg: InboundMessage,
        session_key: str | None = None,
        on_progress: Callable[[str], Awaitable[None]] | None = None,
    ) -> OutboundMessage | None:
        """Process a single inbound message and return the response."""
        # System messages: parse origin from chat_id ("channel:chat_id")
        if msg.channel == "system":
            channel, chat_id = (msg.chat_id.split(":", 1) if ":" in msg.chat_id
                                else ("cli", msg.chat_id))
            logger.info("Processing system message from {}", msg.sender_id)
            sid = f"{channel}:{chat_id}"
            plugin_ctx = self._collect_plugin_context()
            initial = await self.conversation.build_prompt(sid, msg.content, channel=channel, chat_id=chat_id, plugin_context=plugin_ctx or None)
            final_content, _, all_msgs = await self._run_agent_loop(initial)
            await self.conversation.record(sid, all_msgs[len(initial) - 1:])
            return OutboundMessage(channel=channel, chat_id=chat_id,
                                   content=final_content or "Background task completed.")

        preview = msg.content[:80] + "..." if len(msg.content) > 80 else msg.content
        logger.info("Processing message from {}:{}: {}", msg.channel, msg.sender_id, preview)

        sid = session_key or msg.session_key

        # Slash commands
        cmd = msg.content.strip().lower()
        if cmd == "/new":
            success = await self.conversation.clear(sid)
            content = "New session started." if success else (
                "Memory archival failed, session not cleared. Please try again."
            )
            return OutboundMessage(channel=msg.channel, chat_id=msg.chat_id, content=content)

        if cmd == "/help":
            return OutboundMessage(
                channel=msg.channel, chat_id=msg.chat_id,
                content="🦀 exoclaw commands:\n/new — Start a new conversation\n/stop — Stop the current task\n/help — Show available commands",
            )

        self._notify_tools_inbound(msg)
        self._current_ctx = ToolContext(session_key=sid, channel=msg.channel, chat_id=msg.chat_id)

        plugin_ctx = self._collect_plugin_context()
        if self._on_pre_context:
            extra = await self._on_pre_context(msg.content, sid, msg.channel, msg.chat_id)
            if extra:
                plugin_ctx.append(extra)

        initial = await self.conversation.build_prompt(
            sid, msg.content,
            channel=msg.channel, chat_id=msg.chat_id,
            media=msg.media if msg.media else None,
            plugin_context=plugin_ctx or None,
        )

        async def _bus_progress(content: str, *, tool_hint: bool = False) -> None:
            meta = dict(msg.metadata or {})
            meta["_progress"] = True
            meta["_tool_hint"] = tool_hint
            await self.bus.publish_outbound(OutboundMessage(
                channel=msg.channel, chat_id=msg.chat_id, content=content, metadata=meta,
            ))

        final_content, _, all_msgs = await self._run_agent_loop(
            initial, on_progress=on_progress or _bus_progress,
        )

        if final_content is None:
            final_content = "I've completed processing but have no response to give."

        # new_msgs = user message + everything the loop added
        new_msgs = all_msgs[len(initial) - 1:]
        await self.conversation.record(sid, new_msgs)
        if self._on_post_turn and new_msgs:
            asyncio.ensure_future(self._on_post_turn(new_msgs, sid, msg.channel, msg.chat_id))

        if any(getattr(t, "sent_in_turn", False) for t in self.tools._tools.values()):
            return None

        preview = final_content[:120] + "..." if len(final_content) > 120 else final_content
        logger.info("Response to {}:{}: {}", msg.channel, msg.sender_id, preview)
        return OutboundMessage(
            channel=msg.channel, chat_id=msg.chat_id, content=final_content,
            metadata=msg.metadata or {},
        )

    async def process_direct(
        self,
        content: str,
        session_key: str = "cli:direct",
        channel: str = "cli",
        chat_id: str = "direct",
        on_progress: Callable[[str], Awaitable[None]] | None = None,
    ) -> str:
        """Process a message directly (for CLI or cron usage)."""
        msg = InboundMessage(channel=channel, sender_id="user", chat_id=chat_id, content=content)
        response = await self._process_message(msg, session_key=session_key, on_progress=on_progress)
        return response.content if response else ""
