"""
A2A Protocol Adapter — enables Nexus agents to communicate with external
A2A-compatible agents over HTTP/JSON-RPC.

Design notes:
  - A2AServer: exposes a Nexus agent as an A2A-compatible HTTP endpoint
  - A2AClient: calls external A2A agents from within Nexus
  - Follows the A2A RC v1.0 spec: JSON-RPC 2.0, Task lifecycle, AgentCard
"""

from __future__ import annotations

import logging
from typing import Any

import httpx
from pydantic import BaseModel, Field

from nexus.core.agent import Agent
from nexus.core.task import Task, TaskStatus
from nexus.types import AgentCard, new_id

logger = logging.getLogger("nexus.protocol.a2a")


# ---------------------------------------------------------------------------
# A2A JSON-RPC helpers
# ---------------------------------------------------------------------------


class JsonRpcRequest(BaseModel):
    """A2A JSON-RPC 2.0 request."""
    jsonrpc: str = "2.0"
    method: str
    params: dict[str, Any] = Field(default_factory=dict)
    id: str = Field(default_factory=new_id)


class JsonRpcResponse(BaseModel):
    """A2A JSON-RPC 2.0 response."""
    jsonrpc: str = "2.0"
    result: dict[str, Any] | None = None
    error: dict[str, Any] | None = None
    id: str = ""


# ---------------------------------------------------------------------------
# A2A Client — call external A2A agents
# ---------------------------------------------------------------------------


class A2AClient:
    """
    Client for calling external A2A-compatible agents.

    Usage:
        client = A2AClient()

        # Discover the agent
        card = await client.discover("https://agent.example.com")

        # Send a task
        result = await client.send_message(
            endpoint="https://agent.example.com",
            message="Analyze this dataset",
        )
    """

    def __init__(self, timeout: float = 60.0) -> None:
        self._timeout = timeout
        self._http = httpx.AsyncClient(timeout=timeout)

    async def discover(self, base_url: str) -> AgentCard | None:
        """
        Discover an agent by fetching its AgentCard from the well-known URI.

        A2A spec: GET {base_url}/.well-known/agent-card.json
        """
        url = f"{base_url.rstrip('/')}/.well-known/agent-card.json"
        try:
            response = await self._http.get(url)
            response.raise_for_status()
            data = response.json()
            return AgentCard.model_validate(data)
        except Exception as exc:
            logger.error(f"Failed to discover agent at {url}: {exc}")
            return None

    async def send_message(
        self,
        endpoint: str,
        message: str,
        context_id: str | None = None,
        task_id: str | None = None,
        blocking: bool = True,
    ) -> dict[str, Any]:
        """
        Send a message to an A2A agent.

        A2A spec: POST {endpoint} with JSON-RPC method "message/send"
        """
        request = JsonRpcRequest(
            method="message/send",
            params={
                "message": {
                    "role": "user",
                    "parts": [{"type": "text", "text": message}],
                    "messageId": new_id(),
                },
                "configuration": {
                    "blocking": blocking,
                },
            },
        )

        if context_id:
            request.params["contextId"] = context_id
        if task_id:
            request.params["taskId"] = task_id

        try:
            response = await self._http.post(
                endpoint,
                json=request.model_dump(),
                headers={"Content-Type": "application/json"},
            )
            response.raise_for_status()
            return response.json()
        except Exception as exc:
            logger.error(f"A2A send_message to {endpoint} failed: {exc}")
            return {"error": str(exc)}

    async def get_task(self, endpoint: str, task_id: str) -> dict[str, Any]:
        """
        Get the status of a task.

        A2A spec: POST {endpoint} with JSON-RPC method "tasks/get"
        """
        request = JsonRpcRequest(
            method="tasks/get",
            params={"taskId": task_id},
        )

        try:
            response = await self._http.post(
                endpoint,
                json=request.model_dump(),
                headers={"Content-Type": "application/json"},
            )
            response.raise_for_status()
            return response.json()
        except Exception as exc:
            logger.error(f"A2A get_task failed: {exc}")
            return {"error": str(exc)}

    async def cancel_task(self, endpoint: str, task_id: str) -> dict[str, Any]:
        """Cancel a running task."""
        request = JsonRpcRequest(
            method="tasks/cancel",
            params={"taskId": task_id},
        )
        try:
            response = await self._http.post(
                endpoint,
                json=request.model_dump(),
                headers={"Content-Type": "application/json"},
            )
            response.raise_for_status()
            return response.json()
        except Exception as exc:
            return {"error": str(exc)}

    async def close(self) -> None:
        await self._http.aclose()


# ---------------------------------------------------------------------------
# A2A Server — expose a Nexus agent as A2A endpoint
# ---------------------------------------------------------------------------


class A2AServer:
    """
    Exposes a Nexus Agent as an A2A-compatible HTTP endpoint.

    This is a lightweight adapter that translates between A2A JSON-RPC
    and Nexus's internal task/message model.

    For production use, mount this with an ASGI framework (Starlette/FastAPI).

    Usage:
        agent = Agent(name="MyAgent", ...)
        server = A2AServer(agent)

        # Get the ASGI app to mount
        app = server.build_asgi_app()

        # Or handle a raw request dict
        response = await server.handle_request(request_dict)
    """

    def __init__(self, agent: Agent) -> None:
        self.agent = agent
        self._tasks: dict[str, Task] = {}

    @property
    def agent_card(self) -> AgentCard:
        """Generate the AgentCard for the well-known endpoint."""
        return self.agent.to_agent_card()

    async def handle_request(self, request_data: dict[str, Any]) -> dict[str, Any]:
        """
        Process a JSON-RPC request and return a JSON-RPC response.

        Dispatches to the appropriate handler based on the 'method' field.
        """
        try:
            request = JsonRpcRequest.model_validate(request_data)
        except Exception as exc:
            return self._error_response("", -32700, f"Parse error: {exc}")

        method = request.method
        handlers = {
            "message/send": self._handle_send_message,
            "tasks/get": self._handle_get_task,
            "tasks/cancel": self._handle_cancel_task,
        }

        handler = handlers.get(method)
        if not handler:
            return self._error_response(
                request.id, -32601, f"Method not found: {method}"
            )

        try:
            result = await handler(request.params)
            return JsonRpcResponse(
                id=request.id,
                result=result,
            ).model_dump()
        except Exception as exc:
            logger.error(f"A2A handler error for {method}: {exc}")
            return self._error_response(request.id, -32603, str(exc))

    async def _handle_send_message(self, params: dict[str, Any]) -> dict[str, Any]:
        """Handle message/send — create a task and execute it."""
        msg_data = params.get("message", {})
        parts = msg_data.get("parts", [])
        text = ""
        for part in parts:
            if part.get("type") == "text":
                text += part.get("text", "")

        # Create a Nexus task
        task = Task(
            name="a2a_task",
            description=text,
        )
        self._tasks[task.id] = task

        # Execute (blocking by default for simplicity)
        blocking = params.get("configuration", {}).get("blocking", True)

        if blocking:
            result = await self.agent.run(task)
            return {
                "taskId": task.id,
                "contextId": params.get("contextId", task.conversation.context_id),
                "status": {"state": task.status.value},
                "artifacts": [
                    {
                        "parts": [{"type": "text", "text": result.content}],
                    }
                ],
            }
        else:
            # Non-blocking: start task in background
            import asyncio
            asyncio.create_task(self.agent.run(task))
            return {
                "taskId": task.id,
                "status": {"state": TaskStatus.WORKING.value},
            }

    async def _handle_get_task(self, params: dict[str, Any]) -> dict[str, Any]:
        """Handle tasks/get — return current task status."""
        task_id = params.get("taskId", "")
        task = self._tasks.get(task_id)
        if not task:
            raise KeyError(f"Task {task_id} not found")

        result_artifacts = []
        if task.result:
            result_artifacts.append({
                "parts": [{"type": "text", "text": task.result.content}],
            })

        return {
            "taskId": task.id,
            "status": {"state": task.status.value},
            "artifacts": result_artifacts,
        }

    async def _handle_cancel_task(self, params: dict[str, Any]) -> dict[str, Any]:
        """Handle tasks/cancel."""
        task_id = params.get("taskId", "")
        task = self._tasks.get(task_id)
        if not task:
            raise KeyError(f"Task {task_id} not found")
        task.cancel()
        return {"taskId": task.id, "status": {"state": task.status.value}}

    def _error_response(
        self, request_id: str, code: int, message: str
    ) -> dict[str, Any]:
        return JsonRpcResponse(
            id=request_id,
            error={"code": code, "message": message},
        ).model_dump()
