from __future__ import annotations as _annotations

import os
from typing import overload

from pydantic_ai import ModelProfile
from pydantic_ai._json_schema import JsonSchema, JsonSchemaTransformer
from pydantic_ai.exceptions import UserError
from pydantic_ai.native_tools import SUPPORTED_NATIVE_TOOLS
from pydantic_ai.profiles import merge_profile
from pydantic_ai.profiles.amazon import amazon_model_profile
from pydantic_ai.profiles.anthropic import anthropic_model_profile
from pydantic_ai.profiles.cohere import cohere_model_profile
from pydantic_ai.profiles.deepseek import deepseek_model_profile
from pydantic_ai.profiles.google import google_model_profile
from pydantic_ai.profiles.grok import grok_model_profile
from pydantic_ai.profiles.meta import meta_model_profile
from pydantic_ai.profiles.mistral import mistral_model_profile
from pydantic_ai.profiles.moonshotai import moonshotai_model_profile
from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile, openai_model_profile
from pydantic_ai.profiles.qwen import qwen_model_profile

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


class OpenRouterModelProfile(OpenAIModelProfile, total=False):
    """Profile for models used with OpenRouterModel.

    ALL FIELDS MUST BE `openrouter_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
    """

    openrouter_supports_cache_control: bool
    """Whether the downstream provider supports explicit `cache_control` breakpoints via OpenRouter."""
    openrouter_supports_cache_ttl: bool
    """Whether the downstream provider supports TTL in `cache_control`."""
    openrouter_supports_tool_cache: bool
    """Whether the downstream provider supports `cache_control` on tool definitions."""
    openrouter_supports_dynamic_instruction_cache: bool
    """Whether instruction cache boundaries can exclude later dynamic instruction blocks."""
    openrouter_max_cache_points: int | None
    """Maximum number of `cache_control` breakpoints the downstream provider allows per request.

    Anthropic enforces a limit of 4. When set, excess breakpoints are silently removed
    from messages (newest kept first). `None` means no limit."""
    openrouter_supports_forced_tool_choice_with_thinking: bool
    """Whether the downstream provider accepts a forced `tool_choice` while thinking is enabled.

    Anthropic rejects `tool_choice` `any`/`tool` alongside extended thinking, but OpenRouter swallows the
    incompatibility by dropping `reasoning` from the request instead of erroring, so the response silently
    comes back with no reasoning at all. When False and thinking is enabled, a resolved `required` tool
    choice falls back to `auto` (filtering tools to the requested set), and an explicit
    `tool_choice='required'` (or an explicit list of tools) raises a `UserError`."""


class _OpenRouterGoogleJsonSchemaTransformer(JsonSchemaTransformer):
    """Legacy Google JSON schema transformer for OpenRouter compatibility.

    OpenRouter's compatibility layer doesn't fully support modern JSON Schema features
    like $defs/$ref and anyOf for nullable types. This transformer restores v1.19.0
    behavior by inlining definitions and simplifying nullable unions.

    See: https://github.com/pydantic/pydantic-ai/issues/3617
    """

    def __init__(self, schema: JsonSchema, *, strict: bool | None = None):
        super().__init__(schema, strict=strict, prefer_inlined_defs=True, simplify_nullable_unions=True)

    def transform(self, schema: JsonSchema) -> JsonSchema:
        # Remove properties not supported by Gemini
        schema.pop('$schema', None)
        schema.pop('title', None)
        schema.pop('discriminator', None)
        schema.pop('examples', None)
        schema.pop('exclusiveMaximum', None)
        schema.pop('exclusiveMinimum', None)

        if (const := schema.pop('const', None)) is not None:
            schema['enum'] = [const]

        # Convert enums to string type (legacy Gemini requirement)
        if enum := schema.get('enum'):
            schema['type'] = 'string'
            schema['enum'] = [str(val) for val in enum]

        # Convert oneOf to anyOf for discriminated unions
        if 'oneOf' in schema and 'type' not in schema:
            schema['anyOf'] = schema.pop('oneOf')

        # Handle string format -> description
        type_ = schema.get('type')
        if type_ == 'string' and (fmt := schema.pop('format', None)):
            description = schema.get('description')
            if description:
                schema['description'] = f'{description} (format: {fmt})'
            else:
                schema['description'] = f'Format: {fmt}'

        return schema


def _openrouter_google_model_profile(model_name: str) -> ModelProfile | None:
    """Get the model profile for a Google model accessed via OpenRouter.

    Uses the legacy transformer to maintain compatibility with OpenRouter's
    translation layer, which doesn't fully support modern JSON Schema features.
    """
    profile = google_model_profile(model_name)
    if profile is None:  # pragma: no cover
        return None
    return merge_profile(profile, ModelProfile(json_schema_transformer=_OpenRouterGoogleJsonSchemaTransformer))


class OpenRouterProvider(_OpenAICompatibleProvider):
    """Provider for OpenRouter API."""

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

    @property
    def base_url(self) -> str:
        return 'https://openrouter.ai/api/v1'

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

    @staticmethod
    def model_profile(model_name: str) -> ModelProfile | None:
        provider_to_profile = {
            'google': _openrouter_google_model_profile,
            'openai': openai_model_profile,
            'anthropic': anthropic_model_profile,
            'mistralai': mistral_model_profile,
            'qwen': qwen_model_profile,
            'x-ai': grok_model_profile,
            'cohere': cohere_model_profile,
            'amazon': amazon_model_profile,
            'deepseek': deepseek_model_profile,
            'meta-llama': meta_model_profile,
            'moonshotai': moonshotai_model_profile,
        }

        profile = None

        # OpenRouter identifies models as `provider/model`.
        if '/' not in model_name:
            raise UserError(
                f'OpenRouter model names must be prefixed with the upstream provider, e.g. '
                f'{("openai/" + model_name)!r}, not {model_name!r}. '
                'See https://openrouter.ai/models for the available model names.'
            )

        # OpenRouter exposes latest-model aliases as `~provider/model`; strip the
        # alias marker before using the provider prefix for profile selection.
        provider, model_name = model_name.removeprefix('~').split('/', 1)
        if provider in provider_to_profile:
            model_name, *_ = model_name.split(':', 1)  # drop tags
            if provider == 'anthropic':
                model_name = model_name.replace('.', '-')
            profile = provider_to_profile[provider](model_name)

        # Cache capability flags are set on the gateway layer based on the downstream provider.
        # The TTL / tool-cache / dynamic-instruction flags are kept separate even though they all
        # coincide with `supports_anthropic_cache` today: they model independent OpenRouter cache
        # capabilities that merely happen to line up on the current Anthropic-only provider set, so a
        # future non-Anthropic downstream can enable any of them independently without re-coupling them.
        supports_cache_control = provider in ('anthropic', 'google')
        supports_anthropic_cache = provider == 'anthropic'

        # Three-layer merge:
        # 1. Fallback layer — `OpenAIJsonSchemaTransformer` is the default unless an upstream profile sets one explicitly
        #    (e.g. `_openrouter_google_model_profile` installs `_OpenRouterGoogleJsonSchemaTransformer`).
        # 2. Upstream profile — model-specific traits from the lab's profile function.
        # 3. Gateway-specific overrides — wins on every key it sets, because the upstream profile can't know what
        #    the OpenRouter gateway adds (web plugin, file URLs, custom thinking field, cache capabilities). OpenRouter
        #    accepts `reasoning` universally, so the gate also forces `supports_thinking=True` so the unified `thinking`
        #    setting is always forwarded regardless of the upstream model's own thinking support. OpenRouter only
        #    accepts the older `max_tokens` field, so `openai_chat_supports_max_completion_tokens=False`.
        #    It also silently transforms mid-conversation system messages rather than handling them natively,
        #    so `supports_inline_system_prompts=False` selects the deliberate `<system>`-wrapped fallback.
        return merge_profile(
            OpenAIModelProfile(json_schema_transformer=OpenAIJsonSchemaTransformer),
            profile,
            OpenRouterModelProfile(
                openai_chat_send_back_thinking_parts='field',
                openai_chat_thinking_field='reasoning',
                openai_chat_supports_file_urls=True,
                openai_chat_supports_web_search=True,
                openai_chat_supports_max_completion_tokens=False,
                supports_inline_system_prompts=False,
                supports_thinking=True,
                # OpenRouter's native tools (web search plugin, advisor) are gateway features that
                # work with any underlying model, so the upstream profile's vendor-specific tool
                # gating (e.g. Anthropic's valid-executor list) doesn't apply. Neutralize it here;
                # `OpenRouterModel.supported_native_tools()` caps the effective set via the
                # intersection in `Model.profile`.
                supported_native_tools=SUPPORTED_NATIVE_TOOLS,
                openrouter_supports_cache_control=supports_cache_control,
                openrouter_supports_cache_ttl=supports_anthropic_cache,
                openrouter_supports_tool_cache=supports_anthropic_cache,
                openrouter_supports_dynamic_instruction_cache=supports_anthropic_cache,
                openrouter_max_cache_points=4 if supports_anthropic_cache else None,
                # Anthropic errors on a forced `tool_choice` with thinking enabled; OpenRouter instead
                # drops `reasoning` from the request and returns a response with no reasoning at all.
                # https://platform.claude.com/docs/en/agents-and-tools/tool-use/implement-tool-use#forcing-tool-use
                openrouter_supports_forced_tool_choice_with_thinking=provider != 'anthropic',
            ),
        )

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

    @overload
    def __init__(
        self,
        *,
        api_key: str | None = None,
        app_url: str | None = None,
        app_title: str | None = None,
        openai_client: None = None,
        http_client: _OpenAIHTTPClient | None = None,
    ) -> None: ...

    def __init__(
        self,
        *,
        api_key: str | None = None,
        app_url: str | None = None,
        app_title: str | None = None,
        openai_client: AsyncOpenAI | None = None,
        http_client: _OpenAIHTTPClient | None = None,
    ) -> None:
        """Configure the provider with either an API key or prebuilt client.

        Args:
            api_key: OpenRouter API key. Falls back to `OPENROUTER_API_KEY`
                when omitted and required unless `openai_client` is provided.
            app_url: Optional url for app attribution. Falls back to
                `OPENROUTER_APP_URL` when omitted.
            app_title: Optional title for app attribution. Falls back to
                `OPENROUTER_APP_TITLE` when omitted.
            openai_client: Existing `AsyncOpenAI` client to reuse instead of
                creating one internally.
            http_client: Custom `httpx2.AsyncClient` or legacy `httpx.AsyncClient` to pass into the
                `AsyncOpenAI` constructor when building a client.

        Raises:
            UserError: If no API key is available and no `openai_client` is
                provided.
        """
        api_key = api_key or os.getenv('OPENROUTER_API_KEY')
        if not api_key and openai_client is None:
            raise UserError(
                'Set the `OPENROUTER_API_KEY` environment variable or pass it via `OpenRouterProvider(api_key=...)`'
                ' to use the OpenRouter provider.'
            )

        attribution_headers: dict[str, str] = {}
        if http_referer := app_url or os.getenv('OPENROUTER_APP_URL'):
            attribution_headers['HTTP-Referer'] = http_referer
        if x_title := app_title or os.getenv('OPENROUTER_APP_TITLE'):
            attribution_headers['X-Title'] = x_title

        if openai_client is not None:
            self._client = openai_client
        else:
            self._client = self._create_openai_client(
                base_url=self.base_url, api_key=api_key, http_client=http_client, default_headers=attribution_headers
            )