"""LLM-based routing between simple chat and internal Task mode."""

from __future__ import annotations

import asyncio
import json
from typing import Any

from .models import MainAgentDecision, TaskRecord


class MainAgentRouter:
    """Semantic router for deciding whether a message belongs to a Task."""

    async def classify(
        self,
        message: str,
        *,
        active_task: TaskRecord | None = None,
        provider: Any | None = None,
        model: str | None = None,
        recent_messages: list[dict[str, Any]] | None = None,
        intent_skill: str | None = None,
        thinking_enabled: bool | None = None,
        timeout_seconds: float = 8.0,
    ) -> MainAgentDecision:
        if provider is None:
            return self._apply_active_task_boundary(
                self._fallback(active_task=active_task, reason="router_provider_unavailable"),
                message=message,
                active_task=active_task,
            )
        chat_kwargs: dict[str, Any] = {
            "messages": [
                {
                    "role": "system",
                    "content": (
                        "You are Beaver's Intent Agent. Your only job is to route the user's "
                        "message to simple chat or internal Task mode. Return only compact JSON. "
                        "Do not answer the user. Do not explain."
                    ),
                },
                {
                    "role": "user",
                    "content": self._prompt(
                        message=message,
                        active_task=active_task,
                        recent_messages=recent_messages or [],
                        intent_skill=intent_skill,
                    ),
                },
            ],
            "tools": None,
            "model": model,
            "max_tokens": 256,
            "temperature": 0.0,
        }
        if thinking_enabled is not None:
            chat_kwargs["thinking_enabled"] = thinking_enabled

        last_error: Exception | None = None
        for attempt_timeout in (timeout_seconds, 12.0):
            try:
                response = await asyncio.wait_for(provider.chat(**chat_kwargs), timeout=attempt_timeout)
                return self._apply_active_task_boundary(
                    self.from_json(response.content or "", active_task=active_task),
                    message=message,
                    active_task=active_task,
                )
            except Exception as exc:
                last_error = exc
        return self._apply_active_task_boundary(
            self._fallback(active_task=active_task, reason=f"router_failed: {last_error}"),
            message=message,
            active_task=active_task,
        )

    def from_json(self, text: str, *, active_task: TaskRecord | None = None) -> MainAgentDecision:
        payload = self._parse_json_object(text)
        raw_action = str(payload.get("action") or payload.get("mode") or "").strip().lower()
        reason = str(payload.get("reason") or raw_action or "llm_router")
        short_title = _clean_short_title(payload.get("short_title") or payload.get("title"))

        if raw_action in {"revise_task", "revise", "revision", "needs_revision"}:
            return MainAgentDecision(
                mode="task",
                reason=reason,
                starts_new_task=active_task is None,
                short_title=short_title,
                action="revise_task" if active_task is not None else "create_task",
            )
        if raw_action in {"continue_task", "continue", "task"}:
            return MainAgentDecision(
                mode="task",
                reason=reason,
                starts_new_task=active_task is None,
                short_title=short_title,
                action="continue_task" if active_task is not None else "create_task",
            )
        if raw_action in {"new_task", "new"}:
            return MainAgentDecision(
                mode="task",
                reason=reason,
                starts_new_task=True,
                short_title=short_title,
                action="create_task",
            )
        if raw_action in {"close_task", "close", "done", "finish"}:
            return MainAgentDecision(
                mode="simple",
                reason=reason,
                closes_task=active_task is not None,
                short_title=short_title,
                action="close_task",
            )
        if raw_action in {"abandon_task", "abandon", "cancel_task"}:
            return MainAgentDecision(
                mode="simple",
                reason=reason,
                abandons_task=active_task is not None,
                short_title=short_title,
                action="abandon_task",
            )
        return MainAgentDecision(
            mode="simple",
            reason=reason or "simple_chat",
            short_title=short_title,
            action="simple_chat",
        )

    def _fallback(self, *, active_task: TaskRecord | None, reason: str) -> MainAgentDecision:
        if active_task is not None:
            return MainAgentDecision(mode="task", reason=reason, action="continue_task")
        return MainAgentDecision(mode="simple", reason=reason, action="simple_chat")

    def _apply_active_task_boundary(
        self,
        decision: MainAgentDecision,
        *,
        message: str,
        active_task: TaskRecord | None,
    ) -> MainAgentDecision:
        if active_task is None or decision.action != "continue_task":
            return decision
        if not _looks_like_fresh_task_request(message):
            return decision
        if _looks_like_explicit_task_followup(message):
            return decision
        title = decision.short_title or active_task.metadata.get("short_title")
        return MainAgentDecision(
            mode="task",
            reason=(
                "fresh standalone task request in the same session; "
                "do not attach it to the active task without explicit follow-up wording"
            ),
            starts_new_task=True,
            short_title=title,
            action="create_task",
        )

    @staticmethod
    def _prompt(
        *,
        message: str,
        active_task: TaskRecord | None,
        recent_messages: list[dict[str, Any]],
        intent_skill: str | None,
    ) -> str:
        active_task_payload = None
        if active_task is not None:
            active_task_payload = {
                "task_id": active_task.task_id,
                "description": active_task.description,
                "goal": active_task.goal,
                "status": active_task.status,
                "short_title": active_task.metadata.get("short_title"),
            }
        recent = [
            {"role": item.get("role"), "content": str(item.get("content") or "")[:500]}
            for item in recent_messages[-8:]
            if item.get("role") in {"user", "assistant"}
        ]
        skill_section = (
            f"Intent Agent skill guidance:\n{intent_skill.strip()}\n\n"
            if intent_skill and intent_skill.strip()
            else ""
        )
        return (
            "Decide how to route the current user message.\n\n"
            f"{skill_section}"
            "Actions:\n"
            "- simple_chat: no Task should be created or continued.\n"
            "- continue_task: keep the user in the active Task.\n"
            "- revise_task: user asks to change, correct, refine, expand, reformat, or redo the latest active Task result.\n"
            "- new_task: start a separate new Task.\n"
            "- close_task: user explicitly says the active Task is done/satisfactory/finished.\n"
            "- abandon_task: user explicitly says to stop, cancel, abandon, or no longer do the active Task.\n\n"
            "Critical policy:\n"
            "- A Session is the durable conversation/device/group context. A Task is one unit of work inside that Session. "
            "Do not use an active Task as a reason to merge every later message into the same work item.\n"
            "- If there is an active Task, choose continue_task only when the current message explicitly depends on, extends, "
            "or asks a direct follow-up about that active Task's latest result.\n"
            "- With an active Task, choose simple_chat for unrelated lightweight conversation and new_task for unrelated work "
            "that needs Task capabilities. Either decision starts a new topic.\n"
            "- An unrelated lightweight conversation must not be classified as revise_task merely because the active Task is awaiting acceptance.\n"
            "- Choose revise_task when the active Task is awaiting feedback or needs revision and the user asks for changes "
            "such as '改一下', '加上', '删除', '换成', '再详细点', '格式改成', '不要', or equivalent wording.\n"
            "- Choose continue_task for neutral follow-up questions or additional next steps that refer to the previous result, "
            "for example '顺便查一下深圳', '这个也加上', or '继续'.\n"
            "- A standalone tool-dependent request such as a fresh weather/search/file/run/test request is new_task even when it is "
            "similar to the active Task. Repeating '珠海天气怎么样' later is a new Task unless the user says to revise or continue the old result.\n"
            "- If there is no active Task, choose new_task only for work that requires execution, iteration, tools, files, "
            "implementation, validation, or multi-step completion. Otherwise choose simple_chat.\n"
            "- Requests that need current, real-time, external, user-private, local-file, web, weather, price, news, "
            "calendar, email, or system data require tools. Choose new_task for them because the Intent Agent has no tools.\n"
            "- The Intent Agent must never answer tool-dependent requests itself or apologize for lacking tools. "
            "It only routes the request so the main Task agent can use tools.\n"
            "- short_title must be 5-15 Chinese characters or a similarly short English phrase when a Task is involved.\n\n"
            "Return JSON only with keys: action, reason, short_title.\n\n"
            f"Active task:\n{json.dumps(active_task_payload, ensure_ascii=False)}\n\n"
            f"Recent conversation:\n{json.dumps(recent, ensure_ascii=False)}\n\n"
            f"Current user message:\n{message}"
        )

    @staticmethod
    def _parse_json_object(text: str) -> dict[str, Any]:
        cleaned = text.strip()
        if cleaned.startswith("```"):
            cleaned = cleaned.strip("`")
            if cleaned.lower().startswith("json"):
                cleaned = cleaned[4:].strip()
        start = cleaned.find("{")
        end = cleaned.rfind("}")
        if start >= 0 and end >= start:
            cleaned = cleaned[start : end + 1]
        payload = json.loads(cleaned)
        if not isinstance(payload, dict):
            raise ValueError("router response must be a JSON object")
        return payload


def _clean_short_title(value: Any) -> str | None:
    if value in (None, ""):
        return None
    title = " ".join(str(value).strip().split())
    return title[:40] or None


def _looks_like_explicit_task_followup(message: str) -> bool:
    text = _compact_text(message)
    if not text:
        return False
    markers = (
        "继续",
        "接着",
        "上面",
        "刚才",
        "前面",
        "这个",
        "那个",
        "它",
        "结果",
        "再",
        "也",
        "顺便",
        "补充",
        "加上",
        "加入",
        "删除",
        "去掉",
        "改",
        "换成",
        "重做",
        "详细",
        "展开",
        "格式",
        "continue",
        "same task",
        "previous",
        "above",
        "that result",
        "revise",
        "update it",
        "add",
        "remove",
        "change",
        "also",
    )
    return any(marker in text for marker in markers)


def _looks_like_fresh_task_request(message: str) -> bool:
    text = _compact_text(message)
    if not text:
        return False
    markers = (
        "天气",
        "气温",
        "下雨",
        "降雨",
        "空气质量",
        "预报",
        "查一下",
        "帮我查",
        "搜索",
        "搜一下",
        "看看最新",
        "最新",
        "今天",
        "明天",
        "上传",
        "下载",
        "文件",
        "运行",
        "执行",
        "测试",
        "构建",
        "部署",
        "修复",
        "weather",
        "forecast",
        "temperature",
        "search",
        "look up",
        "latest",
        "today",
        "tomorrow",
        "upload",
        "download",
        "file",
        "run",
        "execute",
        "test",
        "build",
        "deploy",
        "fix",
    )
    return any(marker in text for marker in markers)


def _compact_text(message: str) -> str:
    return " ".join(str(message or "").strip().lower().split())
