"""Blocking subagent coordination tools: delegate and await.

``ctl_delegate`` spawns an ephemeral child and blocks until completion.
``ctl_await`` blocks until all specified children reach terminal state.

Both tools share a common poll loop and result-reading infrastructure
implemented as private helpers in this module.

Spec source: docs/proposals/subagent-coordination.md (Phase 2).
"""

from __future__ import annotations

import asyncio
import json
import time

from deal import post, pre
from returns.result import Failure, Result, Success

from anima.core.lifecycle import is_terminal
from anima.core.schemas import (
    AgentState,
    AgentStateData,
    AnimaError,
    InternalError,
)
from anima.shell.agent_kill import _is_pid_alive
from anima.shell.ctl_child import _validate_child_id
from anima.shell.ctl_spawn import ctl_spawn
from anima.shell.state import read_result, read_state


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------


async def _poll_until_terminal(
    child_id: str,
    timeout: float | None,
    poll_interval: float = 0.5,
) -> Result[AgentStateData, AnimaError]:
    """Poll a child's state.json until it reaches a terminal state.

    Reads the child's ``state.json`` every *poll_interval* seconds.
    Every 5th poll, performs a PID liveness check as a crash-detection
    fallback.  If the child process died without writing a terminal
    state, synthesises a ``failed`` result.

    Args:
        child_id: Instance ID of the child agent to poll.
        timeout: Maximum seconds to wait.  ``None`` means wait
            indefinitely (parent's own budget is the implicit cap).
        poll_interval: Seconds between state.json reads.  Default 0.5.

    Returns:
        Success(AgentStateData) when the child reaches a terminal state,
        or Failure(AnimaError) on timeout or unexpected I/O error.

    >>> import asyncio
    >>> r = asyncio.run(_poll_until_terminal("nonexistent", timeout=0.1))
    >>> from returns.result import Failure
    >>> isinstance(r, Failure)
    True
    """
    start = time.monotonic()
    iteration = 0

    while True:
        state_result = read_state(child_id)
        if isinstance(state_result, Success):
            state_data = state_result.unwrap()
            if is_terminal(state_data.state):
                return Success(state_data)

        # Check timeout after state read but before PID liveness so
        # short timeouts are honoured before crash-detection fires.
        # Use both wall-clock and virtual elapsed time (iteration * poll_interval)
        # to handle mocked asyncio.sleep correctly in tests.
        elapsed_wall = time.monotonic() - start
        elapsed_virtual = iteration * poll_interval
        elapsed = max(elapsed_wall, elapsed_virtual)
        if timeout is not None and elapsed >= timeout:
            return Failure(
                InternalError(
                    f"Timeout after {timeout}s waiting for child {child_id} to reach terminal state",
                    {"child_id": child_id, "timeout": timeout},
                )
            )

        # Every 5th iteration, check PID liveness as crash detection fallback
        if (
            isinstance(state_result, Success)
            and iteration > 0
            and iteration % 5 == 0
        ):
            sd = state_result.unwrap()
            if sd.pid is not None:
                alive_result = _is_pid_alive(sd.pid)
                if isinstance(alive_result, Success) and not alive_result.unwrap():
                    failed_state = sd.model_copy(
                        update={
                            "state": AgentState.failed,
                            "error": f"Child process (pid {sd.pid}) died unexpectedly (crash detected)",
                        }
                    )
                    return Success(failed_state)

        iteration += 1
        await asyncio.sleep(poll_interval)


@pre(lambda child_id, max_response_chars: max_response_chars is None or max_response_chars > 0)
def _read_child_result(
    child_id: str,
    max_response_chars: int | None,
) -> Result[dict[str, object], AnimaError]:
    """Read a child's result.json and build a result envelope dict.

    Reads the child's ``result.json``, extracts ``response_text``,
    ``error``, and ``final_state``, and applies truncation to
    ``response_text`` when it exceeds *max_response_chars*.

    Args:
        child_id: Instance ID of the child whose result to read.
        max_response_chars: Maximum characters for ``response_text``.
            ``None`` disables truncation.

    Returns:
        Success(dict) with keys ``child_id``, ``status``,
        ``response_text``, ``error``, ``summary``.
        Failure(AnimaError) if result.json is missing or unreadable.

    >>> _read_child_result("nonexistent", max_response_chars=100)  # doctest: +ELLIPSIS
    <Failure: ...>
    """
    result = read_result(child_id)
    if isinstance(result, Failure):
        return Failure(
            InternalError(
                f"Failed to read result for child {child_id}: {result.failure()}",
                {"child_id": child_id},
            )
        )

    agent_result = result.unwrap()
    response_text = agent_result.response_text
    if response_text is not None and max_response_chars is not None:
        response_text = _truncate_response(response_text, max_response_chars, child_id)

    status = agent_result.final_state.value
    summary = f"Child {child_id} {status}."
    if agent_result.error:
        summary = f"Child {child_id} {status}: {agent_result.error}"

    envelope: dict[str, object] = {
        "child_id": child_id,
        "status": status,
        "response_text": response_text,
        "error": agent_result.error,
        "summary": summary,
    }
    return Success(envelope)


# @invar:allow shell_result: Pure dict builder, not an I/O boundary
def _build_failure_envelope(
    child_id: str,
    state_data: AgentStateData,
) -> dict[str, object]:
    """Build an enriched failure envelope from a child's terminal state.

    Constructs a dict containing ``child_id``, ``status``, ``error``,
    ``response_text`` (partial, if any), ``current_turn``, and
    ``debug_hint`` pointing to the child's logs.

    Args:
        child_id: Instance ID of the failed child.
        state_data: The child's ``AgentStateData`` at failure time.

    Returns:
        Dict with failure envelope fields suitable for JSON serialisation.

    >>> from anima.core.schemas import AgentState, AgentStateData
    >>> sd = AgentStateData(
    ...     instance_id="inst_fail",
    ...     trace_id="t1",
    ...     persona_name="p",
    ...     persona_ref="p.json",
    ...     spec_id="p",
    ...     state=AgentState.failed,
    ...     children=[],
    ...     depth=1,
    ...     max_depth=3,
    ...     created_at="2025-01-01T00:00:00Z",
    ...     updated_at="2025-01-01T00:00:00Z",
    ...     tools_called=[],
    ...     current_turn=5,
    ...     error="max_turns exceeded",
    ... )
    >>> env = _build_failure_envelope("inst_fail", sd)
    >>> env["status"]
    'failed'
    >>> env["error"]
    'max_turns exceeded'
    >>> env["current_turn"]
    5
    """
    # Attempt to read partial response_text from result.json if available
    response_text: str | None = None
    result_data = read_result(child_id)
    if isinstance(result_data, Success):
        response_text = result_data.unwrap().response_text

    status = state_data.state.value
    error = state_data.error or ""

    # Provide actionable hint for common failure modes
    hint = ""
    if "budget exhausted" in error.lower() or "grace" in error.lower():
        hint = " The child ran out of token budget. Retry without setting max_tokens (default is unlimited), or set a larger value."
    elif "max_turns" in error.lower():
        hint = " The child ran out of turns. Retry without setting max_turns (default is unlimited), or set a larger value."

    summary = f"Child {child_id} {status}: {error}{hint}" if error else f"Child {child_id} {status}."

    return {
        "child_id": child_id,
        "status": status,
        "error": error if error else None,
        "current_turn": state_data.current_turn,
        "response_text": response_text,
        "debug_hint": f"Full logs: anima logs {child_id}",
        "summary": summary,
    }


# @invar:allow shell_result: Pure string transformation, not an I/O boundary
@pre(lambda text, max_chars, child_id: max_chars is None or max_chars > 0)
@post(lambda result: isinstance(result, str))
def _truncate_response(
    text: str,
    max_chars: int | None,
    child_id: str,
) -> str:
    """Truncate response text with a pointer to the full result file.

    When *text* exceeds *max_chars*, returns a truncated version with
    a suffix pointing to ``~/.anima/agents/{child_id}/result.json``.
    When *max_chars* is ``None``, returns *text* unchanged.

    Args:
        text: The full response text to potentially truncate.
        max_chars: Maximum character count.  ``None`` disables truncation.
        child_id: Instance ID used to construct the result.json path hint.

    Returns:
        The original or truncated text string.

    >>> _truncate_response("short", max_chars=100, child_id="inst_x")
    'short'
    >>> result = _truncate_response("A" * 200, max_chars=50, child_id="inst_abc")
    >>> "result.json" in result
    True
    >>> _truncate_response("long text", max_chars=None, child_id="inst_y")
    'long text'
    """
    if max_chars is None:
        return text
    if len(text) <= max_chars:
        return text
    suffix = f"...[truncated. Full response in ~/.anima/agents/{child_id}/result.json]"
    return text[:max_chars] + suffix


# ---------------------------------------------------------------------------
# Public tools
# ---------------------------------------------------------------------------

_DEFAULT_MAX_RESPONSE_CHARS = 4000


# @invar:allow shell_result: Pure string parser extracts child_id, no I/O
def _extract_child_id(spawn_value: str) -> str:
    """Extract child_id from spawn result (JSON envelope or plain string).

    Args:
        spawn_value: The string returned by ctl_spawn_subagent.

    Returns:
        The child instance ID.
    """
    try:
        data = json.loads(spawn_value)
        if isinstance(data, dict) and "child_id" in data:
            return data["child_id"]
    except (json.JSONDecodeError, TypeError):
        pass
    return spawn_value


# @shell_complexity: Blocking spawn + poll + result read orchestration
# @invar:allow function_size: Delegate orchestrates spawn, poll, and result read
@pre(lambda instance_id, task, persona_ref=None, context_mode=None, max_turns=None, max_tokens=None, timeout=None, max_response_chars=_DEFAULT_MAX_RESPONSE_CHARS, mcp_servers=None: len(instance_id.strip()) > 0 and len(task.strip()) > 0)
@post(lambda result: isinstance(result, Result))
async def ctl_delegate(
    instance_id: str,
    task: str,
    persona_ref: str | None = None,
    context_mode: str | None = None,
    max_turns: int | None = None,
    max_tokens: int | None = None,
    timeout: float | None = None,
    max_response_chars: int | None = _DEFAULT_MAX_RESPONSE_CHARS,
    mcp_servers: list[object] | None = None,
) -> Result[str, AnimaError]:
    """Spawn an ephemeral child agent, wait for completion, return result.

    This is the primary synchronous delegation tool (~80% of subagent use).
    Internally spawns a child with ``lifecycle=ephemeral`` and ``auto_rm=False``,
    then polls ``state.json`` until the child reaches a terminal state.

    The parent agent loop is blocked during the wait (via ``asyncio.sleep``
    in the poll loop), matching the existing ``ctl_suspend_until`` pattern.

    Args:
        instance_id: The calling (parent) agent's instance ID.  Injected
            by the tool factory closure -- never exposed in the LLM schema.
        task: Self-contained instruction for the child agent.
        persona_ref: Child persona file path.  ``None`` inherits the
            parent's ``persona_ref``.
        context_mode: ``"isolated"`` (default) or ``"shared"``.
        max_turns: Child turn limit override.
        max_tokens: Child token budget override.
        timeout: Maximum seconds to wait for the child.  ``None`` means
            no timeout (parent's own budget is the implicit cap).
        max_response_chars: Truncate ``response_text`` to this length.
            Default 4000.  ``None`` disables truncation.

    Returns:
        ``Success(json_string)`` containing the result envelope::

            {
                "child_id": "inst_abc123",
                "status": "completed",
                "response_text": "The analysis found 3 issues...",
                "error": null,
                "summary": "Child inst_abc123 completed successfully."
            }

        On failure, the envelope includes ``current_turn``, partial
        ``response_text``, and ``debug_hint``.

        On timeout::

            {"child_id": "...", "status": "timeout", ...}

        ``Failure(AnimaError)`` on spawn failure or I/O error.

    >>> import asyncio, json
    >>> # (requires live agent state and spawn infrastructure -- see integration tests)
    """
    spawn_result = await ctl_spawn(
        instance_id=instance_id,
        persona_ref=persona_ref,
        task=task,
        context_mode=context_mode or "isolated",
        lifecycle="ephemeral",
        max_turns=max_turns,
        max_tokens=max_tokens,
        mcp_servers=mcp_servers,
    )
    if isinstance(spawn_result, Failure):
        return spawn_result

    child_id = _extract_child_id(spawn_result.unwrap())

    poll_result = await _poll_until_terminal(child_id, timeout)

    if isinstance(poll_result, Failure):
        # Timeout case
        envelope: dict[str, object] = {
            "child_id": child_id,
            "status": "timeout",
            "error": str(poll_result.failure()),
            "summary": f"Child {child_id} timed out after {timeout}s.",
        }
        return Success(json.dumps(envelope))

    state_data = poll_result.unwrap()

    if state_data.state == AgentState.completed:
        read_result_val = _read_child_result(child_id, max_response_chars)
        if isinstance(read_result_val, Success):
            return Success(json.dumps(read_result_val.unwrap()))
        # Fall through to failure envelope if result read fails
        envelope = _build_failure_envelope(child_id, state_data)
    else:
        # Failed, killed, or other terminal state
        envelope = _build_failure_envelope(child_id, state_data)

    return Success(json.dumps(envelope))


# @shell_complexity: Concurrent poll of multiple children with timeout handling
# @invar:allow function_size: Await orchestrates concurrent polling and result aggregation
@pre(lambda instance_id, child_ids, timeout=None, max_response_chars=_DEFAULT_MAX_RESPONSE_CHARS: len(instance_id.strip()) > 0 and len(child_ids) > 0)
@post(lambda result: isinstance(result, Result))
async def ctl_await(
    instance_id: str,
    child_ids: list[str],
    timeout: float | None = None,
    max_response_chars: int | None = _DEFAULT_MAX_RESPONSE_CHARS,
) -> Result[str, AnimaError]:
    """Wait for all specified children to reach terminal state.

    Concurrently polls each child's ``state.json`` every 0.5 seconds.
    Performs a PID liveness check every 5th poll per child.  Blocks the
    parent agent loop until ALL children are terminal or *timeout* expires.

    On partial failure (e.g. child 2 of 3 fails), waiting continues for
    the remaining children.  All results -- including failures -- are
    returned together so the parent LLM can decide how to proceed.

    Args:
        instance_id: The calling (parent) agent's instance ID.  Injected
            by the tool factory closure -- never exposed in the LLM schema.
        child_ids: List of child instance IDs previously returned by
            ``ctl_spawn``.  Validated against the parent's ``children``
            list; invalid IDs produce an error envelope.
        timeout: Maximum seconds to wait.  ``None`` means no timeout.
        max_response_chars: Truncate each child's ``response_text`` to
            this length.  Default 4000.  ``None`` disables truncation.

    Returns:
        ``Success(json_string)`` containing the results envelope::

            {
                "results": [
                    {"child_id": "inst_abc", "status": "completed",
                     "response_text": "...", "error": null},
                    {"child_id": "inst_def", "status": "failed",
                     "response_text": "partial...", "error": "max_turns exceeded",
                     "current_turn": 5,
                     "debug_hint": "Full logs: anima logs inst_def"}
                ],
                "summary": "2 children awaited: 1 completed, 1 failed."
            }

        On timeout, completed children show their results while pending
        children show ``"status": "pending"``.

        ``Failure(AnimaError)`` on child_id validation failure or I/O error.

    >>> import asyncio, json
    >>> # (requires live agent state -- see integration tests)
    """
    # Validate all child_ids belong to the parent
    for cid in child_ids:
        validation = _validate_child_id(instance_id, cid)
        if isinstance(validation, Failure):
            return validation

    # Start concurrent polling for all children
    tasks = {
        cid: asyncio.create_task(_poll_until_terminal(cid, timeout))
        for cid in child_ids
    }

    # Wait for all tasks to complete
    await asyncio.gather(*tasks.values(), return_exceptions=True)

    results: list[dict[str, object]] = []
    status_counts: dict[str, int] = {}

    for cid in child_ids:
        task = tasks[cid]
        if task.done() and not task.cancelled() and task.exception() is None:
            poll_result = task.result()
            if isinstance(poll_result, Success):
                state_data = poll_result.unwrap()
                if state_data.state == AgentState.completed:
                    read_res = _read_child_result(cid, max_response_chars)
                    if isinstance(read_res, Success):
                        child_envelope = read_res.unwrap()
                    else:
                        child_envelope = _build_failure_envelope(cid, state_data)
                else:
                    child_envelope = _build_failure_envelope(cid, state_data)
            else:
                # Timeout for this child
                child_envelope = {
                    "child_id": cid,
                    "status": "pending",
                    "error": str(poll_result.failure()),
                    "summary": f"Child {cid} still pending (timeout).",
                }
        else:
            # Task failed or was cancelled
            child_envelope = {
                "child_id": cid,
                "status": "pending",
                "error": "Poll task failed unexpectedly.",
                "summary": f"Child {cid} still pending.",
            }

        results.append(child_envelope)
        status = str(child_envelope.get("status", "unknown"))
        status_counts[status] = status_counts.get(status, 0) + 1

    summary_parts = [f"{len(child_ids)} children awaited:"]
    for status, count in sorted(status_counts.items()):
        summary_parts.append(f"{count} {status}")
    summary = " ".join(summary_parts[:1]) + " " + ", ".join(summary_parts[1:]) + "."

    outer_envelope: dict[str, object] = {
        "results": results,
        "summary": summary,
    }
    return Success(json.dumps(outer_envelope))
