from __future__ import annotations

import asyncio
import json
import shutil
import sys
from collections.abc import Awaitable, Callable
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING, Any

import aiosqlite

from aicd.bus.event_bus import AsyncEventBus
from aicd.db.store import TaskStore
from aicd.tasks.main_inbox import main_inbox_id
from aicd.tasks.safe_ids import validate_topic_slug
from aicd.tools.process_tool import (
    ProcessTool,
    default_shell_argv,
    warn_shell_preflight,
)

if TYPE_CHECKING:
    from aicd.session_log import SessionLogger

# 取消后台任务时 await 单任务的最长等待（避免 task_shell / GUI 子进程卡住导致 /exit 永远不完）
_CANCEL_TASK_AWAIT_SEC = 14.0


def _utc_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _duration_since_started_iso(started_at: str | None) -> float | None:
    if not started_at or not isinstance(started_at, str):
        return None
    try:
        s = started_at.replace("Z", "+00:00")
        t0 = datetime.fromisoformat(s)
        if t0.tzinfo is None:
            t0 = t0.replace(tzinfo=timezone.utc)
        return round((datetime.now(timezone.utc) - t0).total_seconds(), 3)
    except ValueError:
        return None


class TaskManager:
    def __init__(
        self,
        *,
        store: TaskStore,
        conn: aiosqlite.Connection,
        bus: AsyncEventBus,
        tasks_root: Path,
        process_tool: ProcessTool,
        session_log: SessionLogger | None = None,
    ) -> None:
        self._store = store
        self._conn = conn
        self._bus = bus
        self._tasks_root = tasks_root
        self._process = process_tool
        self._session_log = session_log
        self._running: dict[str, asyncio.Task[None]] = {}
        self._ai_runner: Callable[[str, str], Awaitable[None]] | None = None
        self._ai_paths_resolver: Callable[[], tuple[Path, Path]] | None = None

    @property
    def bus(self) -> AsyncEventBus:
        return self._bus

    @property
    def db_conn(self) -> aiosqlite.Connection:
        return self._conn

    @property
    def store(self) -> TaskStore:
        return self._store

    def set_ai_runner(
        self, runner: Callable[[str, str], Awaitable[None]] | None
    ) -> None:
        self._ai_runner = runner

    def set_ai_paths_resolver(
        self, fn: Callable[[], tuple[Path, Path]] | None
    ) -> None:
        """Return (ai_output_root, ai_tmp_root) for scratch/deliverable paths on disk."""
        self._ai_paths_resolver = fn

    def _resolve_ai_paths(self) -> tuple[Path, Path]:
        if self._ai_paths_resolver is not None:
            o, t = self._ai_paths_resolver()
            return o.resolve(), t.resolve()
        base = (Path.cwd() / "aicd").resolve()
        return base, (base / "tmp").resolve()

    def set_process_tool(self, process_tool: ProcessTool) -> None:
        """与主 AI 工作目录对齐时，替换用于 task_shell 等任务的 ProcessTool。"""
        self._process = process_tool

    async def patch_task_payload(
        self, task_id: str, patch: dict[str, Any]
    ) -> None:
        """合并写入 ``tasks.payload_json``（``progress`` 子对象与已有键深度合并）。"""
        row = await self._store.get_task(self._conn, task_id)
        if not row:
            return
        pl = dict(row.payload)
        for k, v in patch.items():
            if k == "progress" and isinstance(v, dict):
                prev = pl.get("progress")
                if isinstance(prev, dict):
                    merged = dict(prev)
                    merged.update(v)
                    pl[k] = merged
                else:
                    pl[k] = dict(v)
            elif k == "shell_process" and isinstance(v, dict):
                prev = pl.get("shell_process")
                if isinstance(prev, dict):
                    merged = dict(prev)
                    merged.update(v)
                    pl[k] = merged
                else:
                    pl[k] = dict(v)
            else:
                pl[k] = v
        await self._store.update_task(self._conn, task_id, payload=pl)

    async def _stamp_started(self, task_id: str) -> None:
        row = await self._store.get_task(self._conn, task_id)
        if not row:
            return
        pl = dict(row.payload)
        if pl.get("started_at"):
            return
        pl["started_at"] = _utc_iso()
        await self._store.update_task(self._conn, task_id, payload=pl)

    async def _stamp_finished(self, task_id: str) -> None:
        row = await self._store.get_task(self._conn, task_id)
        if not row:
            return
        pl = dict(row.payload)
        if pl.get("finished_at"):
            return
        pl["finished_at"] = _utc_iso()
        d = _duration_since_started_iso(pl.get("started_at"))
        if d is not None:
            pl["duration_sec"] = d
        await self._store.update_task(self._conn, task_id, payload=pl)

    async def create_ai_task(
        self,
        instruction: str,
        *,
        parent_id: str | None = None,
        topic_slug: str | None = None,
        read_only_paths: list[str] | None = None,
        parent_handoff_expectation: str | None = None,
        thread_root_id: str | None = None,
        display_title: str | None = None,
        purpose: str | None = None,
        estimated_steps: int | None = None,
    ) -> str:
        if self._ai_runner is None:
            raise RuntimeError("AI runner not configured")
        slog = self._session_log
        slug = validate_topic_slug(topic_slug)
        ai_out, ai_tmp = self._resolve_ai_paths()
        if parent_id and thread_root_id is None:
            prow = await self._store.get_task(self._conn, parent_id)
            if prow:
                tr = prow.payload.get("thread_root_id")
                if isinstance(tr, str) and tr.strip():
                    thread_root_id = tr.strip()
        ro_list: list[str] = []
        if read_only_paths:
            for p in read_only_paths:
                if isinstance(p, str) and p.strip():
                    ro_list.append(p.strip())
        payload: dict[str, Any] = {"instruction": instruction}
        first_line = instruction.strip().split("\n")[0].strip()
        title = (display_title or "").strip() or (first_line[:200] if first_line else "ai task")
        payload["display_title"] = title[:200]
        if purpose and str(purpose).strip():
            payload["purpose"] = str(purpose).strip()[:4000]
        if estimated_steps is not None:
            try:
                es = int(estimated_steps)
                if 1 <= es <= 10_000:
                    payload["estimated_steps"] = es
            except (TypeError, ValueError):
                pass
        if slug:
            ddir = (ai_out / "deliverables" / slug).resolve()
            ddir.mkdir(parents=True, exist_ok=True)
            payload["topic_slug"] = slug
            payload["deliverable_dir"] = str(ddir)
        if ro_list:
            payload["read_only_paths"] = ro_list
        if parent_handoff_expectation and str(parent_handoff_expectation).strip():
            payload["parent_handoff_expectation"] = str(parent_handoff_expectation).strip()[
                :4000
            ]
        if thread_root_id and thread_root_id.strip():
            payload["thread_root_id"] = thread_root_id.strip()

        tid = await self._store.insert_task(
            self._conn,
            task_type="ai_agent",
            status="pending",
            parent_id=parent_id,
            payload=payload,
        )
        scratch_path = (ai_tmp / "tasks" / tid).resolve()
        scratch_path.mkdir(parents=True, exist_ok=True)
        payload["scratch_dir"] = str(scratch_path)
        await self._store.update_task(self._conn, tid, payload=payload)

        ws = self._tasks_root / tid / "workspace"
        ws.mkdir(parents=True, exist_ok=True)
        t = asyncio.create_task(self._run_ai(tid, instruction))
        self._running[tid] = t
        if slog:
            ins = instruction if len(instruction) <= 6000 else instruction[:5997] + "..."
            slog.write_event(
                "task",
                "TaskManager",
                "create_ai_task",
                ins,
                f"task_id={tid}",
            )
        return tid

    async def _run_ai(self, task_id: str, instruction: str) -> None:
        assert self._ai_runner is not None
        await self._store.update_task(self._conn, task_id, status="running")
        await self._stamp_started(task_id)
        await self._emit(task_id, "status", status="running")
        slog = self._session_log
        try:
            await self._ai_runner(task_id, instruction)
            await self._stamp_finished(task_id)
            await self._store.update_task(
                self._conn,
                task_id,
                status="completed",
                result_summary="ai task finished",
            )
            await self._emit(task_id, "status", status="completed")
            await self.notify_main_inbox_top_level(task_id)
            if slog:
                slog.write_event(
                    "task",
                    "TaskManager",
                    "ai_task_done",
                    f"task_id={task_id}",
                    "status=completed",
                )
        except asyncio.CancelledError:
            await self._stamp_finished(task_id)
            await self._store.update_task(
                self._conn, task_id, status="cancelled", error="cancelled"
            )
            await self._emit(task_id, "status", status="cancelled")
            await self.notify_main_inbox_top_level(task_id)
            if slog:
                slog.write_event(
                    "task",
                    "TaskManager",
                    "ai_task_done",
                    f"task_id={task_id}",
                    "status=cancelled",
                )
            raise
        except Exception as e:  # noqa: BLE001
            await self._stamp_finished(task_id)
            await self._store.update_task(
                self._conn, task_id, status="failed", error=str(e)
            )
            await self._emit(task_id, "status", status="failed", error=str(e))
            await self.notify_main_inbox_top_level(task_id)
            if slog:
                slog.write_exception(
                    "task",
                    "TaskManager",
                    "ai_task_failed",
                    e,
                    f"task_id={task_id}",
                )
        finally:
            # 子 Agent 内每次 run_command 均在本协程内 await 结束；失败/取消时无长期占用子进程句柄。
            self._running.pop(task_id, None)

    async def _emit(self, task_id: str, event: str, **kwargs: Any) -> None:
        await self._bus.publish(
            {"topic": "task", "task_id": task_id, "event": event, **kwargs}
        )

    async def notify_main_inbox_top_level(self, task_id: str) -> None:
        """顶层任务进入终态后向主会话虚拟收件箱投递 lifecycle（payload.thread_root_id）。"""
        row = await self._store.get_task(self._conn, task_id)
        if not row or row.parent_id is not None:
            return
        tr = row.payload.get("thread_root_id")
        if not isinstance(tr, str) or not tr.strip():
            return
        sid = tr.strip()
        to_id = main_inbox_id(sid)
        title = row.payload.get("display_title")
        title_s = str(title).strip()[:200] if title else None
        parts = [
            f"子任务 `{task_id}`（{row.type}）状态：**{row.status}**",
        ]
        if title_s:
            parts.append(f"标题：{title_s}")
        if row.result_summary:
            parts.append(f"摘要：{row.result_summary}")
        if row.error:
            parts.append(f"错误：{row.error}")
        body: dict[str, Any] = {
            "phase": row.status,
            "child_task_id": task_id,
            "task_type": row.type,
            "result_summary": row.result_summary,
            "error": row.error,
            "text": "。".join(parts) + "。",
        }
        if title_s:
            body["display_title"] = title_s
        try:
            await self.post_task_message(
                to_task_id=to_id,
                from_task_id=task_id,
                kind="lifecycle",
                body=body,
                thread_root_id=sid,
            )
        except Exception as e:  # noqa: BLE001
            slog = self._session_log
            if slog:
                slog.write_exception(
                    "task",
                    "TaskManager",
                    "notify_main_inbox_top_level",
                    e,
                    f"task_id={task_id}",
                )

    async def post_task_message(
        self,
        *,
        to_task_id: str,
        from_task_id: str,
        kind: str,
        body: dict[str, Any],
        thread_root_id: str | None = None,
    ) -> int:
        mid = await self._store.insert_task_message(
            self._conn,
            to_task_id=to_task_id,
            from_task_id=from_task_id,
            kind=kind,
            body=body,
            thread_root_id=thread_root_id,
        )
        await self._bus.publish(
            {
                "topic": "task",
                "task_id": to_task_id,
                "event": "task_message",
                "message_id": mid,
                "from_task_id": from_task_id,
                "kind": kind,
            }
        )
        return mid

    async def create_shell_task(
        self,
        command: str,
        parent_id: str | None = None,
        cwd: str | None = None,
        *,
        thread_root_id: str | None = None,
    ) -> str:
        if (msg := warn_shell_preflight(command)):
            raise ValueError(msg)
        tr: str | None = thread_root_id
        if parent_id and not (tr and str(tr).strip()):
            prow = await self._store.get_task(self._conn, parent_id)
            if prow:
                t0 = prow.payload.get("thread_root_id")
                if isinstance(t0, str) and t0.strip():
                    tr = t0.strip()
        payload: dict[str, Any] = {"command": command, "cwd": cwd}
        if tr and str(tr).strip():
            payload["thread_root_id"] = str(tr).strip()
        tid = await self._store.insert_task(
            self._conn,
            task_type="shell",
            status="pending",
            parent_id=parent_id,
            payload=payload,
        )
        workspace = self._tasks_root / tid / "workspace"
        workspace.mkdir(parents=True, exist_ok=True)
        t = asyncio.create_task(self._run_shell(tid, command, cwd))
        self._running[tid] = t
        if self._session_log:
            self._session_log.write_event(
                "task",
                "TaskManager",
                "create_shell_task",
                json.dumps(
                    {"command": command[:4000], "cwd": cwd},
                    ensure_ascii=False,
                ),
                f"task_id={tid}",
            )
        return tid

    async def _run_shell(self, task_id: str, command: str, cwd: str | None) -> None:
        await self._store.update_task(
            self._conn, task_id, status="running", result_summary=None, error=None
        )
        await self._stamp_started(task_id)
        await self._emit(task_id, "status", status="running")

        async def on_line(line: str) -> None:
            await self._emit(task_id, "log", line=line)

        async def on_spawned(root_pid: int) -> None:
            await self.patch_task_payload(
                task_id,
                {
                    "shell_process": {
                        "root_pid": root_pid,
                        "platform": sys.platform,
                        "role": "asyncio subprocess root (e.g. cmd.exe / pwsh); "
                        "GUI or child CLIs started under this shell are descendants",
                    }
                },
            )
            await self._emit(task_id, "shell_process", root_pid=root_pid)

        argv = default_shell_argv(command)
        slog = self._session_log
        try:
            result = await self._process.run_command(
                argv,
                cwd=cwd,
                timeout_sec=3600,
                on_line=on_line,
                on_spawned=on_spawned,
                kill_subtree_on_cancel=(sys.platform == "win32"),
            )
        except asyncio.CancelledError:
            await self.patch_task_payload(
                task_id,
                {
                    "shell_process": {
                        "cancelled": True,
                        "windows_subtree_kill_attempted": sys.platform == "win32",
                    }
                },
            )
            await self._stamp_finished(task_id)
            await self._store.update_task(
                self._conn,
                task_id,
                status="cancelled",
                error="user cancel",
                result_summary=None,
            )
            await self._emit(task_id, "status", status="cancelled")
            await self.notify_main_inbox_top_level(task_id)
            if slog:
                slog.write_event(
                    "task",
                    "TaskManager",
                    "shell_task_done",
                    f"task_id={task_id}",
                    "status=cancelled",
                )
            self._running.pop(task_id, None)
            raise
        except Exception as e:  # noqa: BLE001
            await self._stamp_finished(task_id)
            await self._store.update_task(
                self._conn,
                task_id,
                status="failed",
                error=str(e),
                result_summary="exception",
            )
            await self._emit(task_id, "status", status="failed", error=str(e))
            await self.notify_main_inbox_top_level(task_id)
            if slog:
                slog.write_exception(
                    "task",
                    "TaskManager",
                    "shell_task",
                    e,
                    f"task_id={task_id}",
                )
            self._running.pop(task_id, None)
            return
        await self.patch_task_payload(
            task_id,
            {
                "shell_process": {
                    "finished": True,
                    "returncode": result.get("returncode"),
                    "pid": result.get("pid"),
                }
            },
        )
        summary = json.dumps(
            {
                "returncode": result.get("returncode"),
                "ok": result.get("ok"),
                "pid": result.get("pid"),
            },
            ensure_ascii=False,
        )
        st = "completed" if result.get("ok") else "failed"
        await self._stamp_finished(task_id)
        await self._store.update_task(
            self._conn,
            task_id,
            status=st,
            result_summary=summary[:4000],
            error=None if result.get("ok") else (result.get("stderr") or "")[:2000],
        )
        await self._emit(task_id, "status", status=st, result=result)
        await self.notify_main_inbox_top_level(task_id)
        if slog:
            out = summary
            if isinstance(result, dict):
                out = json.dumps(
                    {
                        "ok": result.get("ok"),
                        "returncode": result.get("returncode"),
                        "stdout_chars": len(str(result.get("stdout", ""))),
                        "stderr_chars": len(str(result.get("stderr", ""))),
                    },
                    ensure_ascii=False,
                )
            slog.write_event(
                "task",
                "TaskManager",
                "shell_task_done",
                f"task_id={task_id}",
                out,
            )
        self._running.pop(task_id, None)

    async def create_dummy_long_task(
        self,
        steps: int = 5,
        parent_id: str | None = None,
        *,
        thread_root_id: str | None = None,
    ) -> str:
        tr: str | None = thread_root_id
        if parent_id and not (tr and str(tr).strip()):
            prow = await self._store.get_task(self._conn, parent_id)
            if prow:
                t0 = prow.payload.get("thread_root_id")
                if isinstance(t0, str) and t0.strip():
                    tr = t0.strip()
        payload: dict[str, Any] = {"steps": steps}
        if tr and str(tr).strip():
            payload["thread_root_id"] = str(tr).strip()
        tid = await self._store.insert_task(
            self._conn,
            task_type="dummy_long",
            status="pending",
            parent_id=parent_id,
            payload=payload,
        )
        t = asyncio.create_task(self._run_dummy(tid, steps))
        self._running[tid] = t
        if self._session_log:
            self._session_log.write_event(
                "task",
                "TaskManager",
                "create_dummy_task",
                f"steps={steps}",
                f"task_id={tid}",
            )
        return tid

    async def _run_dummy(self, task_id: str, steps: int) -> None:
        await self._store.update_task(self._conn, task_id, status="running")
        await self._stamp_started(task_id)
        await self._emit(task_id, "status", status="running")
        try:
            for i in range(steps):
                await asyncio.sleep(0.8)
                pct = int((i + 1) / steps * 100)
                await self._emit(
                    task_id, "progress", percent=pct, message=f"step {i+1}/{steps}"
                )
            await self._stamp_finished(task_id)
            await self._store.update_task(
                self._conn,
                task_id,
                status="completed",
                result_summary="dummy completed",
            )
            await self._emit(task_id, "status", status="completed")
            await self.notify_main_inbox_top_level(task_id)
            if self._session_log:
                self._session_log.write_event(
                    "task",
                    "TaskManager",
                    "dummy_task_done",
                    f"task_id={task_id}",
                    "completed",
                )
        except asyncio.CancelledError:
            await self._stamp_finished(task_id)
            await self._store.update_task(
                self._conn, task_id, status="cancelled", error="cancelled"
            )
            await self._emit(task_id, "status", status="cancelled")
            await self.notify_main_inbox_top_level(task_id)
            if self._session_log:
                self._session_log.write_event(
                    "task",
                    "TaskManager",
                    "dummy_task_done",
                    f"task_id={task_id}",
                    "cancelled",
                )
            raise
        finally:
            self._running.pop(task_id, None)

    async def cancel_all_running(self, *, reason: str = "safe_exit") -> None:
        """取消所有仍在运行的后台任务（用于安全退出）。并行取消，避免多个 shell 任务串行拖死退出。"""
        tids = list(self._running.keys())
        if not tids:
            return
        await asyncio.gather(
            *(self.cancel_task(tid, reason=reason) for tid in tids),
            return_exceptions=True,
        )

    def abandon_stale_running_tasks(self) -> list[str]:
        """
        ``shutdown_graceful`` 在 ``cancel_all_running`` 整体超时后调用：不等待协程结束，
        仅从注册表摘除并 ``cancel()``，保证能继续关库。
        """
        leaked: list[str] = []
        for tid, t in list(self._running.items()):
            self._running.pop(tid, None)
            leaked.append(tid)
            if t is not None and not t.done():
                t.cancel()
        return leaked

    async def cancel_task(
        self, task_id: str, *, reason: str = "user cancel"
    ) -> dict[str, Any]:
        t = self._running.get(task_id)
        if t and not t.done():
            t.cancel()
            try:
                await asyncio.wait_for(t, timeout=_CANCEL_TASK_AWAIT_SEC)
            except asyncio.TimeoutError:
                self._running.pop(task_id, None)
                if self._session_log:
                    self._session_log.write_event(
                        "task",
                        "TaskManager",
                        "cancel_task_await_timeout",
                        f"task_id={task_id}",
                        f"timeout_sec={_CANCEL_TASK_AWAIT_SEC}",
                    )
            except asyncio.CancelledError:
                pass
            try:
                await self._stamp_finished(task_id)
                await self._store.update_task(
                    self._conn, task_id, status="cancelled", error=reason
                )
                await self._emit(task_id, "status", status="cancelled")
                await self._notify_main_inbox_top_level_safe(task_id)
            except Exception:
                pass
            if self._session_log:
                self._session_log.write_event(
                    "task",
                    "TaskManager",
                    "cancel_task",
                    f"task_id={task_id}",
                    "ok=true",
                )
            return {"ok": True, "task_id": task_id}
        row = await self._store.get_task(self._conn, task_id)
        if not row:
            return {"ok": False, "error": "unknown task"}
        return {"ok": False, "error": "task not running", "status": row.status}

    async def _notify_main_inbox_top_level_safe(self, task_id: str) -> None:
        """退出阶段 DB 可能已忙/关闭，避免因收件箱通知失败阻断 cancel。"""
        try:
            await self.notify_main_inbox_top_level(task_id)
        except Exception:
            pass

    async def get_task_row(self, task_id: str) -> Any:
        return await self._store.get_task(self._conn, task_id)

    def is_task_running(self, task_id: str) -> bool:
        t = self._running.get(task_id)
        return t is not None and not t.done()

    async def list_tasks(self, limit: int = 50) -> list[Any]:
        return await self._store.list_tasks(self._conn, limit=limit)

    async def clear_task_history(self) -> dict[str, Any]:
        """取消运行中任务，清空 SQLite 任务表与消息表，并删除 HOME/tasks/* 与 AI tmp/tasks/* 目录。"""
        await self.cancel_all_running(reason="task_history_clear")
        counts = await self._store.clear_task_history(self._conn)
        removed_ws = 0
        root = self._tasks_root
        if root.is_dir():
            for child in list(root.iterdir()):
                if child.is_dir():
                    shutil.rmtree(child, ignore_errors=True)
                    removed_ws += 1
        _, ai_tmp = self._resolve_ai_paths()
        scratch_root = (ai_tmp / "tasks").resolve()
        removed_scr = 0
        if scratch_root.is_dir():
            for child in list(scratch_root.iterdir()):
                if child.is_dir():
                    shutil.rmtree(child, ignore_errors=True)
                    removed_scr += 1
        slog = self._session_log
        if slog:
            slog.write_event(
                "task",
                "TaskManager",
                "clear_task_history",
                json.dumps(counts, ensure_ascii=False),
                f"workspace_dirs_removed={removed_ws} scratch_dirs_removed={removed_scr}",
            )
        # TUI 右侧任务表仅在收到 task 总线事件时刷新；主对话 task_history_clear 须显式通知
        await self._bus.publish(
            {
                "topic": "task",
                "task_id": "",
                "event": "history_cleared",
            }
        )
        return {
            "ok": True,
            **counts,
            "workspace_dirs_removed": removed_ws,
            "scratch_dirs_removed": removed_scr,
        }
