from __future__ import annotations as _annotations

import os
from typing import Literal, overload

from typing_extensions import assert_never

from pydantic_ai import ModelProfile
from pydantic_ai.exceptions import UserError
from pydantic_ai.profiles import merge_profile
from pydantic_ai.profiles.openai import OpenAIModelProfile, openai_model_profile
from pydantic_ai.providers._bedrock_model_names import split_bedrock_model_id

try:
    from openai import AsyncBedrockOpenAI, AsyncOpenAI
except ImportError as _import_error:
    raise ImportError(
        'Please install the Bedrock Mantle dependencies to use the Bedrock Mantle provider, '
        'you can use the `bedrock-mantle` optional group — `pip install "pydantic-ai-slim[bedrock-mantle]"`'
    ) from _import_error
else:
    from ._openai_compatible import (
        AsyncHTTPClient as _OpenAIHTTPClient,
        OpenAICompatibleProvider as _OpenAICompatibleProvider,
    )

BedrockMantleInterface = Literal['chat', 'responses', 'openai-responses']
"""The OpenAI-compatible endpoint family a Bedrock Mantle model is served on.

- `'chat'`: Chat Completions at `/v1/chat/completions` (GPT-OSS Safeguard).
- `'responses'`: Responses at `/v1/responses` (GPT-OSS).
- `'openai-responses'`: Responses at `/openai/v1/responses`, the OpenAI-model-specific path (GPT-5.4+).
"""


class BedrockMantleModelProfile(OpenAIModelProfile, total=False):
    """Profile for OpenAI models served through Amazon Bedrock Mantle."""

    bedrock_mantle_interface: BedrockMantleInterface
    """Which Mantle endpoint family serves this model, selecting the model class and base URL."""


def bedrock_mantle_model_profile(model_name: str) -> ModelProfile:
    """Resolve the profile for an OpenAI model served through Bedrock Mantle."""
    provider, base_model_name = split_bedrock_model_id(model_name)
    if provider != 'openai':
        raise UserError(
            f'Model {model_name!r} is not an OpenAI model on Bedrock Mantle. '
            'Bedrock Mantle currently serves OpenAI models through the `bedrock-mantle:` prefix.'
        )
    if base_model_name.startswith('gpt-oss-safeguard'):
        interface: BedrockMantleInterface = 'chat'
    elif base_model_name.startswith('gpt-oss'):
        interface = 'responses'
    else:
        interface = 'openai-responses'
    # Mantle GPT-5.x on the `/openai/v1` Responses endpoint resets tool-call IDs to `call_0` across
    # separate responses (verified live on 5.5 and 5.6); qualify them with the response ID at ingestion
    # so pydantic-ai keeps its history-wide-unique invariant. GPT-OSS on `/v1/responses` is unaffected
    # (globally-unique IDs), so this keys on the endpoint family rather than the model version — the
    # reset follows the endpoint, and a per-version allowlist would silently miss future GPT-5.x. See #6536.
    response_scoped = interface == 'openai-responses'
    return merge_profile(
        openai_model_profile(base_model_name),
        BedrockMantleModelProfile(
            bedrock_mantle_interface=interface,
            openai_responses_tool_call_ids_are_response_scoped=response_scoped,
            # Bedrock Mantle does not serve image output for these models, unlike the direct OpenAI API
            # the base profile is resolved from (per the AWS model cards). Without this override the
            # image-output guard is bypassed and the request fails with an opaque provider error.
            supports_image_output=False,
            # AWS's server-side tool integrations use Lambda or MCP rather than the OpenAI-native tools
            # represented by Pydantic AI (web search, code interpreter, file search, image generation, ...):
            # https://docs.aws.amazon.com/bedrock/latest/userguide/tool-use-server-side.html
            # Disable the base profile's native tools to fail with a clean UserError rather than an opaque
            # provider error.
            supported_native_tools=frozenset(),
        ),
    )


def _mantle_origin(base_url: str) -> str:
    """Return the scheme+host origin of a Mantle base URL, stripping a trailing `/openai/v1` or `/v1`.

    Mantle serves its two OpenAI-compatible endpoint families as siblings under one origin, so the origin
    is enough to derive both `{origin}/v1` and `{origin}/openai/v1` regardless of which one a caller passed.
    """
    origin = base_url.rstrip('/')
    for suffix in ('/openai/v1', '/v1'):
        if origin.endswith(suffix):
            return origin[: -len(suffix)]
    return origin


class BedrockMantleProvider(_OpenAICompatibleProvider):
    """Provider for the Amazon Bedrock Mantle OpenAI-compatible API."""

    @property
    def name(self) -> str:
        return 'bedrock-mantle'

    @property
    def base_url(self) -> str:
        return str(self._client.base_url)

    @property
    def client(self) -> AsyncOpenAI:
        return self._client

    @staticmethod
    def model_profile(model_name: str) -> ModelProfile | None:
        return bedrock_mantle_model_profile(model_name)

    def _client_for_interface(self, interface: BedrockMantleInterface) -> AsyncOpenAI:
        """Return the client for a Mantle interface.

        The `openai-responses` interface (GPT-5.x) is served at `/openai/v1`; `chat` and `responses`
        (GPT-OSS) are served at `/v1`. Both clients are built eagerly in `__init__` and share transport
        and auth, so this is a pure lookup.
        """
        if interface == 'openai-responses':
            return self._client
        elif interface in ('chat', 'responses'):
            return self._v1_client
        else:
            assert_never(interface)

    @overload
    def __init__(self, *, openai_client: AsyncOpenAI) -> None: ...

    @overload
    def __init__(
        self,
        *,
        region_name: str | None = None,
        base_url: str | None = None,
        api_key: str | None = None,
        aws_access_key_id: str | None = None,
        aws_secret_access_key: str | None = None,
        aws_session_token: str | None = None,
        profile_name: str | None = None,
        http_client: _OpenAIHTTPClient | None = None,
    ) -> None: ...

    def __init__(
        self,
        *,
        region_name: str | None = None,
        base_url: str | None = None,
        api_key: str | None = None,
        aws_access_key_id: str | None = None,
        aws_secret_access_key: str | None = None,
        aws_session_token: str | None = None,
        profile_name: str | None = None,
        openai_client: AsyncOpenAI | None = None,
        http_client: _OpenAIHTTPClient | None = None,
    ) -> None:
        """Create a Bedrock Mantle provider.

        Args:
            region_name: The AWS region used to construct the default `bedrock-mantle.{region}.api.aws`
                origin. If not set, the `AWS_DEFAULT_REGION` or `AWS_REGION` environment variable is used.
            base_url: A Mantle base URL. Its origin (scheme + host, with any `/openai/v1` or `/v1` suffix
                stripped) is used to route between `/v1` and `/openai/v1` per model, the same as `region_name`.
            api_key: A Bedrock API key. If omitted, `AWS_BEARER_TOKEN_BEDROCK` is used. Use this or the
                `aws_*` credentials, not both.
            aws_access_key_id: The AWS access key ID for SigV4 authentication.
            aws_secret_access_key: The AWS secret access key for SigV4 authentication.
            aws_session_token: The AWS session token for SigV4 authentication.
            profile_name: The AWS profile name for SigV4 authentication.
            openai_client: An existing OpenAI client. If provided, no other argument may be set; its base
                URL's origin is used to derive both the `/v1` and `/openai/v1` endpoints (preserving its
                auth and transport) so every interface routes correctly.
            http_client: An existing `httpx2.AsyncClient` or legacy `httpx.AsyncClient` used to make requests.
        """
        if openai_client is not None:
            assert region_name is None, 'Cannot provide both `openai_client` and `region_name`'
            assert base_url is None, 'Cannot provide both `openai_client` and `base_url`'
            assert api_key is None, 'Cannot provide both `openai_client` and `api_key`'
            assert http_client is None, 'Cannot provide both `openai_client` and `http_client`'
            assert (aws_access_key_id, aws_secret_access_key, aws_session_token, profile_name) == (
                None,
                None,
                None,
                None,
            ), 'Cannot provide both `openai_client` and AWS credentials'
            base_client = openai_client
            origin = _mantle_origin(str(openai_client.base_url))
        else:
            api_key = api_key or os.getenv('AWS_BEARER_TOKEN_BEDROCK')
            region_name = region_name or os.getenv('AWS_DEFAULT_REGION') or os.getenv('AWS_REGION')

            if base_url is not None:
                origin = _mantle_origin(base_url)
            elif region_name:
                origin = f'https://bedrock-mantle.{region_name}.api.aws'
            else:
                raise UserError(
                    'Set the `AWS_DEFAULT_REGION` or `AWS_REGION` environment variable, pass `region_name`, '
                    'or pass `base_url` to use a Bedrock Mantle model.'
                )

            http_client = self._get_http_client(http_client)

            base_client = AsyncBedrockOpenAI(
                api_key=api_key,
                aws_access_key_id=aws_access_key_id,
                aws_secret_access_key=aws_secret_access_key,
                aws_session_token=aws_session_token,
                aws_profile=profile_name,
                aws_region=region_name,
                base_url=f'{origin}/openai/v1',
                http_client=http_client,  # pyright: ignore[reportArgumentType]
            )

        # `_client` (the default `.client`/`.base_url`) serves the `openai-responses` interface at
        # `/openai/v1` — the endpoint for GPT-5.x, the primary Mantle models. `_v1_client` serves `chat`
        # and `responses` (GPT-OSS) at `/v1`. Both derive from one base client via `with_options`, so they
        # share transport and auth; deriving from the origin means a user-supplied `base_url`/`openai_client`
        # still routes both interfaces correctly regardless of which sibling endpoint it named.
        self._client = base_client.with_options(base_url=f'{origin}/openai/v1')
        self._v1_client = base_client.with_options(base_url=f'{origin}/v1')

    def _set_http_client(self, http_client: _OpenAIHTTPClient) -> None:
        for client in (self._client, self._v1_client):
            client._client = http_client  # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue]
