"""Agent loop: the core message-flow orchestrator.

This module is the thin coordination layer that wires together:
- Message queue decisions (message_queue.py)
- LLM interaction (agent_runner.py)
- Single-turn processing (turn_processor.py)
- Session lifecycle (session_lifecycle.py)
"""

from __future__ import annotations

import asyncio
import weakref
from pathlib import Path
from typing import TYPE_CHECKING, Any, Awaitable, Callable

from loguru import logger

from asedclaw.agent.context import ContextBuilder
from asedclaw.agent.context_engine import ContextEngine
from asedclaw.agent.message_queue import (
    FollowupQueueManager,
    QueueDropPolicy,
    QueueMode,
)
from asedclaw.agent.session_lifecycle import (
    handle_new_session,
    maybe_consolidate,
    maybe_memory_flush,
    maybe_sync_index,
)
from asedclaw.agent.subagent import SubagentManager
from asedclaw.agent.tools.cron import CronTool
from asedclaw.agent.tools.filesystem import FilesystemToolkit
from asedclaw.agent.tools.memory import MemoryToolkit
from asedclaw.agent.tools.send_file import SendFileTool
from asedclaw.agent.tools.registry import ToolRegistry
from asedclaw.agent.tools.shell import ExecTool
from asedclaw.agent.tools.web import WebToolkit
from asedclaw.agent.turn_processor import process_message
from asedclaw.bus.events import InboundMessage, OutboundMessage
from asedclaw.bus.queue import MessageBus
from asedclaw.config.paths import get_media_dir
from asedclaw.providers.base import LLMProvider
from asedclaw.session.manager import Session, SessionManager

if TYPE_CHECKING:
    from asedclaw.config.schema import (
        ContextBudgetConfig,
        ExecToolConfig,
        MessageQueueConfig,
        MemoryFlushConfig,
        ToolResultBudgetConfig,
    )
    from asedclaw.cron.service import CronService


class AgentLoop:
    """Core processing engine: receive → context → LLM → tools → respond."""

    def __init__(
        self,
        bus: MessageBus,
        provider: LLMProvider,
        workspace: Path,
        model: str | None = None,
        temperature: float = 0.1,
        max_tokens: int = 4096,
        memory_window: int = 100,
        reasoning_effort: str | None = None,
        brave_api_key: str | None = None,
        web_proxy: str | None = None,
        exec_config: ExecToolConfig | None = None,
        cron_service: CronService | None = None,
        restrict_to_workspace: bool = False,
        session_manager: SessionManager | None = None,
        mcp_servers: dict | None = None,
        channels_config: Any | None = None,
        max_iterations: int = 40,
        memory_flush_config: MemoryFlushConfig | None = None,
        embedding_config: Any | None = None,
        guardrails_config: Any | None = None,
        context_budget_config: ContextBudgetConfig | None = None,
        tool_result_budget_config: ToolResultBudgetConfig | None = None,
        provider_type: str | None = None,
        message_queue_config: MessageQueueConfig | None = None,
        message_queue_by_channel: dict[str, MessageQueueConfig] | None = None,
    ):
        from asedclaw.config.schema import ExecToolConfig

        self.bus = bus
        self.channels_config = channels_config
        self.provider = provider
        self.workspace = workspace
        self.model = model or provider.get_default_model()
        self.temperature = temperature
        self.memory_window = memory_window
        self.reasoning_effort = reasoning_effort
        self.max_iterations = max_iterations
        self.brave_api_key = brave_api_key
        self.web_proxy = web_proxy
        self.exec_config = exec_config or ExecToolConfig()
        self.cron_service = cron_service
        self.restrict_to_workspace = restrict_to_workspace

        self.context = ContextBuilder(workspace)
        self.sessions = session_manager or SessionManager(workspace)
        self.tools = ToolRegistry()
        self.subagents = SubagentManager(
            provider=provider,
            workspace=workspace,
            bus=bus,
            model=self.model,
            temperature=self.temperature,
            max_tokens=max_tokens,
            reasoning_effort=reasoning_effort,
            brave_api_key=brave_api_key,
            web_proxy=web_proxy,
            exec_config=self.exec_config,
            restrict_to_workspace=restrict_to_workspace,
        )

        # Memory flush config
        if memory_flush_config is None:
            from asedclaw.config.schema import MemoryFlushConfig as _MFC
            memory_flush_config = _MFC()
        self.memory_flush_config = memory_flush_config
        self._flush_done_for_cycle: dict[str, bool] = {}
        self._embedding_config = embedding_config
        self._guardrails_config = guardrails_config
        self._provider_type = provider_type

        # Context Engine for token budget management
        self.context_engine = ContextEngine(
            budget_config=context_budget_config,
            result_budget_config=tool_result_budget_config,
        )

        # Memory index for hybrid search
        self._memory_index = self._create_memory_index()
        self._index_synced_sessions: set[str] = set()

        self._running = False
        self._consolidating: set[str] = set()
        self._consolidation_tasks: set[asyncio.Task] = set()
        self._consolidation_locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary()
        self._session_locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary()
        self._active_tasks: dict[str, list[asyncio.Task]] = {}

        # Message queue for busy-session handling
        self._queue_manager = FollowupQueueManager()
        self._queue_config = message_queue_config
        self._queue_config_by_channel = message_queue_by_channel or {}

        self._register_default_tools()

    # ── Memory index ───────────────────────────────────────────────────

    def _create_memory_index(self):
        """Create MemoryIndexManager using embedding config."""
        try:
            from asedclaw.agent.memory_index import MemoryIndexManager

            emb_cfg = getattr(self, "_embedding_config", None)
            if emb_cfg and emb_cfg.api_base:
                return MemoryIndexManager(
                    workspace=self.workspace,
                    embedding_api_key=emb_cfg.api_key,
                    embedding_api_base=emb_cfg.api_base,
                    embedding_model=emb_cfg.model,
                    embedding_dimensions=emb_cfg.dimensions,
                    embedding_task=emb_cfg.task,
                )
            # No embedding config — still create index for BM25-only search
            return MemoryIndexManager(workspace=self.workspace)
        except Exception:
            logger.debug("Memory index unavailable", exc_info=True)
            return None

    # ── Tool registration ──────────────────────────────────────────────

    def _register_default_tools(self) -> None:
        allowed_dir = self.workspace if self.restrict_to_workspace else None
        self.tools.register_many(
            FilesystemToolkit(workspace=self.workspace, allowed_dir=allowed_dir),
            MemoryToolkit(workspace=self.workspace, index=self._memory_index),
            ExecTool(
                working_dir=str(self.workspace),
                timeout=self.exec_config.timeout,
                restrict_to_workspace=self.restrict_to_workspace,
                path_append=self.exec_config.path_append,
            ),
            WebToolkit(api_key=self.brave_api_key, proxy=self.web_proxy),
            # MessageTool(send_callback=self.bus.publish_outbound, workspace=self.workspace),
            # SpawnTool(manager=self.subagents),
            SendFileTool(send_callback=self.bus.publish_outbound, workspace=self.workspace),
        )
        if self.cron_service:
            self.tools.register(CronTool(self.cron_service))

    # ── Queue decision helpers ────────────────────────────────────────

    def _is_session_busy(self, session_key: str) -> bool:
        """Check if session has an active (non-done) task."""
        tasks = self._active_tasks.get(session_key, [])
        return any(not t.done() for t in tasks)

    def _resolve_queue_mode(self, msg: InboundMessage) -> QueueMode:
        """Resolve the queue mode for an incoming message.

        Priority: inline directive > per-channel config > global config > fallback.
        """
        # 1. Inline directive: /queue <mode>
        inline = self._parse_queue_directive(msg.content)
        if inline is not None:
            return inline

        # 2. Per-channel config
        ch_cfg = self._queue_config_by_channel.get(msg.channel)
        if ch_cfg is not None:
            try:
                return QueueMode(ch_cfg.mode)
            except ValueError:
                pass

        # 3. Global config
        if self._queue_config is not None:
            try:
                return QueueMode(self._queue_config.mode)
            except ValueError:
                pass

        # 4. Fallback
        return QueueMode.FOLLOWUP

    @staticmethod
    def _parse_queue_directive(content: str) -> QueueMode | None:
        """Parse ``/queue <mode>`` directive from message content."""
        stripped = content.strip()
        if not stripped.lower().startswith("/queue "):
            return None
        parts = stripped.split()
        if len(parts) < 2:
            return None
        try:
            return QueueMode(parts[1].lower())
        except ValueError:
            return None

    def _resolve_queue_action(self, msg: InboundMessage) -> str:
        """Decide what to do with an incoming message.

        Returns one of: ``"run"`` | ``"enqueue"`` | ``"interrupt"`` | ``"directive"``.
        """
        # Pure directive message (e.g. "/queue interrupt" with nothing else)
        stripped = msg.content.strip()
        if stripped.lower().startswith("/queue ") and len(stripped.split()) <= 2:
            return "directive"

        session_key = msg.session_key
        is_busy = self._is_session_busy(session_key)
        if not is_busy:
            return "run"

        mode = self._resolve_queue_mode(msg)
        if mode == QueueMode.INTERRUPT:
            return "interrupt"

        # followup / collect / queue → enqueue for later processing
        return "enqueue"

    def _get_queue_settings(self, msg: InboundMessage) -> dict:
        """Build queue settings dict for enqueue call."""
        mode = self._resolve_queue_mode(msg)

        # Start from global defaults
        debounce_ms = 1000
        cap = 20
        drop_policy = QueueDropPolicy.SUMMARIZE

        if self._queue_config is not None:
            debounce_ms = getattr(self._queue_config, "debounce_ms", debounce_ms)
            cap = getattr(self._queue_config, "cap", cap)
            dp = getattr(self._queue_config, "drop_policy", "summarize")
            try:
                drop_policy = QueueDropPolicy(dp)
            except ValueError:
                pass

        # Per-channel override
        ch_cfg = self._queue_config_by_channel.get(msg.channel)
        if ch_cfg is not None:
            debounce_ms = getattr(ch_cfg, "debounce_ms", debounce_ms)
            cap = getattr(ch_cfg, "cap", cap)
            dp = getattr(ch_cfg, "drop_policy", drop_policy.value)
            try:
                drop_policy = QueueDropPolicy(dp)
            except ValueError:
                pass

        return dict(mode=mode, debounce_ms=debounce_ms, cap=cap, drop_policy=drop_policy)

    # ── Main loop ──────────────────────────────────────────────────────

    async def run(self) -> None:
        self._running = True
        logger.info("Agent loop started")
        msg_count = 0

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

            msg_count += 1
            if msg_count % 20 == 0:
                self._cleanup_staged_media()

            cmd_lower = msg.content.strip().lower()
            if cmd_lower == "/stop":
                await self._handle_stop(msg)
                continue
            if cmd_lower == "/stop-current":
                await self._handle_stop_current(msg)
                continue

            action = self._resolve_queue_action(msg)

            if action == "directive":
                # Pure /queue directive — acknowledge and apply
                mode = self._parse_queue_directive(msg.content)
                label = mode.value if mode else "unknown"
                await self.bus.publish_outbound(OutboundMessage(
                    channel=msg.channel, chat_id=msg.chat_id,
                    content=f"Queue mode set to **{label}** for this session.",
                ))
                continue

            if action == "interrupt":
                # Cancel active tasks + clear queue, then dispatch immediately
                await self._interrupt_session(msg.session_key)
                logger.info("Interrupt: cancelled tasks for session={}", msg.session_key)

            if action == "enqueue":
                settings = self._get_queue_settings(msg)
                accepted = self._queue_manager.enqueue(msg.session_key, msg, **settings)
                if accepted:
                    logger.info(
                        "Enqueued message for session={} depth={}",
                        msg.session_key, self._queue_manager.get_depth(msg.session_key),
                    )
                continue

            # action == "run" or "interrupt" (after cancellation)
            self._create_dispatch_task(msg)

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

    def _create_dispatch_task(self, msg: InboundMessage) -> asyncio.Task:
        """Create a dispatch task and register it in _active_tasks."""
        task = asyncio.create_task(self._dispatch(msg))
        self._active_tasks.setdefault(msg.session_key, []).append(task)
        task.add_done_callback(
            lambda t, k=msg.session_key: (
                self._active_tasks.get(k, []) and t in self._active_tasks.get(k, [])
                and self._active_tasks[k].remove(t)
            )
        )
        return task

    async def _interrupt_session(self, session_key: str) -> int:
        """Cancel all active tasks and clear the queue for a session."""
        tasks = self._active_tasks.pop(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 = await self.subagents.cancel_by_session(session_key)
        queue_cleared = self._queue_manager.clear(session_key)
        return cancelled + sub_cancelled + queue_cleared

    # ── Dispatch & process ─────────────────────────────────────────────

    async def _handle_stop_current(self, msg: InboundMessage) -> None:
        """Cancel the current task and subagents but keep the queue intact.

        After cancellation, the queue drain in _dispatch will pick up
        any remaining queued messages automatically.
        """
        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 = await self.subagents.cancel_by_session(msg.session_key)
        total = cancelled + sub_cancelled
        queue_depth = self._queue_manager.get_depth(msg.session_key)
        if queue_depth:
            note = f"Stopped {total} task(s). {queue_depth} queued message(s) will continue."
        else:
            note = 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=note,
            metadata={"_task_cancel": True, "_task_id": msg.metadata.get("_task_id")},
        ))

    async def _handle_stop(self, msg: InboundMessage) -> None:
        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 = await self.subagents.cancel_by_session(msg.session_key)
        queue_cleared = self._queue_manager.clear(msg.session_key)
        total = cancelled + sub_cancelled + queue_cleared
        await self.bus.publish_outbound(OutboundMessage(
            channel=msg.channel, chat_id=msg.chat_id,
            content=f"Stopped {total} task(s)." if total else "No active task to stop.",
        ))

    async def _dispatch(self, msg: InboundMessage) -> None:
        lock = self._session_locks.setdefault(msg.session_key, asyncio.Lock())
        async with 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:
                # Cancel path: _handle_stop_current already sends _task_cancel.
                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.",
                ))

            # Signal task completion so the frontend can finalize the task card.
            # task_id is written into msg.metadata by turn_processor after generation.
            # Runs after normal completion or error, but NOT after CancelledError
            # (cancel path has its own _task_cancel signal via _handle_stop_current).
            task_id = (msg.metadata or {}).get("_task_id")
            if task_id:
                await self.bus.publish_outbound(OutboundMessage(
                    channel=msg.channel, chat_id=msg.chat_id,
                    content="",
                    metadata={"_task_done": True, "_task_id": task_id},
                ))

            # Drain any queued followup messages while we still hold the lock
            try:
                await self._drain_followup(msg.session_key)
            except asyncio.CancelledError:
                raise
            except Exception:
                logger.exception("Error draining followup queue for session {}", msg.session_key)

    # ── Followup queue drain ──────────────────────────────────────────

    async def _drain_followup(self, session_key: str) -> None:
        """Process queued messages after the current turn finishes."""
        if self._queue_manager.get_depth(session_key) == 0:
            return

        mode = self._queue_manager.get_mode(session_key) or QueueMode.FOLLOWUP

        if mode == QueueMode.COLLECT:
            await self._drain_collect(session_key)
        else:
            await self._drain_individual(session_key)

    async def _drain_collect(self, session_key: str) -> None:
        """Collect mode: wait debounce window, then batch all queued messages."""
        debounce_s = self._queue_manager.get_debounce_ms(session_key) / 1000.0
        if debounce_s > 0:
            await asyncio.sleep(debounce_s)

        items = self._queue_manager.drain_all_collected(session_key)
        if not items:
            return

        # Build combined prompt
        prompt_parts = ["[以下是 agent 忙碌期间收到的排队消息]"]
        for i, item in enumerate(items, 1):
            prompt_parts.append(f"---\n排队消息 #{i}\n{item.msg.content}")

        # Include dropped message summary if any
        dropped_lines = self._queue_manager.get_dropped_summary(session_key)
        if dropped_lines:
            prompt_parts.append(
                f"\n[另有 {len(dropped_lines)} 条消息因队列已满被省略，"
                f"摘要: {'; '.join(dropped_lines[:5])}]"
            )

        last_msg = items[-1].msg
        combined = InboundMessage(
            channel=last_msg.channel,
            sender_id=last_msg.sender_id,
            chat_id=last_msg.chat_id,
            content="\n".join(prompt_parts),
            metadata=last_msg.metadata,
            session_key_override=last_msg.session_key_override,
        )

        logger.info(
            "Drain collect: processing {} batched messages for session={}",
            len(items), session_key,
        )
        try:
            response = await self._process_message(combined)
            if response:
                await self.bus.publish_outbound(response)
        except asyncio.CancelledError:
            raise
        except Exception:
            logger.exception("Error processing collected messages for session {}", session_key)

    async def _drain_individual(self, session_key: str) -> None:
        """Followup/queue mode: process queued messages one by one."""
        drained = 0
        while True:
            item = self._queue_manager.drain_next(session_key)
            if item is None:
                break
            drained += 1
            logger.debug("Drain individual #{} for session={}", drained, session_key)
            try:
                response = await self._process_message(item.msg)
                if response:
                    await self.bus.publish_outbound(response)
            except asyncio.CancelledError:
                raise
            except Exception:
                logger.exception("Error draining followup #{} for session {}", drained, session_key)
        if drained > 0:
            logger.info("Drained {} followup messages for session={}", drained, session_key)

    async def _process_message(
        self,
        msg: InboundMessage,
        session_key: str | None = None,
        on_progress: Callable[[str], Awaitable[None]] | None = None,
    ) -> OutboundMessage | None:
        """Delegate to turn_processor.process_message with wired dependencies."""
        return await process_message(
            msg,
            bus=self.bus,
            sessions=self.sessions,
            context=self.context,
            context_engine=self.context_engine,
            provider=self.provider,
            model=self.model,
            temperature=self.temperature,
            reasoning_effort=self.reasoning_effort,
            tools=self.tools,
            workspace=self.workspace,
            memory_window=self.memory_window,
            max_iterations=self.max_iterations,
            guardrails_config=self._guardrails_config,
            provider_type=self._provider_type,
            on_pre_process=self._pre_process_session,
            on_new_session=self._on_new_session,
            session_key=session_key,
            on_progress=on_progress,
        )

    # ── Session lifecycle callbacks ───────────────────────────────────

    async def _pre_process_session(self, session: Session) -> None:
        """Run consolidation, memory flush, and index sync before a turn."""
        maybe_consolidate(
            session,
            memory_window=self.memory_window,
            workspace=self.workspace,
            provider=self.provider,
            model=self.model,
            consolidating=self._consolidating,
            consolidation_locks=self._consolidation_locks,
            consolidation_tasks=self._consolidation_tasks,
            flush_done_for_cycle=self._flush_done_for_cycle,
            memory_index=self._memory_index,
        )
        await maybe_memory_flush(
            session,
            provider=self.provider,
            model=self.model,
            workspace=self.workspace,
            memory_flush_config=self.memory_flush_config,
            flush_done_for_cycle=self._flush_done_for_cycle,
        )
        await maybe_sync_index(
            session,
            memory_index=self._memory_index,
            index_synced_sessions=self._index_synced_sessions,
        )

    async def _on_new_session(self, msg: InboundMessage, session: Session) -> OutboundMessage:
        """Handle /new: delegate to session_lifecycle, then save/invalidate."""
        result = await handle_new_session(
            msg, session,
            workspace=self.workspace,
            provider=self.provider,
            model=self.model,
            memory_window=self.memory_window,
            consolidation_locks=self._consolidation_locks,
            consolidating=self._consolidating,
            memory_index=self._memory_index,
        )
        if result.content == "New session started.":
            self.sessions.save(session)
            self.sessions.invalidate(session.key)
        return result

    def _cleanup_staged_media(self) -> None:
        try:
            from asedclaw.media.store import MediaStore
            MediaStore(get_media_dir()).cleanup()
            MediaStore(get_media_dir("ased")).cleanup()
        except Exception:
            pass

    # ── Public API ─────────────────────────────────────────────────────

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