"""
Explicit per-provider monkey-patching for zero-code instrumentation.

Usage::

    import amplitude_ai

    amplitude_ai.patch(amplitude=amp)
    # All installed providers are instrumented automatically

    amplitude_ai.unpatch()
    # Restores all originals -- critical for test isolation

Each ``patch_*()`` function replaces the provider's client class in its
module with the SDK's wrapper class.  The wrapper inherits from the
original, so ``isinstance()`` checks still pass.

``unpatch_*()`` restores the original reference.  Always call it in test
teardown to prevent state leakage between tests.

Both ``patch_*()`` and ``unpatch_*()`` are idempotent.
"""

from __future__ import annotations

import sys
from typing import Any

from amplitude import Amplitude

# ---------------------------------------------------------------------------
# Internal state
# ---------------------------------------------------------------------------

_original_openai_class: Any = None
_original_async_openai_class: Any = None
_original_anthropic_class: Any = None
_original_azure_openai_class: Any = None
_openai_patched = False
_async_openai_patched = False
_anthropic_patched = False
_azure_openai_patched = False

# Stored config for the patched constructors
_patch_config: dict[str, Any] = {}


# ---------------------------------------------------------------------------
# Debug hook for the patching path
# ---------------------------------------------------------------------------

_debug_hook_installed = False


def _install_debug_hook(amplitude: Amplitude) -> None:
    """Install the debug/stderr hook on *amplitude* so patched clients print events.

    Guarded so it only wraps once, even when called from multiple ``patch_*()`` functions.
    """
    global _debug_hook_installed
    if _debug_hook_installed:
        return

    import sys as _sys

    from .utils.debug import format_debug_line

    original_track = amplitude.track

    def _hooked_track(event: Any) -> None:
        print(format_debug_line(event), file=_sys.stderr)
        original_track(event)

    amplitude.track = _hooked_track  # type: ignore[assignment,unused-ignore]
    _debug_hook_installed = True


# ---------------------------------------------------------------------------
# OpenAI
# ---------------------------------------------------------------------------

def patch_openai(
    amplitude: Amplitude,
    user_id: str | None = None,
    *,
    debug: bool = False,
    **kwargs: Any,
) -> None:
    """Replace ``openai.OpenAI`` with the SDK's instrumented wrapper.

    All subsequent ``openai.OpenAI()`` instantiations will return an
    instrumented client that auto-tracks LLM calls.

    Args:
        amplitude: Amplitude client instance.
        user_id: Default user ID for all tracked events.  Can be
            overridden per-call via ``amplitude_user_id`` kwargs.
        debug: If True, print colored event summaries to stderr.
        **kwargs: Additional kwargs forwarded to the SDK's OpenAI wrapper
            (e.g., ``privacy_mode=True``).
    """
    global _original_openai_class, _openai_patched

    if _openai_patched:
        return

    import openai as openai_mod

    from .providers.openai import OpenAI as AmplitudeOpenAI

    _original_openai_class = openai_mod.OpenAI

    _patch_config["openai"] = {
        "amplitude": amplitude,
        "user_id": user_id,
        **kwargs,
    }

    _amplitude_reserved = frozenset({"amplitude", "api_key", "privacy_mode", "content_mode", "user_id", "debug"})

    class PatchedOpenAI(_original_openai_class):  # type: ignore[misc]
        """Patched OpenAI client that auto-instruments via Amplitude AI SDK."""

        def __init__(self, **init_kwargs: Any) -> None:
            cfg = _patch_config.get("openai", {})
            amp = cfg.get("amplitude")
            if amp is None:
                super().__init__(**init_kwargs)
                return

            merged = {**cfg, **init_kwargs}

            passthrough = {k: v for k, v in merged.items() if k not in _amplitude_reserved}
            wrapper = AmplitudeOpenAI(
                amplitude=amp,
                api_key=merged.get("api_key"),
                privacy_mode=merged.get("privacy_mode", False),
                content_mode=merged.get("content_mode"),
                **passthrough,
            )
            self.__dict__.update(wrapper.__dict__)
            self.__class__ = wrapper.__class__

            default_uid = init_kwargs.get("user_id") or cfg.get("user_id")
            if default_uid:
                self._default_amplitude_user_id = default_uid

    openai_mod.OpenAI = PatchedOpenAI
    _openai_patched = True

    if debug:
        _install_debug_hook(amplitude)


def unpatch_openai() -> None:
    """Restore the original ``openai.OpenAI`` class.  Idempotent."""
    global _original_openai_class, _openai_patched

    if not _openai_patched:
        return

    import openai as openai_mod

    openai_mod.OpenAI = _original_openai_class
    _original_openai_class = None
    _openai_patched = False
    _patch_config.pop("openai", None)


# ---------------------------------------------------------------------------
# Async OpenAI
# ---------------------------------------------------------------------------

def patch_async_openai(
    amplitude: Amplitude,
    user_id: str | None = None,
    *,
    debug: bool = False,
    **kwargs: Any,
) -> None:
    """Replace ``openai.AsyncOpenAI`` with the SDK's instrumented async wrapper.

    All subsequent ``openai.AsyncOpenAI()`` instantiations will return an
    instrumented client that auto-tracks LLM calls.

    Args:
        amplitude: Amplitude client instance.
        user_id: Default user ID for all tracked events.
        debug: If True, print colored event summaries to stderr.
        **kwargs: Additional kwargs forwarded to the SDK's AsyncOpenAI wrapper.
    """
    global _original_async_openai_class, _async_openai_patched

    if _async_openai_patched:
        return

    import openai as openai_mod

    from .providers.openai import AsyncOpenAI as AmplitudeAsyncOpenAI

    _original_async_openai_class = openai_mod.AsyncOpenAI

    _patch_config["async_openai"] = {
        "amplitude": amplitude,
        "user_id": user_id,
        **kwargs,
    }

    _async_amplitude_reserved = frozenset({"amplitude", "api_key", "privacy_mode", "content_mode", "user_id", "debug"})

    class PatchedAsyncOpenAI(_original_async_openai_class):  # type: ignore[misc]
        """Patched AsyncOpenAI client that auto-instruments via Amplitude AI SDK."""

        def __init__(self, **init_kwargs: Any) -> None:
            cfg = _patch_config.get("async_openai", {})
            amp = cfg.get("amplitude")
            if amp is None:
                super().__init__(**init_kwargs)
                return

            merged = {**cfg, **init_kwargs}

            passthrough = {k: v for k, v in merged.items() if k not in _async_amplitude_reserved}
            wrapper = AmplitudeAsyncOpenAI(
                amplitude=amp,
                api_key=merged.get("api_key"),
                privacy_mode=merged.get("privacy_mode", False),
                content_mode=merged.get("content_mode"),
                **passthrough,
            )
            self.__dict__.update(wrapper.__dict__)
            self.__class__ = wrapper.__class__

            default_uid = init_kwargs.get("user_id") or cfg.get("user_id")
            if default_uid:
                self._default_amplitude_user_id = default_uid

    openai_mod.AsyncOpenAI = PatchedAsyncOpenAI
    _async_openai_patched = True

    if debug:
        _install_debug_hook(amplitude)


def unpatch_async_openai() -> None:
    """Restore the original ``openai.AsyncOpenAI`` class.  Idempotent."""
    global _original_async_openai_class, _async_openai_patched

    if not _async_openai_patched:
        return

    import openai as openai_mod

    openai_mod.AsyncOpenAI = _original_async_openai_class
    _original_async_openai_class = None
    _async_openai_patched = False
    _patch_config.pop("async_openai", None)


# ---------------------------------------------------------------------------
# Anthropic
# ---------------------------------------------------------------------------

def patch_anthropic(
    amplitude: Amplitude,
    user_id: str | None = None,
    *,
    debug: bool = False,
    **kwargs: Any,
) -> None:
    """Replace ``anthropic.Anthropic`` with the SDK's instrumented wrapper.

    All subsequent ``anthropic.Anthropic()`` instantiations will return an
    instrumented client that auto-tracks LLM calls.

    Args:
        amplitude: Amplitude client instance.
        user_id: Default user ID for all tracked events.
        debug: If True, print colored event summaries to stderr.
        **kwargs: Additional kwargs forwarded to the SDK's Anthropic wrapper.
    """
    global _original_anthropic_class, _anthropic_patched

    if _anthropic_patched:
        return

    import anthropic as anthropic_mod

    from .providers.anthropic import Anthropic as AmplitudeAnthropic

    _original_anthropic_class = anthropic_mod.Anthropic

    _patch_config["anthropic"] = {
        "amplitude": amplitude,
        "user_id": user_id,
        **kwargs,
    }

    _anthropic_reserved = frozenset({"amplitude", "api_key", "privacy_mode", "content_mode", "user_id", "debug"})

    class PatchedAnthropic(_original_anthropic_class):  # type: ignore[misc]
        """Patched Anthropic client that auto-instruments via Amplitude AI SDK."""

        def __init__(self, **init_kwargs: Any) -> None:
            cfg = _patch_config.get("anthropic", {})
            amp = cfg.get("amplitude")
            if amp is None:
                super().__init__(**init_kwargs)
                return

            merged = {**cfg, **init_kwargs}

            passthrough = {k: v for k, v in merged.items() if k not in _anthropic_reserved}
            wrapper = AmplitudeAnthropic(
                amplitude=amp,
                api_key=merged.get("api_key"),
                privacy_mode=merged.get("privacy_mode", False),
                content_mode=merged.get("content_mode"),
                **passthrough,
            )
            self.__dict__.update(wrapper.__dict__)
            self.__class__ = wrapper.__class__

            default_uid = init_kwargs.get("user_id") or cfg.get("user_id")
            if default_uid:
                self._default_amplitude_user_id = default_uid

    anthropic_mod.Anthropic = PatchedAnthropic
    _anthropic_patched = True

    if debug:
        _install_debug_hook(amplitude)


def unpatch_anthropic() -> None:
    """Restore the original ``anthropic.Anthropic`` class.  Idempotent."""
    global _original_anthropic_class, _anthropic_patched

    if not _anthropic_patched:
        return

    import anthropic as anthropic_mod

    anthropic_mod.Anthropic = _original_anthropic_class
    _original_anthropic_class = None
    _anthropic_patched = False
    _patch_config.pop("anthropic", None)


# ---------------------------------------------------------------------------
# Azure OpenAI
# ---------------------------------------------------------------------------

def patch_azure_openai(
    amplitude: Amplitude,
    user_id: str | None = None,
    *,
    debug: bool = False,
    **kwargs: Any,
) -> None:
    """Replace ``openai.AzureOpenAI`` with the SDK's instrumented wrapper.

    All subsequent ``openai.AzureOpenAI()`` instantiations will return an
    instrumented client that auto-tracks LLM calls.

    Args:
        amplitude: Amplitude client instance.
        user_id: Default user ID for all tracked events.
        debug: If True, print colored event summaries to stderr.
        **kwargs: Additional kwargs forwarded to the SDK's AzureOpenAI wrapper.
    """
    global _original_azure_openai_class, _azure_openai_patched

    if _azure_openai_patched:
        return

    import openai as openai_mod

    from .providers.azure_openai import AzureOpenAI as AmplitudeAzureOpenAI

    _original_azure_openai_class = openai_mod.AzureOpenAI

    _patch_config["azure_openai"] = {
        "amplitude": amplitude,
        "user_id": user_id,
        **kwargs,
    }

    _azure_reserved = frozenset({
        "amplitude", "api_key", "azure_endpoint", "api_version",
        "privacy_mode", "content_mode", "user_id", "debug",
    })

    class PatchedAzureOpenAI(_original_azure_openai_class):  # type: ignore[misc,unused-ignore]
        """Patched AzureOpenAI client that auto-instruments via Amplitude AI SDK."""

        def __init__(self, **init_kwargs: Any) -> None:
            cfg = _patch_config.get("azure_openai", {})
            amp = cfg.get("amplitude")
            if amp is None:
                super().__init__(**init_kwargs)
                return

            merged = {**cfg, **init_kwargs}

            passthrough = {k: v for k, v in merged.items() if k not in _azure_reserved}
            wrapper = AmplitudeAzureOpenAI(
                amplitude=amp,
                api_key=merged.get("api_key"),
                azure_endpoint=merged.get("azure_endpoint"),
                api_version=merged.get("api_version"),
                privacy_mode=merged.get("privacy_mode", False),
                content_mode=merged.get("content_mode"),
                **passthrough,
            )
            self.__dict__.update(wrapper.__dict__)
            self.__class__ = wrapper.__class__

            default_uid = init_kwargs.get("user_id") or cfg.get("user_id")
            if default_uid:
                self._default_amplitude_user_id = default_uid

    openai_mod.AzureOpenAI = PatchedAzureOpenAI
    _azure_openai_patched = True

    if debug:
        _install_debug_hook(amplitude)


def unpatch_azure_openai() -> None:
    """Restore the original ``openai.AzureOpenAI`` class.  Idempotent."""
    global _original_azure_openai_class, _azure_openai_patched

    if not _azure_openai_patched:
        return

    import openai as openai_mod

    openai_mod.AzureOpenAI = _original_azure_openai_class
    _original_azure_openai_class = None
    _azure_openai_patched = False
    _patch_config.pop("azure_openai", None)


# ---------------------------------------------------------------------------
# Gemini (google-genai)
# ---------------------------------------------------------------------------

_original_genai_client_class: Any = None
_gemini_patched = False


def patch_gemini(
    amplitude: Amplitude,
    user_id: str | None = None,
    *,
    debug: bool = False,
    **kwargs: Any,
) -> None:
    """Replace ``google.genai.Client`` with an instrumented wrapper.

    All subsequent ``genai.Client()`` instantiations will return a client
    whose ``models.generate_content`` calls are auto-tracked.

    Args:
        amplitude: Amplitude client instance.
        user_id: Default user ID for all tracked events.
        debug: If True, print colored event summaries to stderr.
    """
    global _original_genai_client_class, _gemini_patched

    if _gemini_patched:
        return

    from google import genai as genai_mod

    _original_genai_client_class = genai_mod.Client

    _patch_config["gemini"] = {
        "amplitude": amplitude,
        "user_id": user_id,
        **kwargs,
    }

    class PatchedGenAIClient(_original_genai_client_class):  # type: ignore[misc]
        """Patched google.genai.Client that auto-instruments generate_content."""

        def __init__(self, **init_kwargs: Any) -> None:
            super().__init__(**init_kwargs)
            cfg = _patch_config.get("gemini", {})
            amp = cfg.get("amplitude")
            if amp is not None:
                self._amplitude = amp
                self._amplitude_user_id = cfg.get("user_id")
                self._amplitude_privacy_mode = cfg.get("privacy_mode", False)
                self._amplitude_content_mode = cfg.get("content_mode")
                # Wrap models.generate_content
                original_generate = self.models.generate_content

                def _tracked_generate(*, model: str, contents: Any, **gkw: Any) -> Any:
                    return _gemini_tracked_generate(
                        self, original_generate, model=model, contents=contents, **gkw,
                    )
                self.models.generate_content = _tracked_generate

    genai_mod.Client = PatchedGenAIClient
    _gemini_patched = True

    if debug:
        _install_debug_hook(amplitude)


def _gemini_tracked_generate(
    client: Any,
    original_fn: Any,
    *,
    model: str,
    contents: Any,
    **gkwargs: Any,
) -> Any:
    """Shared tracking logic for patched Gemini generate_content calls."""
    import time as _time
    import uuid as _uuid

    from .core.privacy import PrivacyConfig
    from .core.tracking import track_ai_message, track_user_message
    from .providers.base import _apply_session_context
    from .utils.costs import calculate_cost

    amp = client._amplitude
    amplitude_user_id = gkwargs.pop("amplitude_user_id", None) or client._amplitude_user_id
    amplitude_session_id = gkwargs.pop("amplitude_session_id", None)
    amplitude_trace_id = gkwargs.pop("amplitude_trace_id", None)
    amplitude_turn_id = gkwargs.pop("amplitude_turn_id", None)
    amplitude_event_properties = gkwargs.pop("amplitude_event_properties", None)
    amplitude_groups = gkwargs.pop("amplitude_groups", None)

    (
        amplitude_user_id, amplitude_session_id,
        amplitude_trace_id, amplitude_turn_id,
        amplitude_event_properties, amplitude_groups,
    ) = _apply_session_context(
        user_id=amplitude_user_id,
        session_id=amplitude_session_id,
        trace_id=amplitude_trace_id,
        turn_id=amplitude_turn_id,
        event_properties=amplitude_event_properties,
        groups=amplitude_groups,
    )

    privacy_config = PrivacyConfig(
        privacy_mode=client._amplitude_privacy_mode,
        content_mode=getattr(client, "_amplitude_content_mode", None),
    )
    trace_id = amplitude_trace_id or str(_uuid.uuid4())
    content_str = str(contents) if not isinstance(contents, str) else contents

    _gemini_parent_agent_id = (
        amplitude_event_properties.get("parent_agent_id")
        if amplitude_event_properties
        else None
    )
    if amplitude_user_id:
        track_user_message(
            amplitude=amp,
            user_id=amplitude_user_id,
            message_content=content_str,
            session_id=amplitude_session_id,
            trace_id=trace_id,
            turn_id=amplitude_turn_id,
            message_source="agent" if _gemini_parent_agent_id else "user",
            privacy_config=privacy_config,
            groups=amplitude_groups,
        )

    start = _time.time()
    is_error = False
    error_msg = None
    response = None
    try:
        response = original_fn(model=model, contents=contents, **gkwargs)
    except Exception as exc:
        is_error = True
        error_msg = str(exc)
        raise
    finally:
        latency_ms = (_time.time() - start) * 1000
        resp_text = ""
        input_tokens = output_tokens = total_tokens = 0
        if response is not None and not is_error:
            try:
                resp_text = response.text if hasattr(response, "text") else ""
            except (ValueError, AttributeError):
                resp_text = ""
            usage = getattr(response, "usage_metadata", None)
            if usage:
                input_tokens = getattr(usage, "prompt_token_count", 0) or 0
                output_tokens = getattr(usage, "candidates_token_count", 0) or 0
                total_tokens = getattr(usage, "total_token_count", 0) or 0

        if amplitude_user_id:
            cost = calculate_cost(
                model_name=model,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                default_provider="google",
            )
            raw_tools = gkwargs.get("tools")
            _tool_defs = raw_tools if isinstance(raw_tools, list) and raw_tools else None
            track_ai_message(
                amplitude=amp,
                user_id=amplitude_user_id,
                response_content=resp_text,
                model_name=model,
                provider="gemini",
                latency_ms=latency_ms,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                total_tokens=total_tokens,
                total_cost_usd=cost,
                is_error=is_error,
                error_message=error_msg,
                session_id=amplitude_session_id,
                trace_id=trace_id,
                turn_id=amplitude_turn_id,
                tool_definitions=_tool_defs,
                privacy_config=privacy_config,
                groups=amplitude_groups,
            )

    return response


def unpatch_gemini() -> None:
    """Restore the original ``google.genai.Client`` class.  Idempotent."""
    global _original_genai_client_class, _gemini_patched

    if not _gemini_patched:
        return

    from google import genai as genai_mod

    genai_mod.Client = _original_genai_client_class
    _original_genai_client_class = None
    _gemini_patched = False
    _patch_config.pop("gemini", None)


# ---------------------------------------------------------------------------
# Mistral
# ---------------------------------------------------------------------------

_original_mistral_class: Any = None
_mistral_patched = False


def patch_mistral(
    amplitude: Amplitude,
    user_id: str | None = None,
    *,
    debug: bool = False,
    **kwargs: Any,
) -> None:
    """Replace ``mistralai.Mistral`` with an instrumented wrapper.

    All subsequent ``mistralai.Mistral()`` instantiations will return a client
    that auto-tracks chat completions.

    Args:
        amplitude: Amplitude client instance.
        user_id: Default user ID for all tracked events.
        debug: If True, print colored event summaries to stderr.
    """
    global _original_mistral_class, _mistral_patched

    if _mistral_patched:
        return

    import mistralai as mistral_mod

    _original_mistral_class = mistral_mod.Mistral

    _patch_config["mistral"] = {
        "amplitude": amplitude,
        "user_id": user_id,
        **kwargs,
    }

    from .providers.mistral import Mistral as AmplitudeMistral, _WrappedChat

    class PatchedMistral(_original_mistral_class):  # type: ignore[misc]
        """Patched Mistral client that auto-instruments via Amplitude AI SDK."""

        def __init__(self, **init_kwargs: Any) -> None:
            super().__init__(**init_kwargs)
            cfg = _patch_config.get("mistral", {})
            amp = cfg.get("amplitude")
            if amp is None:
                return

            from .core.privacy import PrivacyConfig

            # Save the native chat namespace before replacing it.
            native_chat = self.chat  # type: ignore[has-type]

            # Graft the attributes that _WrappedChat expects on its owner.
            self.amplitude = amp
            self.privacy_config = PrivacyConfig(
                privacy_mode=cfg.get("privacy_mode", False),
                content_mode=cfg.get("content_mode"),
            )

            # _WrappedChat delegates native calls to self._client.chat.
            # Point _client at a proxy that routes to the saved native chat.
            class _NativeProxy:
                def __init__(self, chat: Any) -> None:
                    self.chat = chat

            self._client = _NativeProxy(native_chat)

            # Bind tracking methods from the provider wrapper class.
            self._track_completion = AmplitudeMistral._track_completion.__get__(self)
            self._extract_response_content = AmplitudeMistral._extract_response_content
            self._extract_reasoning_content = AmplitudeMistral._extract_reasoning_content
            self._extract_finish_reason = AmplitudeMistral._extract_finish_reason
            self._extract_tool_calls = AmplitudeMistral._extract_tool_calls

            self.chat = _WrappedChat(self)

            default_uid = init_kwargs.get("user_id") or cfg.get("user_id")
            if default_uid:
                self._default_amplitude_user_id = default_uid

    mistral_mod.Mistral = PatchedMistral
    _mistral_patched = True

    if debug:
        _install_debug_hook(amplitude)


def unpatch_mistral() -> None:
    """Restore the original ``mistralai.Mistral`` class.  Idempotent."""
    global _original_mistral_class, _mistral_patched

    if not _mistral_patched:
        return

    import mistralai as mistral_mod

    mistral_mod.Mistral = _original_mistral_class
    _original_mistral_class = None
    _mistral_patched = False
    _patch_config.pop("mistral", None)


# ---------------------------------------------------------------------------
# One-liner: auto-detect and patch all installed providers
# ---------------------------------------------------------------------------

def patch(
    amplitude: Amplitude,
    *,
    debug: bool = False,
    **kwargs: Any,
) -> list[str]:
    """Auto-detect installed LLM providers and patch them all.

    This is the recommended way to instrument your application -- a single
    call patches every provider whose SDK is importable::

        amplitude_ai.patch(amplitude=Amplitude("YOUR_API_KEY"))

    Returns a list of provider names that were successfully patched
    (e.g. ``["openai", "anthropic", "gemini"]``).

    Args:
        amplitude: Amplitude client instance.
        debug: If True, print colored event summaries to stderr.
        **kwargs: Forwarded to each ``patch_*()`` call (e.g. ``privacy_mode=True``).
    """
    patched: list[str] = []

    for name, patch_fn in [
        ("openai", patch_openai),
        ("async_openai", patch_async_openai),
        ("anthropic", patch_anthropic),
        ("azure_openai", patch_azure_openai),
        ("gemini", patch_gemini),
        ("mistral", patch_mistral),
    ]:
        try:
            patch_fn(amplitude=amplitude, debug=debug, **kwargs)
            patched.append(name)
        except ImportError:
            pass
        except Exception as exc:
            print(
                f"amplitude-ai: failed to patch {name}: {exc}",
                file=sys.stderr,
            )

    return patched


def unpatch() -> None:
    """Unpatch all currently patched providers.  Idempotent."""
    global _debug_hook_installed

    unpatch_openai()
    unpatch_async_openai()
    unpatch_anthropic()
    unpatch_azure_openai()
    unpatch_gemini()
    unpatch_mistral()
    _debug_hook_installed = False


# ---------------------------------------------------------------------------
# Query state
# ---------------------------------------------------------------------------

def patched_providers() -> list[str]:
    """Return a list of currently patched provider names."""
    result: list[str] = []
    if _openai_patched:
        result.append("openai")
    if _async_openai_patched:
        result.append("async_openai")
    if _anthropic_patched:
        result.append("anthropic")
    if _azure_openai_patched:
        result.append("azure-openai")
    if _gemini_patched:
        result.append("gemini")
    if _mistral_patched:
        result.append("mistral")
    return result
