from __future__ import annotations

import asyncio
import json
import os
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING, Any

import aiosqlite

from aicd.agent.chat_context import (
    build_hydrated_messages,
    longest_valid_openai_suffix,
    repair_openai_chat_messages,
)
from aicd.agent.chat_serialize import tool_calls_to_jsonable
from aicd.agent.llm import LLMClient
from aicd.agent.idle_pulse import sub_agent_idle_pulse_user_text
from aicd.agent.prompts import MAIN_SYSTEM, SUB_SYSTEM
from aicd.env_probe import RUNTIME_ENV_FILENAME
from aicd.agent.tool_router import ToolRouter
from aicd.agent.tools_openai import openai_tools
from aicd.bus.event_bus import AsyncEventBus
from aicd.config import AppConfig
from aicd.db.store import TaskStore
from aicd.runtime_task_context import current_task_id
from aicd.tasks.manager import TaskManager

if TYPE_CHECKING:
    from aicd.session_log import SessionLogger


def _trim_chat_messages_to_last_user_turn(messages: list[dict[str, Any]]) -> None:
    """
    原地裁剪：只保留「最后一条 user」及其后的消息。

    ``chat_new_session`` 会切换 DB session_id；若在工具循环中途 ``messages.clear()``，
    会删掉带 ``tool_calls`` 的 assistant，只留下孤立的 ``role=tool``，触发 API 400。
    因此在整轮 tool 结果都追加后，再删掉更早轮次，保留 ``user → assistant(tool_calls) → tool*`` 链。
    """
    last_u = -1
    for i, m in enumerate(messages):
        if m.get("role") == "user":
            last_u = i
    if last_u > 0:
        del messages[:last_u]


def _sub_agent_cancel_requested() -> bool:
    ct = asyncio.current_task()
    if ct is None:
        return False
    fn = getattr(ct, "cancelling", None)
    return callable(fn) and bool(fn())


def _format_task_inbox_block(rows: list[Any]) -> str:
    parts: list[str] = []
    for m in rows:
        body = (
            json.dumps(m.body, ensure_ascii=False)
            if getattr(m, "body", None) is not None
            else "{}"
        )
        parts.append(
            f"[task_message id={m.id} from={m.from_task_id} kind={m.kind}]\n{body}"
        )
    return "\n\n".join(parts)


def _trim_sub_agent_messages(messages: list[dict[str, Any]], *, max_msgs: int) -> None:
    """
    子 Agent 上下文过长时裁剪为最近若干条。

    禁止简单 ``tail = messages[-k:]``：可能切断 ``assistant(tool_calls)`` 与其 ``tool`` 结果，
    导致 DashScope/OpenAI 兼容接口报
    ``must be a \"tool\" message``。改为取**合法**的最长后缀。
    """
    if not messages:
        return
    sys0 = messages[0]
    body = messages[1:]
    max_body = max(0, max_msgs - 1)
    if max_body <= 0:
        messages.clear()
        messages.append(sys0)
        return
    if len(body) <= max_body:
        repair_openai_chat_messages(body)
        messages.clear()
        messages.extend([sys0, *body])
        return
    picked = longest_valid_openai_suffix(body, max_body)
    repair_openai_chat_messages(picked)
    messages.clear()
    messages.extend([sys0, *picked])


def _msg_assistant_dict(msg: Any) -> dict[str, Any]:
    out: dict[str, Any] = {
        "role": "assistant",
        "content": msg.content,
    }
    if msg.tool_calls:
        out["tool_calls"] = [
            {
                "id": tc.id,
                "type": "function",
                "function": {
                    "name": tc.function.name,
                    "arguments": tc.function.arguments or "{}",
                },
            }
            for tc in msg.tool_calls
        ]
    return out


class AgentLoop:
    def __init__(
        self,
        cfg: AppConfig,
        llm: LLMClient,
        router: ToolRouter,
        task_manager: TaskManager,
        store: TaskStore,
        conn: aiosqlite.Connection,
        layout: dict[str, Path],
        bus: AsyncEventBus,
        session_id: str,
        session_log: SessionLogger | None = None,
        env_capabilities_block: str = "",
    ) -> None:
        self._cfg = cfg
        self._llm = llm
        self._router = router
        self._tasks = task_manager
        self._store = store
        self._conn = conn
        self._layout = layout
        self._bus = bus
        self._session_id = session_id
        self._session_log = session_log
        self._env_capabilities_block = (env_capabilities_block or "").strip()
        self._messages: list[dict[str, Any]] = []
        self._strip_chat_prefix_after_tool_round = False
        self._workspace_path: Path = router.cwd
        o, t = router.get_ai_paths()
        self._ai_out: Path = o
        self._ai_tmp: Path = t

    def apply_chat_session_id(self, new_session_id: str) -> None:
        """切换主会话键；新消息写入新 session_id（旧记录在 DB 中仍保留）。

        不在此处清空 ``_messages``：工具 ``chat_new_session`` 可能在同一轮 assistant
        ``tool_calls`` 执行过程中被调用，立即清空会导致下一条请求只有孤立的 ``tool`` 角色。
        改在本轮所有 tool 结果写入后，再裁剪到「最后一条 user」起的后缀。
        """
        self._session_id = new_session_id
        self._strip_chat_prefix_after_tool_round = True

    def _log_llm_usage(self, resp: Any) -> None:
        slog = self._session_log
        if not slog:
            return
        u = getattr(resp, "usage", None)
        if u is not None:
            slog.write_event("llm", "LLM", "chat.completions", "", str(u))

    async def _prepend_main_inbox_if_any(self) -> None:
        """主对话：将 main-inbox-<session> 中未投递的 task_message 注入为 user（并 mark read）。"""
        from aicd.tasks.main_inbox import main_inbox_id

        try:
            to_id = main_inbox_id(self._session_id)
        except ValueError:
            return
        rows = await self._store.list_task_messages_unread(
            self._conn,
            to_task_id=to_id,
            limit=50,
        )
        if not rows:
            return
        block = _format_task_inbox_block(rows)
        await self._store.mark_task_messages_read(self._conn, [m.id for m in rows])
        self._messages.append(
            {
                "role": "user",
                "content": (
                    "【主收件箱】以下为未读消息（含子任务 **lifecycle** 与其它上行 **task_message**）。"
                    "请视情况 **task_result**、**task_session_children_digest** 或发起新动作。\n\n"
                    + block
                ),
            },
        )

    def set_workspace(self, path: Path) -> None:
        """与 ToolRouter.set_cwd 同步，用于系统提示中的工作目录说明。"""
        self._workspace_path = path.resolve()

    def set_ai_storage_paths(self, output: Path, tmp: Path) -> None:
        self._ai_out = output.resolve()
        self._ai_tmp = tmp.resolve()

    def _runtime_env_file_str(self) -> str:
        return str((self._layout["data"] / RUNTIME_ENV_FILENAME).resolve())

    async def hydrate_chat_from_store(self) -> None:
        """从 SQLite 恢复本会话消息（按条数尾部 + token 预算 + 可选摘记）。"""
        total = await self._store.count_chat_messages(self._conn, self._session_id)
        if total == 0:
            self._messages = []
            return
        max_m = self._cfg.agent.chat_history_max_messages
        truncated = total > max_m
        tail = min(total, max_m)
        offset = max(0, total - tail)
        rows = await self._store.list_chat_messages(
            self._conn, self._session_id, limit=tail, offset=offset
        )
        latest_summary: str | None = None
        if truncated:
            sums = await self._store.list_chat_summaries(
                self._conn, self._session_id, limit=50
            )
            if sums:
                latest_summary = sums[-1].content
        self._messages = build_hydrated_messages(
            rows=rows,
            truncated_by_count=truncated,
            latest_summary=latest_summary,
            context_budget_tokens=self._cfg.agent.context_budget_tokens,
        )

    def _system_content(self) -> str:
        base = (
            f"{MAIN_SYSTEM}\n\n"
            f"AI workspace: {self._workspace_path}\n"
            f"AI output directory: {self._ai_out}\n"
            f"AI tmp (scratch/intermediate, self-maintain): {self._ai_tmp}\n"
            f"Main chat session_id: {self._session_id}\n"
            f"Runtime env file: {self._runtime_env_file_str()}"
        )
        if self._env_capabilities_block:
            base = f"{base}\n\n{self._env_capabilities_block}"
        return base

    async def run_sub_agent_worker(self, task_id: str, instruction: str) -> None:
        slog = self._session_log
        if slog:
            ins = instruction if len(instruction) <= 8000 else instruction[:7997] + "..."
            slog.write_event(
                "task",
                "AgentLoop",
                "sub_agent_start",
                f"task_id={task_id}",
                ins,
            )
        ws = self._layout["tasks"] / task_id / "workspace"
        ws.mkdir(parents=True, exist_ok=True)
        out, tmp = self._router.get_ai_paths()
        row = await self._store.get_task(self._conn, task_id)
        payload: dict[str, Any] = row.payload if row else {}
        scratch_s = payload.get("scratch_dir")
        sub_tmp = (
            Path(scratch_s).resolve()
            if isinstance(scratch_s, str) and scratch_s.strip()
            else (tmp / "tasks" / task_id).resolve()
        )
        sub_tmp.mkdir(parents=True, exist_ok=True)
        token = current_task_id.set(task_id)
        try:
            sub = ToolRouter(
                self._cfg,
                layout=self._layout,
                task_manager=self._tasks,
                cwd=ws,
                session_log=slog,
                chat_session_id=None,
            )
            sub.set_ai_paths(out, sub_tmp)
            sub.set_vision_llm(self._llm)
            sub_sys = SUB_SYSTEM.format(
                cwd=str(ws),
                ai_out=str(out),
                ai_tmp=str(sub_tmp),
                runtime_env_file=self._runtime_env_file_str(),
            )
            extras: list[str] = []
            dd = payload.get("deliverable_dir")
            if isinstance(dd, str) and dd.strip():
                extras.append(
                    f"User-visible deliverable directory (prefer final artifacts here): {dd.strip()}"
                )
            ro = payload.get("read_only_paths")
            if isinstance(ro, list) and ro:
                paths = [str(x) for x in ro if isinstance(x, str) and x.strip()]
                if paths:
                    extras.append(
                        "Read-only reference paths (do not modify; you may read): "
                        + "; ".join(paths[:20])
                    )
            he = payload.get("parent_handoff_expectation")
            if isinstance(he, str) and he.strip():
                extras.append(f"Parent handoff expectation: {he.strip()[:1200]}")
            if extras:
                sub_sys = f"{sub_sys}\n\n" + "\n".join(extras)
            if self._env_capabilities_block:
                sub_sys = f"{sub_sys}\n\n{self._env_capabilities_block}"
            messages: list[dict[str, Any]] = [
                {"role": "system", "content": sub_sys},
                {"role": "user", "content": instruction},
            ]
            tools = openai_tools()
            max_r = self._cfg.agent.sub_agent_max_tool_rounds
            total_cap = max(300, max_r * 6)
            total_rounds = 0
            last_inbox_id = 0
            max_ctx_msgs = 100

            while True:
                if _sub_agent_cancel_requested():
                    raise asyncio.CancelledError
                rounds_left = min(max_r, total_cap - total_rounds)
                if rounds_left <= 0:
                    if slog:
                        slog.write_event(
                            "task",
                            "AgentLoop",
                            "sub_agent_total_cap",
                            f"task_id={task_id}",
                            f"cumulative={total_rounds} cap={total_cap}",
                        )
                    await self._bus.publish(
                        {
                            "topic": "task",
                            "task_id": task_id,
                            "event": "log",
                            "line": f"[ai] 子 Agent 已达累计 LLM 轮次上限 {total_cap}，结束",
                        }
                    )
                    return

                got_final = False
                for round_i in range(rounds_left):
                    if _sub_agent_cancel_requested():
                        raise asyncio.CancelledError
                    total_rounds += 1
                    burst_idx = round_i + 1
                    _trim_sub_agent_messages(messages, max_msgs=max_ctx_msgs)
                    await self._tasks.patch_task_payload(
                        task_id,
                        {
                            "sub_agent_trace": {
                                "phase": "llm_turn",
                                "llm_round": burst_idx,
                                "max_llm_rounds": max_r,
                                "cumulative_turns": total_rounds,
                                "cumulative_cap": total_cap,
                                "updated_at": datetime.now(timezone.utc).isoformat(),
                            }
                        },
                    )
                    await self._bus.publish(
                        {
                            "topic": "task",
                            "task_id": task_id,
                            "event": "progress",
                            "phase": "sub_agent",
                            "llm_round": burst_idx,
                            "max_llm_rounds": max_r,
                            "message": (
                                f"子 Agent：本轮第 {burst_idx}/{rounds_left} 次 LLM"
                                f"（累计 {total_rounds}/{total_cap}）"
                            ),
                        }
                    )
                    resp = await self._llm.chat(messages, tools)
                    self._log_llm_usage(resp)
                    choice = resp.choices[0]
                    msg = choice.message
                    if not msg.tool_calls:
                        final = msg.content or ""
                        messages.append({"role": "assistant", "content": final})
                        await self._bus.publish(
                            {
                                "topic": "task",
                                "task_id": task_id,
                                "event": "log",
                                "line": f"[ai] {final[:2000]}",
                            }
                        )
                        if slog:
                            slog.write_event(
                                "task",
                                "AgentLoop",
                                "sub_agent_reply",
                                f"task_id={task_id}",
                                final if len(final) <= 12000 else final[:11997] + "...",
                            )
                        got_final = True
                        break
                    messages.append(_msg_assistant_dict(msg))
                    tool_names: list[str] = []
                    for tc in msg.tool_calls:
                        if _sub_agent_cancel_requested():
                            raise asyncio.CancelledError
                        name = tc.function.name
                        tool_names.append(name)
                        raw_args = tc.function.arguments or "{}"
                        result = await sub.run(name, raw_args)
                        line = f"tool {name} -> {str(result)[:800]}"
                        await self._bus.publish(
                            {
                                "topic": "task",
                                "task_id": task_id,
                                "event": "log",
                                "line": line,
                            }
                        )
                        messages.append(
                            {
                                "role": "tool",
                                "tool_call_id": tc.id,
                                "content": json.dumps(result, ensure_ascii=False)[
                                    :24_000
                                ],
                            }
                        )
                    last_tool = tool_names[-1] if tool_names else ""
                    await self._tasks.patch_task_payload(
                        task_id,
                        {
                            "sub_agent_trace": {
                                "phase": "after_tools",
                                "llm_round": burst_idx,
                                "max_llm_rounds": max_r,
                                "cumulative_turns": total_rounds,
                                "tools_in_round": tool_names,
                                "last_tool": last_tool,
                                "updated_at": datetime.now(timezone.utc).isoformat(),
                            }
                        },
                    )
                    await self._bus.publish(
                        {
                            "topic": "task",
                            "task_id": task_id,
                            "event": "progress",
                            "phase": "sub_agent",
                            "llm_round": burst_idx,
                            "max_llm_rounds": max_r,
                            "last_tool": last_tool,
                            "tools": ",".join(tool_names[:12]),
                            "message": (
                                f"子 Agent：本轮 {burst_idx}/{rounds_left} 已执行 "
                                f"{len(tool_names)} 个工具"
                                + (f"（末次 {last_tool}）" if last_tool else "")
                            ),
                        }
                    )

                if not got_final and slog:
                    slog.write_event(
                        "task",
                        "AgentLoop",
                        "sub_agent_burst_exhausted",
                        f"task_id={task_id}",
                        f"burst_rounds={rounds_left}",
                    )

                await asyncio.sleep(0.12)
                inbox_rows = await self._store.list_task_messages(
                    self._conn,
                    to_task_id=task_id,
                    since_id=last_inbox_id,
                    limit=50,
                )
                hb_sub = self._cfg.agent.sub_agent_idle_heartbeat_sec
                if not inbox_rows and hb_sub > 0:
                    children = await self._store.list_tasks_by_parent(
                        self._conn, task_id
                    )
                    if any(r.status in ("running", "pending") for r in children):
                        sleep_s = max(5, min(3600, hb_sub))
                        await asyncio.sleep(sleep_s)
                        inbox_rows = await self._store.list_task_messages(
                            self._conn,
                            to_task_id=task_id,
                            since_id=last_inbox_id,
                            limit=50,
                        )

                if inbox_rows:
                    last_inbox_id = max(m.id for m in inbox_rows)
                    block = _format_task_inbox_block(inbox_rows)
                    messages.append(
                        {
                            "role": "user",
                            "content": (
                                "以下是通过 **task_message** 送达的新指令/补充"
                                "（上级主 AI 或并列任务）。请据此继续执行；若与当前工作冲突，"
                                "以最新指令为准并简要说明取舍。\n\n"
                                + block
                            ),
                        }
                    )
                    await self._store.mark_task_messages_read(
                        self._conn, [m.id for m in inbox_rows]
                    )
                    await self._bus.publish(
                        {
                            "topic": "task",
                            "task_id": task_id,
                            "event": "log",
                            "line": f"[inbox] 收到 {len(inbox_rows)} 条消息，继续执行",
                        }
                    )
                    await self._tasks.patch_task_payload(
                        task_id,
                        {
                            "sub_agent_trace": {
                                "phase": "inbox_continue",
                                "last_inbox_id": last_inbox_id,
                                "updated_at": datetime.now(timezone.utc).isoformat(),
                            }
                        },
                    )
                    continue

                if hb_sub <= 0:
                    return
                children2 = await self._store.list_tasks_by_parent(
                    self._conn, task_id
                )
                if not any(r.status in ("running", "pending") for r in children2):
                    return
                messages.append(
                    {
                        "role": "user",
                        "content": sub_agent_idle_pulse_user_text(),
                    }
                )
                await self._bus.publish(
                    {
                        "topic": "task",
                        "task_id": task_id,
                        "event": "log",
                        "line": "[子 Agent 定时心跳] 仍有未完成子任务，继续协调",
                    }
                )
                continue
        except Exception as e:
            if slog:
                slog.write_exception(
                    "task",
                    "AgentLoop",
                    "sub_agent_worker",
                    e,
                    f"task_id={task_id}",
                )
            raise
        finally:
            current_task_id.reset(token)

    async def chat(self, user_text: str) -> str:
        slog = self._session_log
        if slog:
            u = user_text if len(user_text) <= 12000 else user_text[:11997] + "..."
            slog.write_event("agent", "AgentLoop", "chat_user", u, "")
        await self._store.append_chat_message(
            self._conn, session_id=self._session_id, role="user", content=user_text
        )
        self._messages.append({"role": "user", "content": user_text})
        await self._prepend_main_inbox_if_any()
        tools = openai_tools()
        try:
            for _ in range(self._cfg.agent.max_tool_rounds):
                await self._prepend_main_inbox_if_any()
                batch = [{"role": "system", "content": self._system_content()}, *self._messages]
                await self._bus.publish(
                    {
                        "topic": "agent",
                        "event": "phase",
                        "phase": "llm",
                        "detail": "等待大模型响应…",
                    }
                )
                resp = await self._llm.chat(batch, tools)
                self._log_llm_usage(resp)
                msg = resp.choices[0].message
                if not msg.tool_calls:
                    text = msg.content or ""
                    self._messages.append({"role": "assistant", "content": text})
                    await self._store.append_chat_message(
                        self._conn,
                        session_id=self._session_id,
                        role="assistant",
                        content=text,
                    )
                    if slog:
                        t2 = text if len(text) <= 12000 else text[:11997] + "..."
                        slog.write_event("agent", "AgentLoop", "chat_assistant", "", t2)
                    return text
                ad = _msg_assistant_dict(msg)
                self._messages.append(ad)
                await self._store.append_chat_message(
                    self._conn,
                    session_id=self._session_id,
                    role="assistant",
                    content=msg.content or "",
                    tool_calls=tool_calls_to_jsonable(msg.tool_calls),
                )
                for tc in msg.tool_calls:
                    name = tc.function.name
                    raw_args = tc.function.arguments or "{}"
                    await self._bus.publish(
                        {
                            "topic": "agent",
                            "event": "phase",
                            "phase": "tool",
                            "name": name,
                            "detail": f"执行工具 {name}…",
                        }
                    )
                    result = await self._router.run(name, raw_args)
                    await self._bus.publish(
                        {
                            "topic": "agent",
                            "event": "tool",
                            "name": name,
                            "preview": str(result)[:500],
                        }
                    )
                    body = json.dumps(result, ensure_ascii=False)[:24_000]
                    self._messages.append(
                        {
                            "role": "tool",
                            "tool_call_id": tc.id,
                            "content": body,
                        }
                    )
                    await self._store.append_chat_message(
                        self._conn,
                        session_id=self._session_id,
                        role="tool",
                        content=body,
                        tool_call_id=tc.id,
                        tool_name=name,
                    )
                if self._strip_chat_prefix_after_tool_round:
                    _trim_chat_messages_to_last_user_turn(self._messages)
                    self._strip_chat_prefix_after_tool_round = False
            if slog:
                slog.write_event(
                    "agent",
                    "AgentLoop",
                    "chat_stopped",
                    "",
                    "too_many_tool_rounds",
                )
            return "Stopped: too many tool rounds."
        except Exception as e:
            if slog:
                slog.write_exception("agent", "AgentLoop", "chat", e)
            raise


def merge_env_api_key(cfg: AppConfig) -> None:
    if cfg.llm.api_key:
        return
    proto = (cfg.llm.proto or "openai").strip().lower()
    if proto == "dashscope":
        cfg.llm.api_key = os.environ.get("DASHSCOPE_API_KEY", "") or os.environ.get(
            "OPENAI_API_KEY", ""
        )
    elif proto == "ollama":
        cfg.llm.api_key = (
            os.environ.get("OLLAMA_API_KEY", "")
            or os.environ.get("OPENAI_API_KEY", "")
            or "ollama"
        )
    else:
        cfg.llm.api_key = os.environ.get("OPENAI_API_KEY", "")
