"""
Convenience ``wrap()`` function for dependency injection cases.

When the customer has already created a provider client, ``wrap()``
extracts credentials and returns an instrumented SDK wrapper::

    from openai import OpenAI
    client = OpenAI(api_key="sk-...")
    wrapped = amplitude_ai.wrap(client, amplitude=amp, user_id="u1")
    # wrapped is a real amplitude_ai.OpenAI instance

This is credential extraction, not proxying.  The returned object is a
genuine wrapper class instance -- ``isinstance()`` works, destructuring
works, no proxy magic.

Supported providers (Phase 1): OpenAI, Anthropic, AzureOpenAI.
Unsupported types raise :class:`AmplitudeAIWrapError`.
"""

from __future__ import annotations

from typing import Any

from amplitude import Amplitude

from .exceptions import AmplitudeAIError


class AmplitudeAIWrapError(AmplitudeAIError):
    """Raised when ``wrap()`` receives an unsupported client type."""


def wrap(
    client: Any,
    amplitude: Amplitude,
    user_id: str | None = None,
    **kwargs: Any,
) -> Any:
    """Wrap an existing provider client with Amplitude instrumentation.

    Extracts credentials from *client* and returns a real SDK wrapper
    instance.

    Args:
        client: An existing provider client (e.g., ``openai.OpenAI()``,
            ``anthropic.Anthropic()``, ``openai.AzureOpenAI()``).
        amplitude: Amplitude client instance.
        user_id: Default user ID for tracked events.
        **kwargs: Additional kwargs forwarded to the wrapper constructor
            (e.g., ``privacy_mode=True``).

    Returns:
        An instrumented wrapper instance of the matching provider type.

    Raises:
        AmplitudeAIWrapError: If *client* is not a recognized provider type.
    """
    client_type = type(client)
    module = client_type.__module__ or ""

    # --- OpenAI ---
    try:
        import openai as openai_mod

        if isinstance(client, openai_mod.AzureOpenAI):
            from .providers.azure_openai import AzureOpenAI as AmpAzureOpenAI

            wrapper = AmpAzureOpenAI(
                amplitude=amplitude,
                api_key=getattr(client, "api_key", None),
                azure_endpoint=str(getattr(client, "base_url", "")),
                api_version=getattr(client, "_custom_query", {}).get("api-version"),
                **kwargs,
            )
            if user_id:
                wrapper._default_amplitude_user_id = user_id
            return wrapper

        # AsyncAzureOpenAI extends AsyncOpenAI — check it first to avoid
        # silently wrapping Azure clients as plain AsyncOpenAI.
        if hasattr(openai_mod, "AsyncAzureOpenAI") and isinstance(
            client, openai_mod.AsyncAzureOpenAI
        ):
            raise AmplitudeAIWrapError(
                "wrap() does not yet support openai.AsyncAzureOpenAI. "
                "Use the synchronous openai.AzureOpenAI client with "
                "amplitude_ai.wrap() or amplitude_ai.patch_azure_openai() instead. "
                "Async Azure OpenAI support is not yet available."
            )

        if isinstance(client, openai_mod.AsyncOpenAI):
            from .providers.openai import AsyncOpenAI as AmpAsyncOpenAI

            init_kwargs: dict[str, Any] = {
                "amplitude": amplitude,
                "api_key": getattr(client, "api_key", None),
            }
            base_url = getattr(client, "base_url", None)
            if base_url is not None:
                init_kwargs["base_url"] = str(base_url)
            init_kwargs.update(kwargs)
            wrapper = AmpAsyncOpenAI(**init_kwargs)
            if user_id:
                wrapper._default_amplitude_user_id = user_id
            return wrapper

        if isinstance(client, openai_mod.OpenAI):
            from .providers.openai import OpenAI as AmpOpenAI

            init_kwargs = {
                "amplitude": amplitude,
                "api_key": getattr(client, "api_key", None),
            }
            base_url = getattr(client, "base_url", None)
            if base_url is not None:
                init_kwargs["base_url"] = str(base_url)
            init_kwargs.update(kwargs)
            wrapper = AmpOpenAI(**init_kwargs)
            if user_id:
                wrapper._default_amplitude_user_id = user_id
            return wrapper
    except ImportError:
        pass

    # --- Anthropic ---
    try:
        import anthropic as anthropic_mod

        if isinstance(client, anthropic_mod.Anthropic):
            from .providers.anthropic import Anthropic as AmpAnthropic

            wrapper = AmpAnthropic(
                amplitude=amplitude,
                api_key=getattr(client, "api_key", None),
                **kwargs,
            )
            if user_id:
                wrapper._default_amplitude_user_id = user_id
            return wrapper
    except ImportError:
        pass

    # --- Unsupported ---
    supported = "openai.OpenAI, openai.AsyncOpenAI, openai.AzureOpenAI, anthropic.Anthropic"
    full_name = f"{client_type.__module__}.{client_type.__qualname__}"

    hint = ""
    if "openai" in (client_type.__module__ or ""):
        hint = (
            " Did you mean patch_openai()? For zero-code instrumentation of the "
            "OpenAI module, use amplitude_ai.patch_openai(amplitude=...) instead."
        )
    elif "anthropic" in (client_type.__module__ or ""):
        hint = (
            " Did you mean patch_anthropic()? For zero-code instrumentation of the "
            "Anthropic module, use amplitude_ai.patch_anthropic(amplitude=...) instead."
        )

    raise AmplitudeAIWrapError(
        f"wrap() does not support {full_name}. "
        f"Supported types: {supported}. "
        f"For Gemini, Bedrock, and Mistral, use the SDK's provider classes directly."
        f"{hint}"
    )
