"""API routes for the web chat UI."""

from collections.abc import Mapping, Sequence
from typing import Literal, TypeVar

from pydantic import BaseModel
from pydantic.alias_generators import to_camel
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import JSONResponse, Response
from starlette.routing import Route

from pydantic_ai import Agent
from pydantic_ai.capabilities import NativeTool
from pydantic_ai.exceptions import UserError
from pydantic_ai.models import KnownModelName, Model, infer_model
from pydantic_ai.native_tools import SUPPORTED_NATIVE_TOOLS, AbstractNativeTool
from pydantic_ai.settings import ModelSettings
from pydantic_ai.ui.vercel_ai import VercelAIAdapter

AgentDepsT = TypeVar('AgentDepsT')
OutputDataT = TypeVar('OutputDataT')

# Type alias for models parameter - accepts model names/instances or a dict mapping labels to models
ModelsParam = Sequence[Model | KnownModelName | str] | Mapping[str, Model | KnownModelName | str] | None

# The bundled chat UI ships v7 of the Vercel AI SDK, so the API path targets `sdk_version=7` to match
# it. v7's data-stream protocol equals v6's (same wire, including the tool-approval chunks that enable
# tool-approval streaming), so 7 emits identically to 6 today; targeting 7 keeps the value aligned with
# the UI's real SDK major and reserves it for future v7-only chunks. `to_web()` controls both ends
# (server + bundled UI). See `VercelAIAdapter.sdk_version`.
BUNDLED_UI_SDK_VERSION: Literal[7] = 7

# `/chat` requires exactly this content type. This is a CSRF control, not content negotiation: a
# browser can send the three CORS-safelisted content types (`text/plain`, `multipart/form-data`,
# `application/x-www-form-urlencoded`) — or no content type at all — cross-origin with no preflight,
# and every one of them can carry a JSON body. Without this check, a page the developer happens to
# visit while running the web UI could start an agent run on their machine, and execute whatever
# tools the served agent exposes. `application/json` is not safelisted, so requiring it forces a
# preflight, which `options_chat` refuses.
#
# Keep this an allowlist. A denylist of the safelisted types would miss the no-content-type case,
# and would silently stop covering anything added to the safelist later.
#
# The bundled UI — and the Vercel AI SDK's `DefaultChatTransport` generally — always sends
# `application/json`; other clients need to set the header explicitly.
JSON_MEDIA_TYPE = 'application/json'


class ModelInfo(BaseModel, alias_generator=to_camel, populate_by_name=True):
    """Defines an AI model with its associated built-in tools."""

    id: str
    name: str
    builtin_tools: list[str]


class BuiltinToolInfo(BaseModel, alias_generator=to_camel, populate_by_name=True):
    """Serializable info about a builtin tool for frontend config."""

    id: str
    name: str


class ConfigureFrontend(BaseModel, alias_generator=to_camel, populate_by_name=True):
    """Response model for frontend configuration."""

    models: list[ModelInfo]
    builtin_tools: list[BuiltinToolInfo]


class ChatRequestExtra(BaseModel, extra='ignore', alias_generator=to_camel):
    """Extra data extracted from chat request."""

    model: str | None = None
    """Model ID selected by the user, e.g. 'openai:gpt-5'. Maps to JSON field 'model'."""
    builtin_tools: list[str] = []
    """Tool IDs selected by the user, e.g. ['web_search', 'code_execution']. Maps to JSON field 'builtinTools'."""


def validate_request_options(
    extra_data: ChatRequestExtra,
    model_ids: set[str],
    builtin_tool_ids: set[str],
) -> str | None:
    """Validate that requested model and tools are in the allowed lists.

    Returns an error message if validation fails, or None if valid.
    """
    if extra_data.model and extra_data.model not in model_ids:
        return f'Model "{extra_data.model}" is not in the allowed models list'

    # base model also validates this but makes sense to have an api check, since one could be a UI bug/misbehavior
    # the other would be a pydantic-ai bug
    # also as future proofing since we don't know how users will use this feature in the future
    invalid_tools = [t for t in extra_data.builtin_tools if t not in builtin_tool_ids]
    if invalid_tools:
        return f'Builtin tool(s) {invalid_tools} not in the allowed tools list'

    return None


def create_api_app(
    agent: Agent[AgentDepsT, OutputDataT],
    models: ModelsParam = None,
    native_tools: Sequence[AbstractNativeTool] | None = None,
    deps: AgentDepsT = None,
    model_settings: ModelSettings | None = None,
    instructions: str | None = None,
    sdk_version: Literal[5, 6, 7] = BUNDLED_UI_SDK_VERSION,
) -> Starlette:
    """Create API app for the web chat UI.

    This sub-app carries no `Host` validation of its own: `create_web_app()` applies
    `HostValidationMiddleware` to the outer app it mounts this into. Serving it directly, or
    re-mounting `create_web_app(...).router` (which drops application middleware), exposes the chat
    endpoint to any `Host` — mount the app `create_web_app()` returns instead.

    Args:
        agent: Agent instance.
        models: Models to make available in the UI. Can be:
            - A sequence of model names/instances (e.g., `['openai:gpt-5', Model(...)]`)
            - A dict mapping display labels to model names/instances
            If not provided, the UI will have no model options.
        native_tools: Optional list of additional native tools to make available in the UI.
            Tools already configured on the agent are always included but won't appear as options.
        deps: Optional dependencies to use for all requests.
        model_settings: Optional settings to use for all model requests.
        instructions: Optional extra instructions to pass to each agent run.
        sdk_version: Vercel AI SDK version to target on the chat endpoint: 5, 6, or 7. Defaults to
            `7` to match the bundled v7 UI, which enables tool-approval streaming (7 emits the same
            wire as 6, since v7's data-stream protocol equals v6's).

    Returns:
        A Starlette application with the API endpoints.
    """
    # Build model ID → original reference mapping and ModelInfo list for frontend
    model_id_to_ref: dict[str, Model | str] = {}
    model_infos: list[ModelInfo] = []

    # Filter out tools that are already configured on the agent (they're always included)
    agent_tool_ids = {t.unique_id for t in agent._cap_native_tools if isinstance(t, AbstractNativeTool)}  # pyright: ignore[reportPrivateUsage]
    ui_native_tools = [t for t in (native_tools or []) if t.unique_id not in agent_tool_ids]

    # Build combined models: agent's model first (if exists), then provided models
    all_models: list[tuple[str | None, Model | str]] = []
    if agent.model is not None:
        all_models.append((None, agent.model))
    items = list(models.items()) if isinstance(models, Mapping) else [(None, m) for m in (models or [])]
    all_models.extend(items)

    seen_model_ids: set[str] = set()
    for label, model_ref in all_models:
        try:
            model = infer_model(model_ref)
        except UserError:
            # A capability resolver may intentionally use an ID that built-in inference does
            # not understand. Resolution needs run dependencies, so leave custom references
            # untouched here and let Agent resolve them when the request is dispatched.
            if agent._root_capability.has_resolve_model_id:  # pyright: ignore[reportPrivateUsage]
                model = None
            else:
                raise
        # Use original string if provided to preserve openai-chat: vs openai-responses: distinction
        model_id = model_ref if isinstance(model_ref, str) else model_ref.model_id
        if model_id in seen_model_ids:
            continue
        seen_model_ids.add(model_id)
        display_name = label or (model.label if model is not None else model_id)
        if model is None:
            supported_tool_ids = []
        else:
            model_supported_tools = model.profile.get('supported_native_tools', SUPPORTED_NATIVE_TOOLS)
            supported_tool_ids = [t.unique_id for t in ui_native_tools if type(t) in model_supported_tools]

        model_id_to_ref[model_id] = model_ref
        model_infos.append(ModelInfo(id=model_id, name=display_name, builtin_tools=supported_tool_ids))

    model_ids = set(model_id_to_ref.keys())
    allowed_tool_ids = {tool.unique_id for tool in ui_native_tools}

    async def options_chat(request: Request) -> Response:
        """Answer CORS preflight requests without granting cross-origin access.

        Deliberately carries no `Access-Control-Allow-*` header, so a browser refuses any
        cross-origin request that needs a preflight. Together with the `application/json`
        requirement in `post_chat()`, that is what keeps a page the developer visits from reaching
        the local web UI: requiring a non-safelisted content type forces the preflight, and this
        response denies it.

        Do not add `Access-Control-Allow-Origin` here without an accompanying CSRF control — on its
        own it re-opens the endpoint to any website the developer has open.
        """
        return Response()

    async def configure_frontend(request: Request) -> Response:
        """Endpoint to configure the frontend with available models and tools."""
        config = ConfigureFrontend(
            models=model_infos,
            builtin_tools=[BuiltinToolInfo(id=tool.unique_id, name=tool.label) for tool in ui_native_tools],
        )
        return JSONResponse(config.model_dump(by_alias=True))

    async def health(request: Request) -> Response:
        """Health check endpoint."""
        return JSONResponse({'ok': True})

    async def post_chat(request: Request) -> Response:
        """Handle chat requests via Vercel AI Adapter."""
        if (media_type := request.headers.get('content-type', '').split(';')[0].strip().lower()) != JSON_MEDIA_TYPE:
            # Checked up front so a cross-origin-forgeable request is turned away before the body is
            # parsed and before the agent is dispatched — the attacker never needs to read the
            # response, so the run itself is the damage.
            return JSONResponse(
                {'error': f'Expected `Content-Type: {JSON_MEDIA_TYPE}`, got {media_type or "no content type"}'},
                status_code=415,
            )

        adapter = await VercelAIAdapter[AgentDepsT, OutputDataT].from_request(
            request, agent=agent, sdk_version=sdk_version
        )
        extra_data = ChatRequestExtra.model_validate(adapter.run_input.__pydantic_extra__)

        if error := validate_request_options(extra_data, model_ids, allowed_tool_ids):
            return JSONResponse({'error': error}, status_code=400)

        model_ref = model_id_to_ref.get(extra_data.model) if extra_data.model else None
        request_native_tools = [tool for tool in ui_native_tools if tool.unique_id in extra_data.builtin_tools]
        request_capabilities: list[NativeTool[AgentDepsT]] = [NativeTool(tool) for tool in request_native_tools]
        streaming_response = await VercelAIAdapter[AgentDepsT, OutputDataT].dispatch_request(
            request,
            agent=agent,
            sdk_version=sdk_version,
            model=model_ref,
            capabilities=request_capabilities,
            deps=deps,
            model_settings=model_settings,
            instructions=instructions,
        )
        return streaming_response

    routes = [
        Route('/chat', options_chat, methods=['OPTIONS']),
        Route('/chat', post_chat, methods=['POST']),
        Route('/configure', configure_frontend, methods=['GET']),
        Route('/health', health, methods=['GET']),
    ]
    return Starlette(routes=routes)