from __future__ import annotations

import shutil
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any

from aicd.agent.tool_args_coerce import coerce_tool_json_array
from aicd.runtime_task_context import current_task_id
from aicd.tasks.main_inbox import (
    MAIN_INBOX_PREFIX,
    is_main_inbox_task_id,
    main_inbox_id,
)
from aicd.tasks.safe_ids import is_safe_task_id, validate_topic_slug

if TYPE_CHECKING:
    from aicd.agent.tool_router import ToolRouter

_ALLOWED_MSG_KINDS = frozenset(
    {
        "handoff",
        "status",
        "question",
        "artifact_index",
        "alert",  # 子任务向上级报告异常/风险（与 status 区分：更需尽快处理）
        "directive",  # 上级/主会话对用户新意图的转发，子 Agent 收件箱驱动继续执行
        "lifecycle",  # 运行时写入主收件箱：顶层子任务完成/失败/取消（也可由工具显式发送）
    }
)


def _parent_id_from_context() -> str | None:
    return current_task_id.get()


async def _resolve_thread_root_for_task_chain(
    router: ToolRouter, start_task_id: str
) -> str | None:
    """沿 parent 链查找 payload.thread_root_id（嵌套子任务对齐主会话）。"""
    seen: set[str] = set()
    tid = start_task_id.strip()
    while tid and tid not in seen:
        seen.add(tid)
        row = await router._tasks.get_task_row(tid)
        if row is None:
            break
        tr = row.payload.get("thread_root_id")
        if isinstance(tr, str) and tr.strip():
            return tr.strip()
        pid = row.parent_id
        tid = pid.strip() if isinstance(pid, str) and pid.strip() else ""
    return None


async def _task_shell(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    parent_id = _parent_id_from_context()
    thread_root_id: str | None = None
    if parent_id is None:
        sid = router._chat_session_id
        if isinstance(sid, str) and sid.strip():
            thread_root_id = sid.strip()
    try:
        tid = await router._tasks.create_shell_task(
            args["command"],
            parent_id=parent_id,
            cwd=args.get("cwd"),
            thread_root_id=thread_root_id,
        )
    except ValueError as e:
        return {"ok": False, "error": str(e)}
    return {
        "ok": True,
        "task_id": tid,
        "note": "Use task_result for payload.shell_process (root_pid of the shell wrapper). "
        "On Windows, task_cancel runs taskkill /T on that pid so GUI children exit too.",
    }


async def _task_dummy(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    parent_id = _parent_id_from_context()
    thread_root_id: str | None = None
    if parent_id is None:
        sid = router._chat_session_id
        if isinstance(sid, str) and sid.strip():
            thread_root_id = sid.strip()
    tid = await router._tasks.create_dummy_long_task(
        steps=int(args.get("steps", 5)),
        parent_id=parent_id,
        thread_root_id=thread_root_id,
    )
    return {"ok": True, "task_id": tid}


async def _task_ai(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    parent_id = _parent_id_from_context()
    thread_root_id: str | None = None
    if parent_id is None:
        sid = router._chat_session_id
        if isinstance(sid, str) and sid.strip():
            thread_root_id = sid.strip()

    ro = args.get("read_only_paths")
    ro_list: list[str] | None = None
    if isinstance(ro, list):
        ro_list = [str(x) for x in ro if isinstance(x, str) and x.strip()]
        if not ro_list:
            ro_list = None

    topic_slug = args.get("topic_slug")
    ts_raw = topic_slug if isinstance(topic_slug, str) else None
    try:
        ts = validate_topic_slug(ts_raw)
    except ValueError as e:
        return {"ok": False, "error": str(e)}

    ph = args.get("parent_handoff_expectation")
    ph_exp = ph if isinstance(ph, str) else None

    disp = args.get("display_title")
    disp_s = str(disp).strip()[:200] if isinstance(disp, str) and disp.strip() else None
    purp = args.get("purpose")
    purp_s = str(purp).strip()[:4000] if isinstance(purp, str) and purp.strip() else None
    est = args.get("estimated_steps")
    est_i: int | None = None
    if est is not None:
        try:
            est_i = int(est)
        except (TypeError, ValueError):
            est_i = None

    try:
        tid = await router._tasks.create_ai_task(
            args["instruction"],
            parent_id=parent_id,
            topic_slug=ts,
            read_only_paths=ro_list,
            parent_handoff_expectation=ph_exp,
            thread_root_id=thread_root_id,
            display_title=disp_s,
            purpose=purp_s,
            estimated_steps=est_i,
        )
    except ValueError as e:
        return {"ok": False, "error": str(e)}
    return {"ok": True, "task_id": tid}


def _task_timeline_payload(p: dict[str, Any]) -> dict[str, Any]:
    keys = (
        "started_at",
        "finished_at",
        "duration_sec",
        "progress",
        "sub_agent_trace",
        "display_title",
        "purpose",
        "estimated_steps",
        "shell_process",
    )
    return {k: p[k] for k in keys if k in p and p[k] is not None}


async def _task_list(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    rows = await router._tasks.list_tasks(limit=int(args.get("limit", 30)))
    tm = router._tasks
    return {
        "ok": True,
        "tasks": [
            {
                "id": r.id,
                "type": r.type,
                "status": r.status,
                "running_now": tm.is_task_running(r.id),
                "parent_id": r.parent_id,
                "created_at": r.created_at,
                "updated_at": r.updated_at,
                "summary": r.result_summary,
                "error": r.error,
                "timeline": _task_timeline_payload(r.payload),
            }
            for r in rows
        ],
    }


async def _task_cancel(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    return await router._tasks.cancel_task(args["task_id"])


async def _task_result(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    r = await router._tasks.get_task_row(args["task_id"])
    if not r:
        return {"ok": False, "error": "not found"}
    tm = router._tasks
    return {
        "ok": True,
        "task": {
            "id": r.id,
            "type": r.type,
            "status": r.status,
            "running_now": tm.is_task_running(r.id),
            "parent_id": r.parent_id,
            "created_at": r.created_at,
            "updated_at": r.updated_at,
            "payload": r.payload,
            "timeline": _task_timeline_payload(r.payload),
            "result_summary": r.result_summary,
            "error": r.error,
        },
    }


def _task_matches_thread(row: Any, thread_root_id: str) -> bool:
    if row.id == thread_root_id:
        return True
    p = row.payload or {}
    tr = p.get("thread_root_id")
    return isinstance(tr, str) and tr.strip() == thread_root_id


async def _task_progress(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    tid_arg = args.get("task_id")
    if isinstance(tid_arg, str) and tid_arg.strip():
        tid = tid_arg.strip()
    else:
        ctx = current_task_id.get()
        tid = ctx.strip() if isinstance(ctx, str) and ctx.strip() else ""
    if not tid:
        return {
            "ok": False,
            "error": "task_id required when not running inside a task_ai worker",
        }
    if not is_safe_task_id(tid):
        return {"ok": False, "error": "invalid task_id"}
    row = await router._tasks.get_task_row(tid)
    if not row:
        return {"ok": False, "error": "unknown task_id"}
    try:
        step_index = int(args["step_index"])
    except (KeyError, TypeError, ValueError):
        return {"ok": False, "error": "step_index (integer) is required"}
    step_total = args.get("step_total")
    st: int | None = None
    if step_total is not None:
        try:
            st = int(step_total)
        except (TypeError, ValueError):
            st = None
    cf = args.get("current_focus")
    cf_s = str(cf).strip()[:2000] if isinstance(cf, str) else ""
    note = args.get("note")
    note_s = str(note).strip()[:2000] if isinstance(note, str) else ""
    prog = {
        "step_index": step_index,
        "updated_at": datetime.now(timezone.utc).isoformat(),
    }
    if st is not None and st > 0:
        prog["step_total"] = st
    if cf_s:
        prog["current_focus"] = cf_s
    if note_s:
        prog["note"] = note_s
    await router._tasks.patch_task_payload(tid, {"progress": prog})
    await router._tasks._emit(
        tid,
        "progress",
        phase="explicit",
        step_index=step_index,
        step_total=st,
        current_focus=cf_s[:240],
        message=(
            f"进度 {step_index}"
            + (f"/{st}" if st else "")
            + (f" — {cf_s[:120]}" if cf_s else "")
        ),
    )
    return {"ok": True, "task_id": tid, "progress": prog}


async def _task_thread_stats(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    tr = args.get("thread_root_id")
    tr_s = tr.strip() if isinstance(tr, str) and tr.strip() else ""
    if not tr_s:
        sid = router._chat_session_id
        tr_s = sid.strip() if isinstance(sid, str) and sid.strip() else ""
    if not tr_s:
        return {
            "ok": False,
            "error": "thread_root_id required (or main chat session unavailable)",
        }
    rows = await router._tasks.list_tasks(limit=int(args.get("limit", 250)))
    matched = [r for r in rows if _task_matches_thread(r, tr_s)]
    running = sum(1 for r in matched if r.status == "running")
    completed = sum(1 for r in matched if r.status == "completed")
    failed = sum(1 for r in matched if r.status == "failed")
    cancelled = sum(1 for r in matched if r.status == "cancelled")
    sum_dur = 0.0
    for r in matched:
        ds = r.payload.get("duration_sec")
        if isinstance(ds, (int, float)):
            sum_dur += float(ds)
    earliest: str | None = None
    latest_finish: str | None = None
    for r in matched:
        sa = r.payload.get("started_at")
        if isinstance(sa, str) and sa.strip():
            if earliest is None or sa < earliest:
                earliest = sa
        fa = r.payload.get("finished_at")
        if isinstance(fa, str) and fa.strip():
            if latest_finish is None or fa > latest_finish:
                latest_finish = fa
    wall_sec: float | None = None
    if earliest and latest_finish:
        try:
            t0 = datetime.fromisoformat(earliest.replace("Z", "+00:00"))
            t1 = datetime.fromisoformat(latest_finish.replace("Z", "+00:00"))
            if t0.tzinfo is None:
                t0 = t0.replace(tzinfo=timezone.utc)
            if t1.tzinfo is None:
                t1 = t1.replace(tzinfo=timezone.utc)
            wall_sec = round((t1 - t0).total_seconds(), 3)
        except ValueError:
            wall_sec = None
    return {
        "ok": True,
        "thread_root_id": tr_s,
        "task_count": len(matched),
        "by_status": {
            "running": running,
            "completed": completed,
            "failed": failed,
            "cancelled": cancelled,
        },
        "sum_duration_sec": round(sum_dur, 3),
        "wall_span_sec": wall_sec,
        "earliest_started_at": earliest,
        "latest_finished_at": latest_finish,
    }


def _serialize_task_message(m: Any) -> dict[str, Any]:
    return {
        "id": m.id,
        "thread_root_id": m.thread_root_id,
        "to_task_id": m.to_task_id,
        "from_task_id": m.from_task_id,
        "kind": m.kind,
        "body": m.body,
        "created_at": m.created_at,
        "read_at": m.read_at,
    }


async def _task_message_send(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    to_raw = str(args["to_task_id"]).strip()
    kind = str(args["kind"]).strip()
    if kind not in _ALLOWED_MSG_KINDS:
        return {
            "ok": False,
            "error": f"kind must be one of: {sorted(_ALLOWED_MSG_KINDS)}",
        }
    body = args.get("body")
    if body is None:
        body_d: dict[str, Any] = {}
    elif isinstance(body, dict):
        body_d = body
    else:
        return {"ok": False, "error": "body must be a JSON object"}

    from_tid = args.get("from_task_id")
    if isinstance(from_tid, str) and from_tid.strip():
        fid = from_tid.strip()
    else:
        ctx = current_task_id.get()
        fid = ctx.strip() if isinstance(ctx, str) and ctx.strip() else ""
    if not fid:
        sid = router._chat_session_id
        if isinstance(sid, str) and sid.strip():
            fid = f"main-{sid.strip()}"
        else:
            return {
                "ok": False,
                "error": "from_task_id required when not in task_ai and no chat session",
            }

    resolved_to = to_raw
    if to_raw.lower() == "main":
        ctx2 = current_task_id.get()
        if ctx2 and str(ctx2).strip():
            trt = await _resolve_thread_root_for_task_chain(
                router, str(ctx2).strip()
            )
            if not trt:
                return {
                    "ok": False,
                    "error": "无法解析主会话：当前任务链上缺少 payload.thread_root_id",
                }
            resolved_to = main_inbox_id(trt)
        else:
            msid = router._chat_session_id
            if not (isinstance(msid, str) and msid.strip()):
                return {
                    "ok": False,
                    "error": "主对话发送 to_task_id=main 需要有效 chat session",
                }
            resolved_to = main_inbox_id(msid.strip())
    elif is_main_inbox_task_id(to_raw):
        resolved_to = to_raw

    if not is_safe_task_id(resolved_to) or not is_safe_task_id(fid):
        return {"ok": False, "error": "invalid task id format"}

    trow = await router._tasks.get_task_row(resolved_to)
    if not is_main_inbox_task_id(resolved_to):
        if not trow:
            return {"ok": False, "error": "unknown to_task_id"}

    tr = args.get("thread_root_id")
    thread_root: str | None = None
    if isinstance(tr, str) and tr.strip():
        thread_root = tr.strip()
    if thread_root is None:
        srow = await router._tasks.get_task_row(fid)
        if srow:
            tr2 = srow.payload.get("thread_root_id")
            if isinstance(tr2, str) and tr2.strip():
                thread_root = tr2.strip()
    if thread_root is None and isinstance(fid, str) and fid.startswith("main-"):
        rest = fid[5:].strip()
        if rest:
            thread_root = rest
    if thread_root is None and trow is not None:
        tr3 = trow.payload.get("thread_root_id")
        if isinstance(tr3, str) and tr3.strip():
            thread_root = tr3.strip()
    if thread_root is None and is_main_inbox_task_id(resolved_to):
        thread_root = resolved_to[len(MAIN_INBOX_PREFIX) :]

    mid = await router._tasks.post_task_message(
        to_task_id=resolved_to,
        from_task_id=fid,
        kind=kind,
        body=body_d,
        thread_root_id=thread_root,
    )
    return {
        "ok": True,
        "message_id": mid,
        "resolved_to_task_id": resolved_to,
    }


async def _task_message_inbox(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    to_arg = args.get("to_task_id")
    if isinstance(to_arg, str) and to_arg.strip():
        to_tid = to_arg.strip()
    else:
        ctx = current_task_id.get()
        to_tid = ctx.strip() if isinstance(ctx, str) and ctx.strip() else ""
    if not to_tid:
        return {
            "ok": False,
            "error": "to_task_id required when not running inside a task_ai worker",
        }
    if not is_safe_task_id(to_tid):
        return {"ok": False, "error": "invalid to_task_id"}
    since = int(args.get("since_id", 0))
    limit = min(int(args.get("limit", 50)), 200)
    rows = await router._tasks.store.list_task_messages(
        router._tasks.db_conn,
        to_task_id=to_tid,
        since_id=since,
        limit=limit,
    )
    return {
        "ok": True,
        "to_task_id": to_tid,
        "messages": [_serialize_task_message(m) for m in rows],
    }


async def _task_message_mark_read(
    router: ToolRouter, args: dict[str, Any]
) -> dict[str, Any]:
    raw = args.get("message_ids")
    co = coerce_tool_json_array(raw, "message_ids", allow_empty=False)
    if isinstance(co, dict):
        return co
    raw = co
    ids: list[int] = []
    for x in raw:
        try:
            ids.append(int(x))
        except (TypeError, ValueError):
            return {"ok": False, "error": "message_ids must be integers"}
    await router._tasks.store.mark_task_messages_read(router._tasks.db_conn, ids)
    return {"ok": True, "marked": len(ids)}


async def _task_main_inbox(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    """主对话：拉取 main-inbox-<session> 中的未读/增量 task_message。"""
    if current_task_id.get() is not None:
        return {
            "ok": False,
            "error": "task_main_inbox 仅可在主对话使用；子 Agent 请用 task_message_send(..., to_task_id=\"main\", ...)",
        }
    sid = router._chat_session_id
    if not (isinstance(sid, str) and sid.strip()):
        return {"ok": False, "error": "无 chat session，无法解析主收件箱"}
    to_id = main_inbox_id(sid.strip())
    since = int(args.get("since_message_id", 0) or 0)
    limit = min(max(int(args.get("limit", 50) or 50), 1), 200)
    rows = await router._tasks.store.list_task_messages(
        router._tasks.db_conn,
        to_task_id=to_id,
        since_id=since,
        limit=limit,
    )
    mark = args.get("mark_read")
    do_mark = True if mark is None else bool(mark)
    if do_mark and rows:
        await router._tasks.store.mark_task_messages_read(
            router._tasks.db_conn, [m.id for m in rows]
        )
    return {
        "ok": True,
        "main_inbox_id": to_id,
        "since_message_id": since,
        "messages": [_serialize_task_message(m) for m in rows],
        "note": "子任务结束时会自动写入 kind=lifecycle；嵌套子任务上行请 to_task_id=\"main\"。"
        " 定时检查周期用 **agent.idleHeartbeatSec**（TUI `/agent heartbeat`）或本工具主动拉取。",
    }


async def _task_session_children_digest(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    """主对话：列出本会话发起的顶层后台任务（parent_id 为空且 thread_root_id 匹配）。"""
    if current_task_id.get() is not None:
        return {
            "ok": False,
            "error": "task_session_children_digest 仅可在主对话使用",
        }
    sid = router._chat_session_id
    if not (isinstance(sid, str) and sid.strip()):
        return {"ok": False, "error": "无 chat session"}
    since_u = args.get("since_updated_at")
    since_s = str(since_u).strip() if since_u is not None and str(since_u).strip() else None
    limit = min(max(int(args.get("limit", 40) or 40), 1), 200)
    rows = await router._tasks.store.list_top_level_tasks_for_session_thread(
        router._tasks.db_conn,
        sid.strip(),
        limit=limit,
        since_updated_at=since_s,
    )
    return {
        "ok": True,
        "session_thread_id": sid.strip(),
        "since_updated_at": since_s,
        "tasks": [
            {
                "id": r.id,
                "type": r.type,
                "status": r.status,
                "running_now": router._tasks.is_task_running(r.id),
                "updated_at": r.updated_at,
                "summary": r.result_summary,
                "error": r.error,
                "display_title": (r.payload or {}).get("display_title"),
            }
            for r in rows
        ],
        "note": "仅包含带 payload.thread_root_id 的顶层任务（主对话发起的 task_ai/task_shell 等）。"
        " 若 since_updated_at 为 ISO 时间字符串，只返回 updated_at 更晚的行。",
    }


async def _task_history_clear(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    if current_task_id.get() is not None:
        return {
            "ok": False,
            "error": "task_history_clear 仅可在主对话中使用，不能在 task_ai 子 Agent 内调用",
        }
    if args.get("confirm") is not True:
        return {
            "ok": False,
            "error": "须设置 confirm=true 才会清空：SQLite 中全部 tasks/task_messages、"
            "AICD_HOME/tasks/<id> 工作区目录，以及当前 AI tmp/tasks/* 草稿目录。",
        }
    return await router._tasks.clear_task_history()


async def _task_scratch_cleanup(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    tid_arg = args.get("task_id")
    if isinstance(tid_arg, str) and tid_arg.strip():
        tid = tid_arg.strip()
    else:
        ctx = current_task_id.get()
        tid = ctx.strip() if isinstance(ctx, str) and ctx.strip() else ""
    if not tid:
        return {
            "ok": False,
            "error": "task_id required when not running inside a task_ai worker",
        }
    if not is_safe_task_id(tid):
        return {"ok": False, "error": "invalid task_id"}

    _out, ai_tmp = router.get_ai_paths()
    base = (ai_tmp / "tasks").resolve()
    root = (base / tid).resolve()
    try:
        root.relative_to(base)
    except ValueError:
        return {"ok": False, "error": "scratch path escaped allowed directory"}
    if not root.is_dir():
        root.mkdir(parents=True, exist_ok=True)
        return {"ok": True, "cleaned": str(root), "note": "was missing; created empty"}
    shutil.rmtree(root)
    root.mkdir(parents=True, exist_ok=True)
    return {"ok": True, "cleaned": str(root)}


HANDLERS: dict[str, Any] = {
    "task_shell": _task_shell,
    "task_dummy": _task_dummy,
    "task_ai": _task_ai,
    "task_list": _task_list,
    "task_cancel": _task_cancel,
    "task_result": _task_result,
    "task_progress": _task_progress,
    "task_thread_stats": _task_thread_stats,
    "task_history_clear": _task_history_clear,
    "task_message_send": _task_message_send,
    "task_message_inbox": _task_message_inbox,
    "task_message_mark_read": _task_message_mark_read,
    "task_main_inbox": _task_main_inbox,
    "task_session_children_digest": _task_session_children_digest,
    "task_scratch_cleanup": _task_scratch_cleanup,
}
