"""FastAPI mounting helpers for LangGraph-backed A2A agents."""

from __future__ import annotations

from collections.abc import Mapping
from typing import Any

from a2a.server.apps.jsonrpc.fastapi_app import A2AFastAPIApplication
from a2a.server.request_handlers.default_request_handler import DefaultRequestHandler
from a2a.server.tasks import InMemoryTaskStore

from a2a_langgraph.agent_card import build_agent_card
from a2a_langgraph.config import A2ALangGraphConfig
from a2a_langgraph.executor import LangGraphAgentExecutor


def _require_fastapi() -> None:
    try:
        import fastapi  # noqa: F401
    except ImportError as e:
        raise ImportError(
            "The FastAPI integration requires optional dependencies. "
            "Install with: pip install 'a2a-langgraph[fastapi]'"
        ) from e


_EXECUTOR_KWARGS = frozenset({"input_mapper", "output_mapper", "invoke_config_builder"})
_CONFIG_FIELDS = frozenset(A2ALangGraphConfig.model_fields.keys())


def _config_from_mapping(data: Mapping[str, Any]) -> A2ALangGraphConfig:
    filtered = {k: v for k, v in data.items() if k in _CONFIG_FIELDS}
    return A2ALangGraphConfig.model_validate(filtered)


def create_a2a_fastapi_app(
    graph: Any,
    *,
    config: A2ALangGraphConfig,
    mount_path: str = "/agent",
    **executor_kwargs: Any,
) -> Any:
    """Build an A2A JSON-RPC FastAPI sub-app (not mounted on a parent app)."""
    _require_fastapi()
    agent_card = build_agent_card(config, mount_path)
    executor_cls = config.executor_cls or LangGraphAgentExecutor
    executor = executor_cls(graph, **executor_kwargs)
    task_store = config.task_store or InMemoryTaskStore()
    handler = DefaultRequestHandler(agent_executor=executor, task_store=task_store)
    a2a = A2AFastAPIApplication(
        agent_card=agent_card,
        http_handler=handler,
        max_content_length=config.max_content_length,
    )
    return a2a.build(**config.build_app_kwargs)


def add_langgraph_fastapi_endpoint(
    app: Any,
    graph: Any,
    mount_path: str = "/agent",
    **kwargs: Any,
) -> None:
    """Mount an A2A JSON-RPC app that runs the given LangGraph on each request."""
    _require_fastapi()
    executor_kwargs = {k: kwargs.pop(k) for k in list(kwargs) if k in _EXECUTOR_KWARGS}
    config = _config_from_mapping(kwargs)
    sub_app = create_a2a_fastapi_app(
        graph,
        config=config,
        mount_path=mount_path,
        **executor_kwargs,
    )
    app.mount(mount_path, sub_app)
