from typing import Any

from helpers.tool import Tool, Response
from helpers.print_style import PrintStyle
from helpers.fasta2a_client import connect_to_agent, is_client_available


A2A_EMPTY_RESPONSE_ERROR = (
    "A2A chat failed: the remote task completed but no assistant text was found. "
    "Expected final.result.history to include an assistant message with a text "
    "part, or a text artifact/status message. Treat this as a failed remote "
    "response, not success."
)


def _session_key(agent_url: str) -> str:
    """Keep root and explicit /a2a URLs in the same conversation cache."""
    normalized = agent_url.rstrip("/")
    if normalized.endswith("/a2a"):
        return normalized[:-4].rstrip("/")
    return normalized


def _text_from_part(part: Any) -> str:
    if not isinstance(part, dict):
        return ""
    for key in ("text", "content"):
        value = part.get(key)
        if isinstance(value, str) and value.strip():
            return value.strip()
    return ""


def _text_from_message(message: Any) -> str:
    if isinstance(message, str):
        return message.strip()
    if not isinstance(message, dict):
        return ""

    parts = message.get("parts")
    if isinstance(parts, list):
        texts = [_text_from_part(part) for part in parts]
        text = "\n".join(text for text in texts if text)
        if text:
            return text

    for key in ("text", "content", "message", "output"):
        value = message.get(key)
        if isinstance(value, str) and value.strip():
            return value.strip()

    return ""


def _extract_latest_assistant_text(task_response: Any) -> str:
    if not isinstance(task_response, dict):
        return ""

    result = task_response.get("result", task_response)
    if not isinstance(result, dict):
        return ""

    history = result.get("history")
    if isinstance(history, list):
        for message in reversed(history):
            if isinstance(message, dict) and message.get("role") == "user":
                continue
            text = _text_from_message(message)
            if text:
                return text

    status = result.get("status")
    if isinstance(status, dict):
        text = _text_from_message(status.get("message"))
        if text:
            return text

    artifacts = result.get("artifacts")
    if isinstance(artifacts, list):
        for artifact in reversed(artifacts):
            text = _text_from_message(artifact)
            if text:
                return text

    return _text_from_message(result)


class A2AChatTool(Tool):
    """Communicate with another FastA2A-compatible agent."""

    async def execute(self, **kwargs):
        if not is_client_available():
            return Response(message="FastA2A client not available on this instance.", break_loop=False)

        agent_url: str | None = kwargs.get("agent_url")  # required
        user_message: str | None = kwargs.get("message")  # required
        attachments = kwargs.get("attachments", None)  # optional list[str]
        reset = bool(kwargs.get("reset", False))
        if not agent_url or not isinstance(agent_url, str):
            return Response(message="agent_url argument missing", break_loop=False)
        if not user_message or not isinstance(user_message, str):
            return Response(message="message argument missing", break_loop=False)

        # Retrieve or create session cache on the Agent instance
        sessions: dict[str, str] = self.agent.get_data("_a2a_sessions") or {}
        cache_key = _session_key(agent_url)

        # Handle reset flag: start fresh conversation
        if reset and cache_key in sessions:
            sessions.pop(cache_key, None)

        context_id = None if reset else sessions.get(cache_key)
        try:
            async with await connect_to_agent(agent_url) as conn:
                task_resp = await conn.send_message(user_message, attachments=attachments, context_id=context_id)
                task_id = task_resp.get("result", {}).get("id")  # type: ignore[index]
                if not task_id:
                    return Response(message="Remote agent failed to create task.", break_loop=False)
                final = await conn.wait_for_completion(task_id)
                new_context_id = final["result"].get("context_id")  # type: ignore[index]
                if isinstance(new_context_id, str):
                    sessions[cache_key] = new_context_id
                    # persist back to agent data
                    self.agent.set_data("_a2a_sessions", sessions)
                assistant_text = _extract_latest_assistant_text(final)
                if not assistant_text:
                    return Response(
                        message=A2A_EMPTY_RESPONSE_ERROR,
                        break_loop=False,
                    )
                return Response(message=assistant_text, break_loop=False)
        except Exception as e:
            PrintStyle.error(f"A2A chat error: {e}")
            return Response(message=f"A2A chat error: {e}", break_loop=False)