"""A2A polling helpers — send a question to an A2A agent and await the answer."""

import asyncio
import uuid
import httpx
from fasta2a.client import A2AClient

# A2A task lifecycle constants (spec §4.1.3)
_TERMINAL_STATES = frozenset({"completed", "failed", "canceled", "rejected"})
_INTERRUPTED_STATES = frozenset({"input_required", "auth_required"})
_POLL_INTERVAL = 0.5   # seconds between get_task calls
_POLL_TIMEOUT = 120    # seconds before giving up on a task


def _make_message(text: str, ctx: dict | None = None) -> dict:
    msg: dict = {
        "role": "user",
        "parts": [{"kind": "text", "text": text}],
        "kind": "message",
        "message_id": str(uuid.uuid4()),
    }
    if ctx:
        msg["metadata"] = ctx
    return msg


def _check_jsonrpc_error(response: dict, context: str) -> None:
    """Raise RuntimeError if the JSON-RPC envelope carries an error object."""
    if "error" in response:
        err = response["error"]
        raise RuntimeError(
            f"A2A {context} error {err.get('code')}: {err.get('message')} "
            f"(data={err.get('data')})"
        )


def _extract_text(task: dict) -> str:
    """Extract the answer text from a completed A2A task.

    Preference order per A2A spec §3.7 (results SHOULD be in artifacts):
    1. Artifacts attached to the task
    2. Message attached to the terminal TaskStatus
    3. Last agent message in history (fallback)
    """
    result = task.get("result", task)

    for artifact in result.get("artifacts") or []:
        for part in artifact.get("parts") or []:
            if part.get("kind") == "text" and part.get("text"):
                return part["text"]

    for part in (result.get("status", {}).get("message") or {}).get("parts") or []:
        if part.get("kind") == "text" and part.get("text"):
            return part["text"]

    for msg in reversed(result.get("history") or []):
        if msg.get("role") in ("agent", "assistant"):
            for part in msg.get("parts") or []:
                if part.get("kind") == "text" and part.get("text"):
                    return part["text"]

    state = result.get("status", {}).get("state", "unknown")
    return f"(no text response; task state: {state})"


async def ask_agent(url: str, question: str, ctx: dict | None = None) -> str:
    """Send *question* to the A2A agent at *url* and return its answer.

    Args:
        ctx: Optional metadata dict forwarded as ``Message.metadata`` so the
             receiving agent's tools can access it via ``request_context.get()``.
    """
    client = A2AClient(url, http_client=httpx.AsyncClient(base_url=url, timeout=120.0))

    send_resp = await client.send_message(_make_message(question, ctx))
    _check_jsonrpc_error(send_resp, "send_message")

    result = send_resp.get("result", {})

    if result.get("kind") == "message" or ("parts" in result and "id" not in result):
        for part in result.get("parts") or []:
            if part.get("kind") == "text" and part.get("text"):
                return part["text"]
        return "(direct message with no text)"

    task_id = result.get("id")
    if not task_id:
        return "(agent returned no task id and no direct message)"

    loop = asyncio.get_event_loop()
    deadline = loop.time() + _POLL_TIMEOUT
    while True:
        if loop.time() > deadline:
            return f"(timed out after {_POLL_TIMEOUT}s waiting for task {task_id})"

        task = await client.get_task(task_id)
        _check_jsonrpc_error(task, "get_task")
        state = task.get("result", {}).get("status", {}).get("state", "")

        if state in _TERMINAL_STATES:
            if state in ("failed", "rejected"):
                raise RuntimeError(
                    f"Agent task {task_id} ended with state '{state}': "
                    f"{_extract_text(task)}"
                )
            return _extract_text(task)

        if state in _INTERRUPTED_STATES:
            return f"[{state.upper()}] {_extract_text(task)}"

        await asyncio.sleep(_POLL_INTERVAL)