"""
Agent endpoints for registering + discovering agents via LiteLLM.

Follows the A2A Spec.

1. Register an agent via POST `/v1/agents`
2. Discover agents via GET `/v1/agents`
3. Get specific agent via GET `/v1/agents/{agent_id}`
"""

import asyncio
import os
import uuid
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Annotated, Final, TypedDict, assert_never

from fastapi import APIRouter, Depends, HTTPException, Query, Request
from typing_extensions import ReadOnly, Required

import litellm
from litellm._logging import verbose_proxy_logger
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.proxy._types import (
    CommonProxyErrors,
    LitellmUserRoles,
    UserAPIKeyAuth,
    user_api_key_has_admin_view,
)
from litellm.proxy.a2a.agent_card import (
    SUPPORTED_A2A_PROTOCOL_VERSIONS,
    merge_agent_card,
    normalize_protocol_version,
)
from litellm.proxy.agent_endpoints.agent_registry import (
    parse_agent_litellm_params,
    redact_sensitive_agent_litellm_params,
)
from litellm.proxy.agent_endpoints.agent_search import (
    DEFAULT_AGENT_SEARCH_TOP_K,
    AgentSearchEmbeddingFailed,
    AgentSearchHits,
    AgentSearchNotConfigured,
    global_agent_search_index,
    search_agents,
)
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import accessible_agents
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user
from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity
from litellm.proxy.utils import get_custom_url
from litellm.types.agents import (
    AgentCard,
    AgentConfig,
    AgentKeySummary,
    AgentMakePublicResponse,
    AgentResponse,
    MakeAgentsPublicRequest,
    PatchAgentRequest,
)
from litellm.types.llms.custom_http import httpxSpecialProvider
from litellm.types.proxy.management_endpoints.common_daily_activity import (
    DailySpendMetadata,
    SpendAnalyticsPaginatedResponse,
)


def _proxy_base_url(http_request: Request) -> str:
    """Return the proxy's public base URL, preferring PROXY_BASE_URL when set."""
    return get_custom_url(str(http_request.base_url), route=None)


def _validate_protocol_version(upstream_card: AgentCard | None) -> None:
    """Reject an agent card pinning an unsupported A2A protocol version."""
    version: Final = upstream_card.get("protocolVersion") if upstream_card else None
    if version is not None and normalize_protocol_version(version) is None:
        raise HTTPException(
            status_code=400,
            detail=(
                f"Unsupported protocolVersion '{version}'. "
                f"Supported versions: {', '.join(SUPPORTED_A2A_PROTOCOL_VERSIONS)}."
            ),
        )


def _build_merged_agent_card(
    upstream_card: AgentCard | None,
    *,
    agent_id: str,
    http_request: Request,
    agent_name: str | None = None,
) -> dict[str, object]:
    """Apply the LiteLLM-fronting merge to ``upstream_card`` for ``agent_id``."""
    proxy_base: Final = _proxy_base_url(http_request)
    _validate_protocol_version(upstream_card)
    # Prefer a card-supplied ``name`` (the discovery UI exposes an editable
    # "Name (shown to API clients)" field that flows into
    # ``agent_card_params.name``) over the internal ``agent_name`` identifier.
    # Fall back to ``agent_name`` only when the card itself has no name.
    card_name: Final = upstream_card.get("name") if upstream_card else None
    return merge_agent_card(
        upstream_card,
        proxy_url=f"{proxy_base}/a2a/{agent_id}",
        proxy_base_url=proxy_base,
        name=card_name or agent_name,
    )


router: Final = APIRouter()


async def _attach_keys_to_agents(agents: Sequence[AgentResponse], prisma_client) -> None:
    """Attach each agent's virtual keys, derived from the key table's agent_id
    foreign key. Mirrors how spend is joined into the agent response so the UI
    never has to cross-reference a full key dump client-side. Only non-secret
    fields are exposed (alias, masked key_name, hashed token)."""
    from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry

    agent_ids: Final = tuple(
        alias_id for agent in agents for alias_id in global_agent_registry.ids_for_agent(agent.agent_id)
    )
    if not agent_ids:
        return
    key_rows: Final = await prisma_client.db.litellm_verificationtoken.find_many(
        where={"agent_id": {"in": agent_ids}},
    )
    keys_by_agent: Final[dict[str, list[AgentKeySummary]]] = {}
    for row in key_rows:
        keys_by_agent.setdefault(row.agent_id, []).append(
            AgentKeySummary(
                token=row.token,
                key_alias=row.key_alias,
                key_name=row.key_name,
            )
        )
    for agent in agents:
        matched_keys = [
            key_summary
            for alias_id in global_agent_registry.ids_for_agent(agent.agent_id)
            for key_summary in keys_by_agent.get(alias_id) or ()
        ]
        agent.keys = matched_keys or None


def _redact_agent_litellm_params_dict(
    litellm_params: Mapping[str, object],
) -> dict[str, object]:  # mutable-ok: AgentResponse.litellm_params is declared as a plain dict, not Mapping
    """Type-narrowing wrapper: a dict in always yields a dict back from
    ``redact_sensitive_agent_litellm_params``, which the function's general
    (possible-JSON-string, possibly-None) signature can't express."""
    return dict(  # mutable-ok: AgentResponse.litellm_params is declared as a plain dict, not Mapping
        parse_agent_litellm_params(redact_sensitive_agent_litellm_params(litellm_params))
    )


def _redact_sensitive_agent_fields(
    agents: Sequence[AgentResponse],
    *,
    is_admin: bool,
) -> list[AgentResponse]:
    """
    Return copies of the given agents with credential-bearing litellm_params
    values replaced by a fixed marker (never returned to ANY caller,
    admin included) and, for non-admin callers, virtual-key and header
    fields stripped entirely. The original objects are not modified.
    """
    redacted: Final[list[AgentResponse]] = []
    for agent in agents:
        copy = agent.model_copy(deep=True)
        if not is_admin:
            copy.static_headers = None
            copy.extra_headers = None
            copy.keys = None
        if copy.litellm_params:
            copy.litellm_params = _redact_agent_litellm_params_dict(copy.litellm_params)
        redacted.append(copy)
    return redacted


def _check_agent_management_permission(user_api_key_dict: UserAPIKeyAuth) -> None:
    """
    Raises HTTP 403 if the caller does not have permission to create, update,
    or delete agents.  Only PROXY_ADMIN users are allowed to perform these
    write operations.
    """
    if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
        raise HTTPException(
            status_code=403,
            detail={
                "error": f"Only proxy admins can create, update, or delete agents. Your role={user_api_key_dict.user_role}"
            },
        )


AGENT_HEALTH_CHECK_TIMEOUT_SECONDS: Final = float(os.environ.get("LITELLM_AGENT_HEALTH_CHECK_TIMEOUT", "5.0"))
AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS = float(os.environ.get("LITELLM_AGENT_HEALTH_CHECK_GATHER_TIMEOUT", "30.0"))


class _AgentHealthResult(TypedDict, total=False):
    agent_id: Required[str]
    healthy: Required[bool]
    error: str


async def _check_agent_url_health(
    agent: AgentResponse,
) -> _AgentHealthResult:
    """
    Perform a GET request against the agent's URL and return the health result.

    Returns a dict with ``agent_id``, ``healthy`` (bool), and an optional
    ``error`` message.
    """
    url: Final = (agent.agent_card_params or {}).get("url")
    if not url:
        return {"agent_id": agent.agent_id, "healthy": True}

    try:
        client: Final = get_async_httpx_client(
            llm_provider=httpxSpecialProvider.AgentHealthCheck,
            params={"timeout": AGENT_HEALTH_CHECK_TIMEOUT_SECONDS},
        )
        response: Final = await client.get(url)
        if response.status_code >= 500:
            return {
                "agent_id": agent.agent_id,
                "healthy": False,
                "error": f"HTTP {response.status_code}",
            }
        return {"agent_id": agent.agent_id, "healthy": True}
    except Exception as exc:
        return {
            "agent_id": agent.agent_id,
            "healthy": False,
            "error": str(exc),
        }


class _AgentSearchErrorDetail(TypedDict):
    error: ReadOnly[str]
    message: ReadOnly[str]


def _agent_search_error(status_code: int, error: str, message: str) -> HTTPException:
    detail: Final[_AgentSearchErrorDetail] = {"error": error, "message": message}
    return HTTPException(status_code=status_code, detail=detail)


async def _rank_agents_by_query(
    query: str, agents: Sequence[AgentResponse], top_k: int, user_api_key_dict: UserAPIKeyAuth
) -> tuple[AgentResponse, ...]:
    from litellm.proxy.proxy_server import llm_router

    outcome: Final = await search_agents(
        query=query,
        agents=agents,
        top_k=top_k,
        router=llm_router,
        embedding_model=litellm.agent_search_embedding_model,
        index=global_agent_search_index,
        user_api_key_dict=user_api_key_dict,
    )
    match outcome:
        case AgentSearchHits(hits):
            return tuple(hit.agent.model_copy(update=MappingProxyType({"search_score": hit.score})) for hit in hits)
        case AgentSearchNotConfigured(reason):
            raise _agent_search_error(400, "agent_search_not_configured", reason)
        case AgentSearchEmbeddingFailed(reason):
            raise _agent_search_error(503, "agent_search_unavailable", reason)
        case _:
            assert_never(outcome)


@router.get(
    "/v1/agents",
    tags=["[beta] A2A Agents"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=list[AgentResponse],
)
async def get_agents(
    request: Request,
    health_check: bool = Query(
        False,
        description="When true, performs a GET request to each agent's URL. Agents with reachable URLs (HTTP status < 500) and agents without a URL are returned; unreachable agents are filtered out.",
    ),
    query: Annotated[
        str | None,
        Query(
            min_length=1,
            description="Describe the task in natural language to rank the agents you can reach by semantic similarity over their name, description, and skills. Each result carries a search_score. Requires litellm_settings.agent_search_embedding_model.",
        ),
    ] = None,
    top_k: Annotated[
        int,
        Query(ge=1, le=100, description="With query: the maximum number of ranked agents to return."),
    ] = DEFAULT_AGENT_SEARCH_TOP_K,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),  # Used for auth
):
    """
    Example usage:
    ```
    curl -X GET "http://localhost:4000/v1/agents" \
      -H "Content-Type: application/json" \
      -H "Authorization: Bearer your-key" \
    ```

    Pass `?health_check=true` to filter out agents whose URL is unreachable:
    ```
    curl -X GET "http://localhost:4000/v1/agents?health_check=true" \
      -H "Content-Type: application/json" \
      -H "Authorization: Bearer your-key" \
    ```

    Pass `?query=<task>` to get the best matching agents ranked by semantic similarity:
    ```
    curl -X GET "http://localhost:4000/v1/agents?query=translate+a+PDF+document&top_k=5" \
      -H "Content-Type: application/json" \
      -H "Authorization: Bearer your-key" \
    ```

    Returns: List[AgentResponse]

    """
    await check_feature_access_for_user(user_api_key_dict, "agents")

    from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry

    try:
        returned_agents: Sequence[AgentResponse] = await accessible_agents(user_api_key_dict)

        # Fetch current spend from DB for all returned agents
        from litellm.proxy.proxy_server import prisma_client

        if prisma_client is not None:
            agent_ids: Final = tuple(
                alias_id
                for agent in returned_agents
                for alias_id in global_agent_registry.ids_for_agent(agent.agent_id)
            )
            if agent_ids:
                db_agents: Final = await agents_table(prisma_client).find_many(
                    where={"agent_id": {"in": agent_ids}},
                )
                spend_map: Final = {a.agent_id: a.spend for a in db_agents}
                for agent in returned_agents:
                    matched_spends = tuple(
                        spend_map[alias_id]
                        for alias_id in global_agent_registry.ids_for_agent(agent.agent_id)
                        if alias_id in spend_map
                    )
                    if matched_spends:
                        agent.spend = sum(matched_spends)
                await _attach_keys_to_agents(returned_agents, prisma_client)

        # add is_public field to each agent - we do it this way, to allow setting config agents as public
        for agent in returned_agents:
            if agent.litellm_params is None:
                agent.litellm_params = {}
            agent.litellm_params["is_public"] = litellm.public_agent_groups is not None and not (
                global_agent_registry.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups)
            )

        # litellm_params secrets are always redacted; keys/headers stay
        # admin-only.
        is_admin: Final = (
            user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
            or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
        )
        returned_agents = _redact_sensitive_agent_fields(returned_agents, is_admin=is_admin)

        if health_check:
            agents_with_url: Final = [agent for agent in returned_agents if (agent.agent_card_params or {}).get("url")]
            agents_without_url = [agent for agent in returned_agents if not (agent.agent_card_params or {}).get("url")]
            try:
                health_results: Sequence[_AgentHealthResult] = await asyncio.wait_for(
                    asyncio.gather(*[_check_agent_url_health(agent) for agent in agents_with_url]),
                    timeout=AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS,
                )
            except asyncio.TimeoutError:
                verbose_proxy_logger.warning(
                    "Agent health check gather timed out after %s seconds",
                    AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS,
                )
                health_results = [
                    {
                        "agent_id": agent.agent_id,
                        "healthy": False,
                        "error": "Health check timed out",
                    }
                    for agent in agents_with_url
                ]
            healthy_ids: Final = {result["agent_id"] for result in health_results if result["healthy"]}
            returned_agents = [agent for agent in agents_with_url if agent.agent_id in healthy_ids] + agents_without_url

        if query is None:
            return returned_agents
        return await _rank_agents_by_query(query, returned_agents, top_k, user_api_key_dict)
    except HTTPException:
        raise
    except Exception as e:
        verbose_proxy_logger.exception("litellm.proxy.agent_endpoints.get_agents(): Exception occurred - %s", e)
        raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"})


#### CRUD ENDPOINTS FOR AGENTS ####

from litellm.proxy.agent_endpoints.agent_registry import (
    agents_table,
)
from litellm.proxy.agent_endpoints.agent_registry import (
    global_agent_registry as AGENT_REGISTRY,
)


@router.post(
    "/v1/agents",
    tags=["[beta] A2A Agents"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=AgentResponse,
)
async def create_agent(
    request: AgentConfig,
    http_request: Request,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Create a new agent

    Example Request:
    ```bash
    curl -X POST "http://localhost:4000/v1/agents" \\
        -H "Authorization: Bearer <your_api_key>" \\
        -H "Content-Type: application/json" \\
        -d '{
            "agent_name": "my-custom-agent",
            "agent_card_params": {
                "protocolVersion": "1.0",
                "name": "Hello World Agent",
                "description": "Just a hello world agent",
                "url": "http://localhost:9999/",
                "version": "1.0.0",
                "defaultInputModes": ["text"],
                "defaultOutputModes": ["text"],
                "capabilities": {
                    "streaming": true
                },
                "skills": [
                    {
                        "id": "hello_world",
                        "name": "Returns hello world",
                        "description": "just returns hello world",
                        "tags": ["hello world"],
                        "examples": ["hi", "hello world"]
                    }
                ]
            },
            "litellm_params": {
                "make_public": true
            }
        }'
    ```
    """
    await check_feature_access_for_user(user_api_key_dict, "agents")

    from litellm.proxy.proxy_server import prisma_client

    _check_agent_management_permission(user_api_key_dict)

    if prisma_client is None:
        raise HTTPException(status_code=500, detail="Prisma client not initialized")

    try:
        # Get the user ID from the API key auth
        created_by: Final = user_api_key_dict.user_id or "unknown"

        # check for naming conflicts
        existing_agent: Final = AGENT_REGISTRY.get_agent_by_name(agent_name=request.get("agent_name"))
        if existing_agent is not None:
            raise HTTPException(
                status_code=400,
                detail=f"Agent with name {request.get('agent_name')} already exists",
            )

        # Apply the LiteLLM-fronting merge only when the admin actually
        # provided an agent card. Plain chat/LLM agents register without
        # ``agent_card_params``, and synthesising a default A2A card for them
        # would advertise capabilities (``supportedInterfaces``, security
        # schemes, default skills) the agent doesn't actually expose.
        upstream_card: Final = request.get("agent_card_params")
        agent_to_create: AgentConfig = request
        new_agent_id: str | None = None
        if upstream_card is not None:
            # Pre-generate the agent_id so the merged card can reference it
            # in ``supportedInterfaces`` before the DB row exists.
            new_agent_id = str(uuid.uuid4())
            merged_card: Final = _build_merged_agent_card(
                upstream_card,
                agent_id=new_agent_id,
                http_request=http_request,
                agent_name=request.get("agent_name"),
            )
            agent_to_create = {**request, "agent_card_params": merged_card}

        result: Final = await AGENT_REGISTRY.add_agent_to_db(
            agent=agent_to_create,
            prisma_client=prisma_client,
            created_by=created_by,
            agent_id=new_agent_id,
        )

        agent_name: Final = result.agent_name
        agent_id: Final = result.agent_id

        # Also register in memory
        try:
            AGENT_REGISTRY.register_agent(agent_config=result)
            verbose_proxy_logger.info("Successfully registered agent '%s' (ID: %s) in memory", agent_name, agent_id)
        except Exception as reg_error:
            verbose_proxy_logger.warning(
                "Failed to register agent '%s' (ID: %s) in memory: %s", agent_name, agent_id, reg_error
            )

        # The caller is a proxy admin (enforced above); litellm_params
        # secrets are still never echoed back in the response.
        return _redact_sensitive_agent_fields((result,), is_admin=True)[0]

    except HTTPException:
        raise
    except Exception as e:
        verbose_proxy_logger.exception("Error adding agent to db: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.get(
    "/v1/agents/{agent_id}",
    tags=["[beta] A2A Agents"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=AgentResponse,
)
async def get_agent_by_id(
    agent_id: str,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Get a specific agent by ID

    Example Request:
    ```bash
    curl -X GET "http://localhost:4000/v1/agents/123e4567-e89b-12d3-a456-426614174000" \\
        -H "Authorization: Bearer <your_api_key>"
    ```
    """
    await check_feature_access_for_user(user_api_key_dict, "agents")

    if not user_api_key_has_admin_view(user_api_key_dict):
        from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
            AgentRequestHandler,
        )

        is_allowed = await AgentRequestHandler.is_agent_allowed(agent_id=agent_id, user_api_key_auth=user_api_key_dict)
        if not is_allowed:
            raise HTTPException(
                status_code=403,
                detail=f"Agent '{agent_id}' is not allowed for your key/team. Contact proxy admin for access.",
            )

    from litellm.proxy.proxy_server import prisma_client

    if prisma_client is None:
        raise HTTPException(status_code=500, detail="Prisma client not initialized")

    try:
        agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id)
        if agent is None:
            agent_row: Final = await agents_table(prisma_client).find_unique(
                where={"agent_id": agent_id},
                include={"object_permission": True},
            )
            if agent_row is not None:
                agent_dict: Final = agent_row.model_dump()
                if agent_row.object_permission is not None:
                    try:
                        agent_dict["object_permission"] = agent_row.object_permission.model_dump()
                    except Exception:
                        agent_dict["object_permission"] = agent_row.object_permission.dict()
                agent = AgentResponse(**agent_dict)
        else:
            # Agent found in memory — refresh spend from DB
            db_row: Final = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id})
            if db_row is not None:
                agent.spend = db_row.spend

        if agent is None:
            raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found")

        await _attach_keys_to_agents([agent], prisma_client)

        # litellm_params secrets are always redacted; keys/headers stay
        # admin-only.
        is_admin = (
            user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
            or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
        )
        agent = _redact_sensitive_agent_fields((agent,), is_admin=is_admin)[0]

        return agent
    except HTTPException:
        raise
    except Exception as e:
        verbose_proxy_logger.exception("Error getting agent from db: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.put(
    "/v1/agents/{agent_id}",
    tags=["[beta] A2A Agents"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=AgentResponse,
)
async def update_agent(
    agent_id: str,
    request: AgentConfig,
    http_request: Request,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Update an existing agent

    Example Request:
    ```bash
    curl -X PUT "http://localhost:4000/v1/agents/123e4567-e89b-12d3-a456-426614174000" \\
        -H "Authorization: Bearer <your_api_key>" \\
        -H "Content-Type: application/json" \\
        -d '{
            "agent_name": "updated-agent",
            "agent_card_params": {
                "protocolVersion": "1.0",
                "name": "Updated Agent",
                "description": "Updated description",
                "url": "http://localhost:9999/",
                "version": "1.1.0",
                "defaultInputModes": ["text"],
                "defaultOutputModes": ["text"],
                "capabilities": {
                    "streaming": true
                },
                "skills": []
            },
            "litellm_params": {
                "make_public": false
            }
        }'
    ```
    """
    await check_feature_access_for_user(user_api_key_dict, "agents")

    from litellm.proxy.proxy_server import prisma_client

    _check_agent_management_permission(user_api_key_dict)

    if prisma_client is None:
        raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)

    try:
        # Check if agent exists
        existing_agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id})
        if existing_agent is not None:
            existing_agent = dict(existing_agent)

        if existing_agent is None:
            raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found")

        # Get the user ID from the API key auth
        updated_by: Final = user_api_key_dict.user_id or "unknown"

        # Re-apply the LiteLLM-fronting merge — an update is a re-registration,
        # so any new upstream card the admin pasted must go through the same
        # transformation as initial create. Plain agents without an
        # ``agent_card_params`` skip the merge so we don't synthesise an A2A
        # card for them.
        upstream_card: Final = request.get("agent_card_params")
        agent_to_update: AgentConfig = request
        if upstream_card is not None:
            merged_card: Final = _build_merged_agent_card(
                upstream_card,
                agent_id=agent_id,
                http_request=http_request,
                agent_name=request.get("agent_name"),
            )
            agent_to_update = {**request, "agent_card_params": merged_card}

        result: Final = await AGENT_REGISTRY.update_agent_in_db(
            agent_id=agent_id,
            agent=agent_to_update,
            prisma_client=prisma_client,
            updated_by=updated_by,
        )

        # deregister in memory
        AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name"))
        # register in memory
        AGENT_REGISTRY.register_agent(agent_config=result)

        verbose_proxy_logger.info(
            "Successfully updated agent '%s' (ID: %s) in memory", existing_agent.get("agent_name"), agent_id
        )

        return _redact_sensitive_agent_fields((result,), is_admin=True)[0]
    except HTTPException:
        raise
    except Exception as e:
        verbose_proxy_logger.exception("Error updating agent: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.patch(
    "/v1/agents/{agent_id}",
    tags=["[beta] A2A Agents"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=AgentResponse,
)
async def patch_agent(
    agent_id: str,
    request: PatchAgentRequest,
    http_request: Request,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Update an existing agent

    Example Request:
    ```bash
    curl -X PATCH "http://localhost:4000/v1/agents/123e4567-e89b-12d3-a456-426614174000" \\
        -H "Authorization: Bearer <your_api_key>" \\
        -H "Content-Type: application/json" \\
        -d '{
            "agent_name": "updated-agent",
            "agent_card_params": {
                "protocolVersion": "1.0",
                "name": "Updated Agent",
                "description": "Updated description",
                "url": "http://localhost:9999/",
                "version": "1.1.0",
                "defaultInputModes": ["text"],
                "defaultOutputModes": ["text"],
                "capabilities": {
                    "streaming": true
                },
                "skills": []
            },
            "litellm_params": {
                "make_public": false
            }
        }'
    ```
    """
    await check_feature_access_for_user(user_api_key_dict, "agents")

    from litellm.proxy.proxy_server import prisma_client

    _check_agent_management_permission(user_api_key_dict)

    if prisma_client is None:
        raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)

    try:
        # Check if agent exists
        existing_agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id})
        if existing_agent is not None:
            existing_agent = dict(existing_agent)

        if existing_agent is None:
            raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found")

        # Get the user ID from the API key auth
        updated_by: Final = user_api_key_dict.user_id or "unknown"

        # Re-merge only when the patch actually touches agent_card_params; a
        # patch updating just litellm_params/rate limits (``agent_card_params``
        # omitted) shouldn't rewrite the stored card. An explicitly provided
        # ``agent_card_params`` — even an empty dict — still goes through the
        # merge so LiteLLM applies its security schemes and supported
        # interfaces instead of storing a bare card.
        patch_payload: PatchAgentRequest = request
        upstream_card: Final = request.get("agent_card_params")
        if upstream_card is not None:
            merged_card: Final = _build_merged_agent_card(
                upstream_card,
                agent_id=agent_id,
                http_request=http_request,
                agent_name=request.get("agent_name"),
            )
            patch_payload = {**request, "agent_card_params": merged_card}

        result: Final = await AGENT_REGISTRY.patch_agent_in_db(
            agent_id=agent_id,
            agent=patch_payload,
            prisma_client=prisma_client,
            updated_by=updated_by,
        )

        # deregister in memory
        AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name"))
        # register in memory
        AGENT_REGISTRY.register_agent(agent_config=result)

        verbose_proxy_logger.info(
            "Successfully updated agent '%s' (ID: %s) in memory", existing_agent.get("agent_name"), agent_id
        )

        return _redact_sensitive_agent_fields((result,), is_admin=True)[0]
    except HTTPException:
        raise
    except Exception as e:
        verbose_proxy_logger.exception("Error updating agent: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.delete(
    "/v1/agents/{agent_id}",
    tags=["Agents"],
    dependencies=[Depends(user_api_key_auth)],
)
async def delete_agent(
    agent_id: str,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Delete an agent

    Example Request:
    ```bash
    curl -X DELETE "http://localhost:4000/v1/agents/123e4567-e89b-12d3-a456-426614174000" \\
        -H "Authorization: Bearer <your_api_key>"
    ```

    Example Response:
    ```json
    {
        "message": "Agent 123e4567-e89b-12d3-a456-426614174000 deleted successfully"
    }
    ```
    """
    await check_feature_access_for_user(user_api_key_dict, "agents")

    from litellm.proxy.proxy_server import prisma_client

    _check_agent_management_permission(user_api_key_dict)

    if prisma_client is None:
        raise HTTPException(status_code=500, detail="Prisma client not initialized")

    try:
        # Check if agent exists
        existing_agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id})
        if existing_agent is not None:
            existing_agent = dict[str, object](existing_agent)

        if existing_agent is None:
            raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found in DB.")

        await AGENT_REGISTRY.delete_agent_from_db(agent_id=agent_id, prisma_client=prisma_client)

        AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name"))

        return {"message": f"Agent {agent_id} deleted successfully"}
    except HTTPException:
        raise
    except Exception as e:
        verbose_proxy_logger.exception("Error deleting agent: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.post(
    "/v1/agents/{agent_id}/make_public",
    tags=["[beta] A2A Agents"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=AgentMakePublicResponse,
)
async def make_agent_public(
    agent_id: str,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Make an agent publicly discoverable

    Example Request:
    ```bash
    curl -X POST "http://localhost:4000/v1/agents/123e4567-e89b-12d3-a456-426614174000/make_public" \\
        -H "Authorization: Bearer <your_api_key>" \\
        -H "Content-Type: application/json"
    ```

    Example Response:
    ```json
    {
        "agent_id": "123e4567-e89b-12d3-a456-426614174000",
        "agent_name": "my-custom-agent",
        "litellm_params": {
            "make_public": true
        },
        "agent_card_params": {...},
        "created_at": "2025-11-15T10:30:00Z",
        "updated_at": "2025-11-15T10:35:00Z",
        "created_by": "user123",
        "updated_by": "user123"
    }
    ```
    """
    from litellm.proxy.proxy_server import prisma_client

    if prisma_client is None:
        raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)

    try:
        # Update the public model groups
        import litellm
        from litellm.proxy.agent_endpoints.agent_registry import (
            global_agent_registry as AGENT_REGISTRY,
        )
        from litellm.proxy.proxy_server import proxy_config

        # Check if user has admin permissions
        if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
            raise HTTPException(
                status_code=403,
                detail={
                    "error": f"Only proxy admins can update public model groups. Your role={user_api_key_dict.user_role}"
                },
            )

        agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id)
        if agent is None:
            # check if agent exists in DB
            agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id})
            if agent is not None:
                agent = AgentResponse(**agent.model_dump())

            if agent is None:
                raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found")

        if litellm.public_agent_groups is None:
            litellm.public_agent_groups = []
        # handle duplicates
        if not AGENT_REGISTRY.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups):
            raise HTTPException(
                status_code=400,
                detail=f"Agent with name {agent.agent_name} already in public agent groups",
            )
        litellm.public_agent_groups.append(agent.agent_id)

        # Load existing config
        config: Final = await proxy_config.get_config()

        # Update config with new settings
        if "litellm_settings" not in config or config["litellm_settings"] is None:
            config["litellm_settings"] = {}

        config["litellm_settings"]["public_agent_groups"] = litellm.public_agent_groups

        # Save the updated config
        await proxy_config.save_config(new_config=config)

        verbose_proxy_logger.debug(
            "Updated public agent groups to: %s by user: %s", litellm.public_agent_groups, user_api_key_dict.user_id
        )

        return {
            "message": "Successfully updated public agent groups",
            "public_agent_groups": litellm.public_agent_groups,
            "updated_by": user_api_key_dict.user_id,
        }
    except HTTPException:
        raise
    except Exception as e:
        verbose_proxy_logger.exception("Error making agent public: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.post(
    "/v1/agents/make_public",
    tags=["[beta] A2A Agents"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=AgentMakePublicResponse,
)
async def make_agents_public(
    request: MakeAgentsPublicRequest,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Make multiple agents publicly discoverable

    Example Request:
    ```bash
    curl -X POST "http://localhost:4000/v1/agents/make_public" \\
        -H "Authorization: Bearer <your_api_key>" \\
        -H "Content-Type: application/json" \\
        -d '{
            "agent_ids": ["123e4567-e89b-12d3-a456-426614174000", "123e4567-e89b-12d3-a456-426614174001"]
        }'
    ```

    Example Response:
    ```json
    {
        "agent_id": "123e4567-e89b-12d3-a456-426614174000",
        "agent_name": "my-custom-agent",
        "litellm_params": {
            "make_public": true
        },
        "agent_card_params": {...},
        "created_at": "2025-11-15T10:30:00Z",
        "updated_at": "2025-11-15T10:35:00Z",
        "created_by": "user123",
        "updated_by": "user123"
    }
    ```
    """
    from litellm.proxy.proxy_server import prisma_client

    if prisma_client is None:
        raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)

    try:
        # Update the public model groups
        import litellm
        from litellm.proxy.agent_endpoints.agent_registry import (
            global_agent_registry as AGENT_REGISTRY,
        )
        from litellm.proxy.proxy_server import proxy_config

        # Load existing config
        config: Final = await proxy_config.get_config()
        # Check if user has admin permissions
        if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
            raise HTTPException(
                status_code=403,
                detail={
                    "error": f"Only proxy admins can update public model groups. Your role={user_api_key_dict.user_role}"
                },
            )

        if litellm.public_agent_groups is None:
            litellm.public_agent_groups = []

        for agent_id in request.agent_ids:
            agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id)
            if agent is None:
                # check if agent exists in DB
                agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id})
                if agent is not None:
                    agent = AgentResponse(**agent.model_dump())

                if agent is None:
                    raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found")

        litellm.public_agent_groups = request.agent_ids

        # Update config with new settings
        if "litellm_settings" not in config or config["litellm_settings"] is None:
            config["litellm_settings"] = {}

        config["litellm_settings"]["public_agent_groups"] = litellm.public_agent_groups

        # Save the updated config
        await proxy_config.save_config(new_config=config)

        verbose_proxy_logger.debug(
            "Updated public agent groups to: %s by user: %s", litellm.public_agent_groups, user_api_key_dict.user_id
        )

        return {
            "message": "Successfully updated public agent groups",
            "public_agent_groups": litellm.public_agent_groups,
            "updated_by": user_api_key_dict.user_id,
        }
    except HTTPException:
        raise
    except Exception as e:
        verbose_proxy_logger.exception("Error making agent public: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.get(
    "/agent/daily/activity",
    tags=["Agent Management"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=SpendAnalyticsPaginatedResponse,
)
async def get_agent_daily_activity(
    agent_ids: str | None = None,
    start_date: str | None = None,
    end_date: str | None = None,
    model: str | None = None,
    api_key: str | None = None,
    page: int = 1,
    page_size: int = 10,
    exclude_agent_ids: str | None = None,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Get daily activity for specific agents or all accessible agents.
    """
    await check_feature_access_for_user(user_api_key_dict, "agents")

    from litellm.proxy.proxy_server import prisma_client

    if prisma_client is None:
        raise HTTPException(
            status_code=500,
            detail={"error": CommonProxyErrors.db_not_connected_error.value},
        )

    agent_ids_list = agent_ids.split(",") if agent_ids else None
    exclude_agent_ids_list: list[str] | None = None
    if exclude_agent_ids:
        exclude_agent_ids_list = exclude_agent_ids.split(",") if exclude_agent_ids else None

    # Without scoping, an empty `agent_ids` query returned every agent's
    # spend/token rows on the proxy. Restrict non-admin callers to the
    # agents they're permitted to invoke (or that they created), and
    # intersect their explicit `agent_ids` filter with the same allowlist.
    from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
        AgentRequestHandler,
        RestrictedAgentAccess,
        UnrestrictedAgentAccess,
    )
    from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view

    where_condition: Final[dict[str, object]] = {}
    if not _user_has_admin_view(user_api_key_dict):
        permitted_agent_ids: list[str] = []
        # An unrestricted caller is not "see everything" for activity scoping. Fall
        # back to the agents the caller created so they cannot enumerate other
        # tenants' agents.
        # Guard against `user_id is None`: a literal None in Prisma
        # `where={"created_by": None}` resolves to ``created_by IS NULL``
        # and would expose every ownerless agent's rows.
        match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict):
            case RestrictedAgentAccess(allowed_agent_ids):
                permitted_agent_ids = list(allowed_agent_ids)
            case UnrestrictedAgentAccess():
                if user_api_key_dict.user_id is not None:
                    owned_records: Final = await agents_table(prisma_client).find_many(
                        where={"created_by": user_api_key_dict.user_id}
                    )
                    permitted_agent_ids = [a.agent_id for a in owned_records]

        if agent_ids_list:
            permitted_agent_id_set: Final = set(permitted_agent_ids)
            agent_ids_list = [aid for aid in agent_ids_list if aid in permitted_agent_id_set]
        else:
            agent_ids_list = list(permitted_agent_ids)

        # No accessible agents → return an empty page without querying.
        if not agent_ids_list:
            return SpendAnalyticsPaginatedResponse(
                results=[],
                metadata=DailySpendMetadata(
                    total_spend=0.0,
                    total_prompt_tokens=0,
                    total_completion_tokens=0,
                    total_tokens=0,
                    total_api_requests=0,
                    total_successful_requests=0,
                    total_failed_requests=0,
                    total_cache_read_input_tokens=0,
                    total_cache_creation_input_tokens=0,
                    total_compression_saved_tokens=0,
                    page=page,
                    total_pages=0,
                    has_more=False,
                ),
            )

    if agent_ids_list:
        where_condition["agent_id"] = {"in": list(agent_ids_list)}

    agent_records: Final = await agents_table(prisma_client).find_many(where=where_condition)
    agent_metadata: Final[Mapping[str, dict[str, object]]] = {
        agent.agent_id: {"agent_name": agent.agent_name} for agent in agent_records
    }

    return await get_daily_activity(
        prisma_client=prisma_client,
        table_name="litellm_dailyagentspend",
        entity_id_field="agent_id",
        entity_id=agent_ids_list,
        entity_metadata_field=agent_metadata,
        exclude_entity_ids=exclude_agent_ids_list,
        start_date=start_date,
        end_date=end_date,
        model=model,
        api_key=api_key,
        page=page,
        page_size=page_size,
    )
