# tools/subagent_tools.py — Delegate tasks to isolated subagents with restricted tools and no personality.
# Supports single-task delegation and batch parallel execution with depth limiting.
from __future__ import annotations

import contextvars
import json
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any
from uuid import uuid4

from langchain_core.tools import tool
from pydantic import BaseModel, Field, field_validator
from tools._logging import _debug_log, _emit_input_status

_MAX_DELEGATE_DEPTH = 2
_BATCH_MAX_WORKERS = 3


class BatchTaskItem(BaseModel):
    task_description: str = Field(description="Clear, specific task for the subagent.")
    agent_type: str = Field(default="general", description="Subagent profile for this task.")
    context_snippet: str = Field(default="", description="Optional brief context.")


class DelegateTaskInput(BaseModel):
    task_description: str = Field(description="Clear, specific task for the subagent to complete.")
    agent_type: str = Field(
        default="research",
        description=(
            "Subagent profile: 'research' (full research pipeline), "
            "'web_research_light' (quick web lookup), "
            "'code_readonly' (read/search files), "
            "'extraction' (PDF/image/web data extraction), "
            "'bash' (shell commands), "
            "'ops_cautious' (careful file/shell ops), "
            "or 'general' (all tools)."
        ),
    )
    context_snippet: str = Field(
        default="",
        description="Optional brief context from the current conversation to help the subagent.",
    )
    handoff_envelope_json: str = Field(
        default="",
        description=(
            "Optional JSON for HandoffEnvelope hints at the delegation seam (single-task only in v1; "
            "batch tasks ignore this field). The tool result begins with delegation_handoff: lines "
            "(job id, blackboard/envelope telemetry) then the subagent output."
        ),
    )
    tasks: str | list[Any] = Field(
        default="",
        description=(
            "Optional: JSON array of batch tasks to run in parallel. Each element: "
            '{"task_description": "...", "agent_type": "...", "context_snippet": "..."}. '
            "Pass as a native JSON array in the tool call, or as a string containing JSON. "
            "When non-empty and valid, the top-level task_description is ignored and each batch task runs concurrently."
        ),
    )
    delegation_job_id: str = Field(
        default="",
        description=(
            "Optional delegation scratch job id. When omitted, a random id is generated. "
            "Correlates delegation scratch blackboard inject when handoff_blackboard_auto is on in mugi_config "
            "[subagent]. Use this to correlate related delegated subtasks."
        ),
    )

    @field_validator("tasks", mode="before")
    @classmethod
    def _coerce_tasks(cls, v: Any) -> str | list[Any]:
        if v is None:
            return ""
        return v


# Tools that child subagents must not have (prevents recursion and unsafe mutations).
def _envelope_json_status(handoff_envelope_json: str, handoff_envelope: object | None) -> str:
    if not (handoff_envelope_json or "").strip():
        return "empty"
    if handoff_envelope is None:
        return "invalid"
    return "ok"


def format_delegate_task_result(result, *, envelope_json_status: str) -> str:
    """Parent-visible digest (no paths or envelope body) plus standard subagent text."""
    job_display = (result.delegation_job_id or "").strip() or "-"
    envelope_prompt = "none"
    if envelope_json_status == "ok" and result.envelope_prompt_injected:
        envelope_prompt = "injected"
    lines = [
        f"delegation_handoff: job_id={job_display}",
        f"delegation_handoff: blackboard_chars={result.blackboard_injected_chars}",
        f"delegation_handoff: blackboard_skip={result.blackboard_skip_reason}",
        f"delegation_handoff: envelope_parse={envelope_json_status}",
        f"delegation_handoff: envelope_prompt={envelope_prompt}",
    ]
    return "\n".join(lines) + "\n---\n\n" + result.to_string()


def _batch_handoff_digest_block(result) -> str:
    return (
        "delegation_handoff: handoff_envelope=v1_not_applicable\n"
        f"delegation_handoff: blackboard_chars={result.blackboard_injected_chars}\n"
        f"delegation_handoff: blackboard_skip={result.blackboard_skip_reason}"
    )


def _strip_markdown_fences(raw: str) -> str:
    t = raw.strip()
    if not t.startswith("```"):
        return t
    lines = t.split("\n")
    if lines and lines[0].strip().startswith("```"):
        lines = lines[1:]
    if lines and lines[-1].strip() == "```":
        lines = lines[:-1]
    return "\n".join(lines).strip()


def _loads_json_maybe_double(s: str) -> Any:
    """Parse JSON; if the top-level value is a string, parse again (double-encoded payloads)."""
    val = json.loads(s)
    if isinstance(val, str):
        val = json.loads(val.strip())
    return val


def _parse_delegate_batch_tasks(tasks: str | list[Any] | None) -> tuple[list[BatchTaskItem], str | None]:
    """Parse batch `tasks` for delegate_task.

    Returns (items, err). *err* is set when the caller clearly intended a batch (non-empty
    ``tasks``) but no valid subtasks could be built — callers should return *err* to the model
    instead of silently falling back to single-task mode.
    """
    if isinstance(tasks, tuple):
        tasks = list(tasks)
    raw_list: list[Any] | None = None
    if tasks is None:
        return [], None
    if isinstance(tasks, list):
        raw_list = tasks
    elif isinstance(tasks, str):
        s = tasks.strip()
        if not s:
            return [], None
        s = _strip_markdown_fences(s)
        try:
            parsed = _loads_json_maybe_double(s)
        except json.JSONDecodeError as e:
            return [], (
                "[delegate_task] `tasks` must be valid JSON (array of objects). "
                f"Parse error: {e.msg}. Example: "
                '[{"task_description":"Do X","agent_type":"research"},{"task_description":"Do Y","agent_type":"general"}] '
                "— or omit `tasks` for a single delegation using `task_description`."
            )
        if not isinstance(parsed, list):
            return [], (
                "[delegate_task] `tasks` must be a JSON array of objects with `task_description` "
                f"(optional `agent_type`, `context_snippet`). Got {type(parsed).__name__}, not an array."
            )
        raw_list = parsed
    else:
        return [], (
            "[delegate_task] `tasks` must be a JSON array or a string of JSON, "
            f"not {type(tasks).__name__}."
        )

    if len(raw_list) == 0:
        return [], None

    batch_items: list[BatchTaskItem] = []
    saw_dict = False
    for item in raw_list[:6]:
        if isinstance(item, dict):
            saw_dict = True
            desc = item.get("task_description", "")
            if not isinstance(desc, str):
                desc = str(desc) if desc is not None else ""
            desc = desc.strip()
            if not desc:
                continue
            at = item.get("agent_type", "general")
            at_str = str(at).strip() if at is not None else ""
            if not at_str:
                at_str = "general"
            cs = item.get("context_snippet", "")
            cs_str = cs.strip() if isinstance(cs, str) else (str(cs) if cs is not None else "")
            batch_items.append(BatchTaskItem(
                task_description=desc,
                agent_type=at_str,
                context_snippet=cs_str,
            ))

    if batch_items:
        return batch_items, None
    if saw_dict:
        return [], (
            "[delegate_task] `tasks` was non-empty but no entry had a non-empty `task_description`. "
            "Fix each object or omit `tasks` and use `task_description` for one job."
        )
    return [], (
        "[delegate_task] `tasks` must be a JSON array of objects. "
        "Each object needs at least `task_description` (string)."
    )


_RESTRICTED_CHILD_TOOLS = frozenset({
    "delegate_task",
    "append_to_memory",
    "write_session_state",
    "summarize_old_memory",
    "archive_old_memory",
    "refresh_memory",
    "skill_manage",
})


@tool(args_schema=DelegateTaskInput)
def delegate_task(
    task_description: str,
    agent_type: str = "research",
    context_snippet: str = "",
    handoff_envelope_json: str = "",
    tasks: str | list[Any] = "",
    delegation_job_id: str = "",
) -> str:
    """Delegate a task to a subagent that runs in isolation with its own tools and context.
    The subagent has no personality or memory -- it focuses solely on the task and returns results.
    Use this for research tasks, shell operations, or any work that benefits from a clean, focused context.

    Return value: lines starting with ``delegation_handoff:`` (job id, blackboard and envelope telemetry),
    then ``---``, then the subagent text. Blackboard injection is automatic when enabled in config and a
    delegation job id resolves; ``handoff_envelope_json`` is optional and applies to single-task calls only
    (batch ``tasks`` ignore it in v1). Cite only what those lines report when describing the handoff.

    For parallel work, pass a JSON array of tasks in the 'tasks' parameter."""
    _debug_log(f"[TOOL] delegate_task called: agent_type={agent_type}", "TOOL")
    _emit_input_status("Delegating task…")

    batch_items, batch_err = _parse_delegate_batch_tasks(tasks)
    if batch_err:
        _debug_log(f"[TOOL] delegate_task batch parse: {batch_err[:200]}", "TOOL")
        _emit_input_status("")
        return batch_err

    from persistence import get_current_agent_thread_id
    parent_thread_id = get_current_agent_thread_id()
    root_job_id = (delegation_job_id or "").strip() or uuid4().hex[:12]
    if batch_items:
        return _run_batch(batch_items, parent_thread_id=parent_thread_id, root_job_id=root_job_id)

    from subagent_runner import run_subagent_structured, get_current_iteration_budget
    from tool_policy import get_current_policy_context
    from utils.handoff_envelope import envelope_from_json_str

    handoff_envelope = envelope_from_json_str(handoff_envelope_json)
    if handoff_envelope_json.strip() and handoff_envelope is None:
        _debug_log("[TOOL] delegate_task ignored invalid handoff_envelope_json", "TOOL")

    env_status = _envelope_json_status(handoff_envelope_json, handoff_envelope)
    result = run_subagent_structured(
        task_description=task_description,
        agent_type=agent_type,
        context_snippet=context_snippet or None,
        handoff_envelope=handoff_envelope,
        policy_context=get_current_policy_context(),
        iteration_budget=get_current_iteration_budget(),
        parent_thread_id=parent_thread_id,
        delegation_job_id=root_job_id,
    )
    _emit_input_status("")
    return format_delegate_task_result(result, envelope_json_status=env_status)


def _run_batch(items: list[BatchTaskItem], *, parent_thread_id: str | None, root_job_id: str) -> str:
    """Execute multiple subagent tasks in parallel."""
    _emit_input_status(f"Running {len(items)} subtasks in parallel…")
    from subagent_runner import run_subagent_structured, get_current_iteration_budget
    shared_budget = get_current_iteration_budget()

    results: dict[int, str] = {}

    def run_one(idx: int, item: BatchTaskItem) -> tuple[int, str]:
        from tool_policy import get_current_policy_context
        sub_job_id = f"{root_job_id}_t{idx + 1}"
        result = run_subagent_structured(
            task_description=item.task_description,
            agent_type=item.agent_type,
            context_snippet=item.context_snippet or None,
            policy_context=get_current_policy_context(),
            iteration_budget=shared_budget,
            parent_thread_id=parent_thread_id,
            delegation_job_id=sub_job_id,
        )
        return idx, (
            f"[delegation_job_id={sub_job_id}]\n"
            f"{_batch_handoff_digest_block(result)}\n"
            f"{result.to_string()}"
        )

    with ThreadPoolExecutor(max_workers=_BATCH_MAX_WORKERS) as executor:
        futures = {
            executor.submit(contextvars.copy_context().run, run_one, i, item): i
            for i, item in enumerate(items)
        }
        for fut in as_completed(futures):
            try:
                idx, output = fut.result()
                results[idx] = output
            except Exception as e:
                idx = futures[fut]
                results[idx] = f"Error: {e}"

    parts: list[str] = []
    for i in range(len(items)):
        desc_short = items[i].task_description[:80]
        output = results.get(i, "No result")
        parts.append(f"### Task {i + 1}: {desc_short}\n{output}")

    _emit_input_status("")
    return "\n\n---\n\n".join(parts)
