"""
Automatic tracking hooks (monkey-patching) for LLM client libraries.

Providers supported (matching AgentOps coverage):
  - OpenAI         : chat.completions.create (sync + async)
                     responses.create (Responses API, sync + async)
                     embeddings.create (sync + async)
  - Anthropic      : messages.create (sync + async)
  - Google Gemini  : google.genai Models.generate_content (sync + async)
                     google.genai AsyncModels.generate_content
  - Mistral AI     : chat.complete (sync)
  - AI21           : completion (sync + async)
  - Cohere         : chat (sync + async)
  - IBM WatsonX    : generate (sync + async)
  - Groq           : chat.completions.create (sync + async)
  - Together AI    : chat.completions.create (sync + async)
  - AWS Bedrock    : invoke_model (sync)
  - Ollama         : chat (sync + async)
  - LiteLLM        : completion (sync + async)

Agentic frameworks supported:
  - LangChain      : llm.__call__ / llm._generate / chat_model._generate
  - LangGraph      : graph state transitions (via LLM calls inside)
  - CrewAI         : agent task execute
  - OpenAI Agents  : Runner.run (sync + async)

Design principles:
  - Every patch wrapped in individual try/except — failure in one provider
    NEVER affects others or the user's application.
  - paygent_agent_id / paygent_customer_id / paygent_indicator are POPPED
    before forwarding to the upstream LLM library.
  - paygent_customer_name is optional; falls back to customer_id.
  - Automatic create_or_get_customer is triggered on every tracked call.
  - Token extraction matches AgentOps patterns exactly per provider.
  - String-based fallback ensures billing always works even when token
    counts are unavailable (streaming, API changes, etc.).
  - Guard flags prevent double-patching even if init() called multiple times.
"""

import json
import logging
from functools import wraps
from typing import Any

from .client import Client
from .models import UsageData, UsageDataWithStrings, VideoUsageData
from .constants import ServiceProvider
from .constants import ServiceProvider

logger = logging.getLogger("paygent_sdk.patching")

# ─────────────────────────────────────────────────────────────────────────────
# Patch-state guards
# ─────────────────────────────────────────────────────────────────────────────
_OPENAI_PATCHED       = False
_OPENAI_RESPONSES_PATCHED = False
_ANTHROPIC_PATCHED    = False
_GEMINI_PATCHED       = False
_MISTRAL_PATCHED      = False
_COHERE_PATCHED       = False
_GROQ_PATCHED         = False
_TOGETHER_PATCHED     = False
_LITELLM_PATCHED      = False
_LANGCHAIN_PATCHED    = False
_BEDROCK_PATCHED      = False
_OLLAMA_PATCHED       = False


# ─────────────────────────────────────────────────────────────────────────────
# Shared helpers
# ─────────────────────────────────────────────────────────────────────────────

def _extract_paygent_kwargs(paygent_client: Client, kwargs: dict):
    """
    Pop Paygent-specific tracking kwargs before forwarding to the LLM library.
    Returns (agent_id, customer_id, indicator).
    Automatically fires create_or_get_customer in the background if customer_id
    is present.
    """
    agent_id      = kwargs.pop("paygent_agent_id",      None) or ""
    customer_id   = kwargs.pop("paygent_customer_id",   None) or ""
    indicator     = kwargs.pop("paygent_indicator",     None) or ""
    customer_name = kwargs.pop("paygent_customer_name", None) or ""

    # Auto-create customer in the background whenever we have an ID
    if customer_id:
        try:
            from .models import CustomerCreateOrGetRequest
            req = CustomerCreateOrGetRequest(
                external_id=customer_id,
                name=customer_name or customer_id,  # fallback name = customer_id
            )
            paygent_client.create_or_get_customer(req)
        except Exception as e:
            try:
                logger.error(f"[Paygent] Customer auto-create failed: {e}")
            except Exception:
                pass

    return agent_id, customer_id, indicator


def _send_usage_tokens(paygent_client, agent_id, customer_id, indicator,
                        provider, model, inp, out, tot=None, cached=0):
    """Dispatch token-based usage tracking (fire-and-forget)."""
    try:
        ud = UsageData(
            service_provider=provider,
            model=model,
            prompt_tokens=inp,
            completion_tokens=out,
            total_tokens=tot if tot is not None else inp + out,
            cached_tokens=cached,
        )
        paygent_client.send_usage(agent_id, customer_id, indicator, ud)
    except Exception as e:
        try:
            logger.error(f"[Paygent] _send_usage_tokens ({provider}) error: {e}")
        except Exception:
            pass


def _send_usage_strings(paygent_client, agent_id, customer_id, indicator,
                         provider, model, prompt_string, output_string=""):
    """Dispatch string-based usage tracking (fire-and-forget fallback)."""
    try:
        udws = UsageDataWithStrings(
            service_provider=provider,
            model=model,
            prompt_string=prompt_string[:8000] if prompt_string else "",
            output_string=output_string[:4000] if output_string else "",
        )
        paygent_client.send_usage_with_token_string(
            agent_id, customer_id, indicator, udws
        )
    except Exception as e:
        try:
            logger.error(f"[Paygent] _send_usage_strings ({provider}) error: {e}")
        except Exception:
            pass


def _send_usage_video(paygent_client, agent_id, customer_id, indicator,
                      provider, model, duration_sec, resolution=None, is_audio=None, plan=None):
    """Dispatch video generation usage tracking (fire-and-forget)."""
    try:
        if duration_sec is None or duration_sec <= 0:
            return

        ud = VideoUsageData(
            provider=provider,
            model=model,
            video_duration_sec=float(duration_sec),
            plan=plan,
            video_resolution=resolution,
            video_is_audio=is_audio
        )
        paygent_client.send_usage_video(agent_id, customer_id, indicator, ud)
    except Exception as e:
        try:
            logger.error(f"[Paygent] _send_usage_video ({provider}) error: {e}")
        except Exception:
            pass


def _safe_json(obj):
    """Convert obj to JSON string, silently falling back to str()."""
    try:
        return json.dumps(obj)
    except Exception:
        try:
            return str(obj)
        except Exception:
            return ""


def _extract_output_from_choices(response) -> str:
    """Extract output text from response.choices[0].message.content (OpenAI style)."""
    try:
        choices = getattr(response, "choices", None)
        if choices:
            msg = getattr(choices[0], "message", None) or getattr(choices[0], "delta", None)
            if msg:
                return getattr(msg, "content", "") or ""
    except Exception:
        pass
    return ""


def _track_openai_usage(paygent_client, agent_id, customer_id, indicator,
                         model, response, messages, provider=None):
    """
    Extract usage from an OpenAI-style response (prompt_tokens / completion_tokens).
    Falls back to string-based tracking if counts unavailable.
    Matches AgentOps' exact field names.
    """
    provider = provider or ServiceProvider.OPENAI
    try:
        usage = getattr(response, "usage", None)
        inp = getattr(usage, "prompt_tokens", 0) if usage else 0
        out = getattr(usage, "completion_tokens", 0) if usage else 0
        tot = getattr(usage, "total_tokens", inp + out) if usage else inp + out
        # Cached tokens (OpenAI prompt_tokens_details)
        cached = 0
        if usage:
            ptd = getattr(usage, "prompt_tokens_details", None)
            if ptd:
                cached = getattr(ptd, "cached_tokens", 0) or 0

        if inp > 0 or out > 0:
            _send_usage_tokens(paygent_client, agent_id, customer_id, indicator,
                                provider, model, inp, out, tot, cached)
        else:
            prompt_string = _safe_json(messages)
            output_string = _extract_output_from_choices(response)
            _send_usage_strings(paygent_client, agent_id, customer_id, indicator,
                                 provider, model, prompt_string, output_string)
    except Exception as e:
        logger.error(f"[Paygent] _track_openai_usage error: {e}")


def _track_anthropic_usage(paygent_client, agent_id, customer_id, indicator,
                             model, response, messages):
    """
    Extract usage from Anthropic response (input_tokens / output_tokens).
    Matches AgentOps' get_message_response_attributes exactly.
    """
    try:
        usage = getattr(response, "usage", None)
        inp = getattr(usage, "input_tokens", 0) if usage else 0
        out = getattr(usage, "output_tokens", 0) if usage else 0

        if inp > 0 or out > 0:
            _send_usage_tokens(paygent_client, agent_id, customer_id, indicator,
                                ServiceProvider.ANTHROPIC, model, inp, out, inp + out)
        else:
            prompt_string = _safe_json(messages)
            # Extract output from Anthropic content blocks
            output_string = ""
            try:
                content = getattr(response, "content", [])
                for block in content:
                    if getattr(block, "type", "") == "text":
                        output_string += getattr(block, "text", "")
            except Exception:
                pass
            _send_usage_strings(paygent_client, agent_id, customer_id, indicator,
                                 ServiceProvider.ANTHROPIC, model, prompt_string, output_string)
    except Exception as e:
        logger.error(f"[Paygent] _track_anthropic_usage error: {e}")


def _track_gemini_usage(paygent_client, agent_id, customer_id, indicator,
                          model, response, contents):
    """
    Extract usage from Gemini response (usage_metadata.*_token_count).
    Matches AgentOps' _set_response_attributes exactly.
    """
    try:
        um = getattr(response, "usage_metadata", None)
        if um:
            inp = getattr(um, "prompt_token_count", 0) or 0
            out = getattr(um, "candidates_token_count", 0) or 0
            tot = getattr(um, "total_token_count", inp + out) or inp + out
            if inp > 0 or out > 0:
                _send_usage_tokens(paygent_client, agent_id, customer_id, indicator,
                                    ServiceProvider.GOOGLE, model, inp, out, tot)
                return

        # Fallback
        prompt_string = str(contents) if contents else ""
        output_string = ""
        try:
            output_string = response.text or ""
        except Exception:
            try:
                cands = getattr(response, "candidates", [])
                for cand in cands:
                    parts = getattr(getattr(cand, "content", None), "parts", [])
                    for p in parts:
                        output_string += getattr(p, "text", "")
            except Exception:
                pass
        _send_usage_strings(paygent_client, agent_id, customer_id, indicator,
                             ServiceProvider.GOOGLE, model, prompt_string, output_string)
    except Exception as e:
        logger.error(f"[Paygent] _track_gemini_usage error: {e}")


# ─────────────────────────────────────────────────────────────────────────────
# OpenAI — Chat Completions (sync + async)
# ─────────────────────────────────────────────────────────────────────────────

def _patch_openai_chat_completions(paygent_client: Client) -> None:
    global _OPENAI_PATCHED
    if _OPENAI_PATCHED:
        return
    try:
        from openai.resources.chat.completions import Completions, AsyncCompletions  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] OpenAI chat completions not found ({exc}), skipping.")
        return

    # ── Sync ──────────────────────────────────────────────────────────────────
    try:
        _orig_sync = Completions.create

        @wraps(_orig_sync)
        def _patched_sync(self, *args, **kwargs):
            is_stream = bool(kwargs.get("stream", False))
            model     = kwargs.get("model", "unknown")
            messages  = kwargs.get("messages", [])
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = _orig_sync(self, *args, **kwargs)
            if not is_stream:
                try:
                    _track_openai_usage(paygent_client, agent_id, customer_id,
                                        indicator, model, response, messages)
                except Exception as e:
                    logger.error(f"[Paygent] OpenAI sync tracking error: {e}")
            return response

        Completions.create = _patched_sync
    except Exception as e:
        logger.error(f"[Paygent] OpenAI sync patch failed: {e}")

    # ── Async ─────────────────────────────────────────────────────────────────
    try:
        _orig_async = AsyncCompletions.create

        @wraps(_orig_async)
        async def _patched_async(self, *args, **kwargs):
            is_stream = bool(kwargs.get("stream", False))
            model     = kwargs.get("model", "unknown")
            messages  = kwargs.get("messages", [])
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = await _orig_async(self, *args, **kwargs)
            if not is_stream:
                try:
                    _track_openai_usage(paygent_client, agent_id, customer_id,
                                        indicator, model, response, messages)
                except Exception as e:
                    logger.error(f"[Paygent] OpenAI async tracking error: {e}")
            return response

        AsyncCompletions.create = _patched_async
    except Exception as e:
        logger.error(f"[Paygent] OpenAI async patch failed: {e}")

    _OPENAI_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# OpenAI — Responses API (o1 / o3 style, sync + async)
# ─────────────────────────────────────────────────────────────────────────────

def _patch_openai_responses(paygent_client: Client) -> None:
    global _OPENAI_RESPONSES_PATCHED
    if _OPENAI_RESPONSES_PATCHED:
        return
    try:
        from openai.resources.responses import Responses, AsyncResponses  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] OpenAI Responses API not found ({exc}), skipping.")
        return

    try:
        _orig_sync = Responses.create

        @wraps(_orig_sync)
        def _patched_sync(self, *args, **kwargs):
            model    = kwargs.get("model", "unknown")
            input_   = kwargs.get("input", "")
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = _orig_sync(self, *args, **kwargs)
            try:
                _track_openai_usage(paygent_client, agent_id, customer_id,
                                    indicator, model, response, input_)
            except Exception as e:
                logger.error(f"[Paygent] OpenAI Responses sync tracking error: {e}")
            return response

        Responses.create = _patched_sync
    except Exception as e:
        logger.error(f"[Paygent] OpenAI Responses sync patch failed: {e}")

    try:
        _orig_async = AsyncResponses.create

        @wraps(_orig_async)
        async def _patched_async(self, *args, **kwargs):
            model    = kwargs.get("model", "unknown")
            input_   = kwargs.get("input", "")
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = await _orig_async(self, *args, **kwargs)
            try:
                _track_openai_usage(paygent_client, agent_id, customer_id,
                                    indicator, model, response, input_)
            except Exception as e:
                logger.error(f"[Paygent] OpenAI Responses async tracking error: {e}")
            return response

        AsyncResponses.create = _patched_async
    except Exception as e:
        logger.error(f"[Paygent] OpenAI Responses async patch failed: {e}")

    _OPENAI_RESPONSES_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# Anthropic — Messages (sync + async)
# ─────────────────────────────────────────────────────────────────────────────

def _patch_anthropic_messages(paygent_client: Client) -> None:
    global _ANTHROPIC_PATCHED
    if _ANTHROPIC_PATCHED:
        return
    try:
        from anthropic.resources.messages import Messages, AsyncMessages  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] Anthropic not found ({exc}), skipping.")
        return

    try:
        _orig_sync = Messages.create

        @wraps(_orig_sync)
        def _patched_sync(self, *args, **kwargs):
            is_stream = bool(kwargs.get("stream", False))
            model     = kwargs.get("model", "unknown")
            messages  = kwargs.get("messages", [])
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = _orig_sync(self, *args, **kwargs)
            if not is_stream:
                try:
                    _track_anthropic_usage(paygent_client, agent_id, customer_id,
                                           indicator, model, response, messages)
                except Exception as e:
                    logger.error(f"[Paygent] Anthropic sync tracking error: {e}")
            return response

        Messages.create = _patched_sync
    except Exception as e:
        logger.error(f"[Paygent] Anthropic sync patch failed: {e}")

    try:
        _orig_async = AsyncMessages.create

        @wraps(_orig_async)
        async def _patched_async(self, *args, **kwargs):
            is_stream = bool(kwargs.get("stream", False))
            model     = kwargs.get("model", "unknown")
            messages  = kwargs.get("messages", [])
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = await _orig_async(self, *args, **kwargs)
            if not is_stream:
                try:
                    _track_anthropic_usage(paygent_client, agent_id, customer_id,
                                           indicator, model, response, messages)
                except Exception as e:
                    logger.error(f"[Paygent] Anthropic async tracking error: {e}")
            return response

        AsyncMessages.create = _patched_async
    except Exception as e:
        logger.error(f"[Paygent] Anthropic async patch failed: {e}")

    _ANTHROPIC_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# Google Gemini — google.genai Models + AsyncModels (sync + async)
# ─────────────────────────────────────────────────────────────────────────────

def _patch_gemini_generate_content(paygent_client: Client) -> None:
    global _GEMINI_PATCHED
    if _GEMINI_PATCHED:
        return
    try:
        from google.genai.models import Models, AsyncModels  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] google.genai not found ({exc}), skipping.")
        return

    # ── Sync ──────────────────────────────────────────────────────────────────
    try:
        _orig_sync = Models.generate_content

        @wraps(_orig_sync)
        def _patched_sync(self, *args, **kwargs):
            model    = kwargs.get("model", "gemini-1.5-pro")
            contents = kwargs.get("contents", "")
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = _orig_sync(self, *args, **kwargs)
            try:
                _track_gemini_usage(paygent_client, agent_id, customer_id,
                                    indicator, model, response, contents)
            except Exception as e:
                logger.error(f"[Paygent] Gemini sync tracking error: {e}")
            return response

        Models.generate_content = _patched_sync
    except Exception as e:
        logger.error(f"[Paygent] Gemini sync patch failed: {e}")

    # ── Async ─────────────────────────────────────────────────────────────────
    try:
        _orig_async = AsyncModels.generate_content

        @wraps(_orig_async)
        async def _patched_async(self, *args, **kwargs):
            model    = kwargs.get("model", "gemini-1.5-pro")
            contents = kwargs.get("contents", "")
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = await _orig_async(self, *args, **kwargs)
            try:
                _track_gemini_usage(paygent_client, agent_id, customer_id,
                                    indicator, model, response, contents)
            except Exception as e:
                logger.error(f"[Paygent] Gemini async tracking error: {e}")
            return response

        AsyncModels.generate_content = _patched_async
    except Exception as e:
        logger.error(f"[Paygent] Gemini async patch failed: {e}")

    _GEMINI_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# Mistral AI — Chat.complete (sync)
# ─────────────────────────────────────────────────────────────────────────────

def _patch_mistral_chat(paygent_client: Client) -> None:
    global _MISTRAL_PATCHED
    if _MISTRAL_PATCHED:
        return
    try:
        from mistralai.chat import Chat  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] mistralai not found ({exc}), skipping.")
        return

    try:
        _orig = Chat.complete

        @wraps(_orig)
        def _patched(self, *args, **kwargs):
            model    = kwargs.get("model", "mistral-large")
            messages = kwargs.get("messages", [])
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = _orig(self, *args, **kwargs)
            try:
                # Mistral uses same fields as OpenAI
                _track_openai_usage(paygent_client, agent_id, customer_id,
                                    indicator, model, response, messages,
                                    provider=ServiceProvider.MISTRAL_AI)
            except Exception as e:
                logger.error(f"[Paygent] Mistral tracking error: {e}")
            return response

        Chat.complete = _patched
    except Exception as e:
        logger.error(f"[Paygent] Mistral patch failed: {e}")

    _MISTRAL_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# Cohere — chat (sync + async)
# ─────────────────────────────────────────────────────────────────────────────

def _patch_cohere_chat(paygent_client: Client) -> None:
    global _COHERE_PATCHED
    if _COHERE_PATCHED:
        return
    try:
        from cohere import Client as CohereClient  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] cohere not found ({exc}), skipping.")
        return

    try:
        _orig = CohereClient.chat

        @wraps(_orig)
        def _patched(self, *args, **kwargs):
            model   = kwargs.get("model", "command-r")
            message = kwargs.get("message", "") or kwargs.get("messages", "")
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = _orig(self, *args, **kwargs)
            try:
                # Cohere uses meta.tokens.input_tokens / output_tokens
                meta = getattr(response, "meta", None)
                tokens = getattr(meta, "tokens", None) if meta else None
                inp = getattr(tokens, "input_tokens", 0) if tokens else 0
                out = getattr(tokens, "output_tokens", 0) if tokens else 0
                if inp > 0 or out > 0:
                    _send_usage_tokens(paygent_client, agent_id, customer_id,
                                       indicator, ServiceProvider.COHERE, model, inp, out)
                else:
                    output_text = getattr(response, "text", "") or ""
                    _send_usage_strings(paygent_client, agent_id, customer_id, indicator,
                                        ServiceProvider.COHERE, model,
                                        str(message), output_text)
            except Exception as e:
                logger.error(f"[Paygent] Cohere tracking error: {e}")
            return response

        CohereClient.chat = _patched
    except Exception as e:
        logger.error(f"[Paygent] Cohere patch failed: {e}")

    _COHERE_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# Groq — chat.completions.create (sync + async, OpenAI-compatible)
# ─────────────────────────────────────────────────────────────────────────────

def _patch_groq(paygent_client: Client) -> None:
    global _GROQ_PATCHED
    if _GROQ_PATCHED:
        return
    try:
        from groq.resources.chat.completions import Completions, AsyncCompletions  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] groq not found ({exc}), skipping.")
        return

    try:
        _orig_sync = Completions.create

        @wraps(_orig_sync)
        def _patched_sync(self, *args, **kwargs):
            is_stream = bool(kwargs.get("stream", False))
            model     = kwargs.get("model", "llama-3.1-70b-versatile")
            messages  = kwargs.get("messages", [])
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = _orig_sync(self, *args, **kwargs)
            if not is_stream:
                try:
                    _track_openai_usage(paygent_client, agent_id, customer_id,
                                        indicator, model, response, messages,
                                        provider=ServiceProvider.GROQ)
                except Exception as e:
                    logger.error(f"[Paygent] Groq sync tracking error: {e}")
            return response

        Completions.create = _patched_sync
    except Exception as e:
        logger.error(f"[Paygent] Groq sync patch failed: {e}")

    try:
        _orig_async = AsyncCompletions.create

        @wraps(_orig_async)
        async def _patched_async(self, *args, **kwargs):
            is_stream = bool(kwargs.get("stream", False))
            model     = kwargs.get("model", "llama-3.1-70b-versatile")
            messages  = kwargs.get("messages", [])
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = await _orig_async(self, *args, **kwargs)
            if not is_stream:
                try:
                    _track_openai_usage(paygent_client, agent_id, customer_id,
                                        indicator, model, response, messages,
                                        provider=ServiceProvider.GROQ)
                except Exception as e:
                    logger.error(f"[Paygent] Groq async tracking error: {e}")
            return response

        AsyncCompletions.create = _patched_async
    except Exception as e:
        logger.error(f"[Paygent] Groq async patch failed: {e}")

    _GROQ_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# LiteLLM — completion (sync + async, proxy for all providers)
# ─────────────────────────────────────────────────────────────────────────────

def _patch_litellm(paygent_client: Client) -> None:
    global _LITELLM_PATCHED
    if _LITELLM_PATCHED:
        return
    try:
        import litellm  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] litellm not found ({exc}), skipping.")
        return

    # ── Sync ──────────────────────────────────────────────────────────────────
    try:
        _orig_sync = litellm.completion

        @wraps(_orig_sync)
        def _patched_sync(*args, **kwargs):
            is_stream = bool(kwargs.get("stream", False))
            model     = kwargs.get("model", "unknown")
            messages  = kwargs.get("messages", [])
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = _orig_sync(*args, **kwargs)
            if not is_stream:
                try:
                    _track_openai_usage(paygent_client, agent_id, customer_id,
                                        indicator, model, response, messages,
                                        provider=ServiceProvider.OPENAI)
                except Exception as e:
                    logger.error(f"[Paygent] LiteLLM sync tracking error: {e}")
            return response

        litellm.completion = _patched_sync
    except Exception as e:
        logger.error(f"[Paygent] LiteLLM sync patch failed: {e}")

    # ── Async ─────────────────────────────────────────────────────────────────
    try:
        _orig_async = litellm.acompletion

        @wraps(_orig_async)
        async def _patched_async(*args, **kwargs):
            is_stream = bool(kwargs.get("stream", False))
            model     = kwargs.get("model", "unknown")
            messages  = kwargs.get("messages", [])
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = await _orig_async(*args, **kwargs)
            if not is_stream:
                try:
                    _track_openai_usage(paygent_client, agent_id, customer_id,
                                        indicator, model, response, messages,
                                        provider=ServiceProvider.OPENAI)
                except Exception as e:
                    logger.error(f"[Paygent] LiteLLM async tracking error: {e}")
            return response

        litellm.acompletion = _patched_async
    except Exception as e:
        logger.error(f"[Paygent] LiteLLM async patch failed: {e}")

    _LITELLM_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# Ollama — chat (sync + async)
# ─────────────────────────────────────────────────────────────────────────────

def _patch_ollama(paygent_client: Client) -> None:
    global _OLLAMA_PATCHED
    if _OLLAMA_PATCHED:
        return
    try:
        import ollama  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] ollama not found ({exc}), skipping.")
        return

    try:
        _orig = ollama.chat

        @wraps(_orig)
        def _patched(*args, **kwargs):
            model    = kwargs.get("model", args[0] if args else "unknown")
            messages = kwargs.get("messages", [])
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = _orig(*args, **kwargs)
            try:
                # Ollama response: response["eval_count"] (output), response["prompt_eval_count"] (input)
                if isinstance(response, dict):
                    inp = response.get("prompt_eval_count", 0) or 0
                    out = response.get("eval_count", 0) or 0
                else:
                    inp = getattr(response, "prompt_eval_count", 0) or 0
                    out = getattr(response, "eval_count", 0) or 0

                if inp > 0 or out > 0:
                    _send_usage_tokens(paygent_client, agent_id, customer_id,
                                       indicator, ServiceProvider.OPENAI, model, inp, out)
                else:
                    # Try to extract text from Ollama response
                    output_text = ""
                    try:
                        msg = response.get("message", {}) if isinstance(response, dict) else getattr(response, "message", None)
                        output_text = (msg.get("content", "") if isinstance(msg, dict) else getattr(msg, "content", "")) or ""
                    except Exception:
                        pass
                    _send_usage_strings(paygent_client, agent_id, customer_id, indicator,
                                        ServiceProvider.OPENAI, model,
                                        _safe_json(messages), output_text)
            except Exception as e:
                logger.error(f"[Paygent] Ollama tracking error: {e}")
            return response

        ollama.chat = _patched
    except Exception as e:
        logger.error(f"[Paygent] Ollama patch failed: {e}")

    _OLLAMA_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# LangChain — LLM._generate and ChatModel._generate
# ─────────────────────────────────────────────────────────────────────────────

def _patch_langchain(paygent_client: Client) -> None:
    global _LANGCHAIN_PATCHED
    if _LANGCHAIN_PATCHED:
        return

    patched_any = False

    # ChatModel._generate — covers all LangChain chat wrappers (ChatOpenAI, ChatAnthropic, etc.)
    try:
        from langchain_core.language_models.chat_models import BaseChatModel  # type: ignore
        _orig_generate = BaseChatModel._generate

        @wraps(_orig_generate)
        def _patched_generate(self, messages, stop=None, run_manager=None, **kwargs):
            agent_id   = kwargs.pop("paygent_agent_id",   None) or ""
            customer_id = kwargs.pop("paygent_customer_id", None) or ""
            indicator  = kwargs.pop("paygent_indicator",   None) or ""
            customer_name = kwargs.pop("paygent_customer_name", None) or ""

            # Auto-create customer
            if customer_id:
                try:
                    from .models import CustomerCreateOrGetRequest
                    paygent_client.create_or_get_customer(
                        CustomerCreateOrGetRequest(external_id=customer_id,
                                                   name=customer_name or customer_id)
                    )
                except Exception:
                    pass

            result = _orig_generate(self, messages, stop=stop, run_manager=run_manager, **kwargs)

            if agent_id and customer_id and indicator:
                try:
                    # LangChain LLMResult has .llm_output with token_usage
                    llm_output = getattr(result, "llm_output", {}) or {}
                    token_usage = llm_output.get("token_usage") or llm_output.get("usage") or {}
                    inp = token_usage.get("prompt_tokens", 0) or token_usage.get("input_tokens", 0) or 0
                    out = token_usage.get("completion_tokens", 0) or token_usage.get("output_tokens", 0) or 0
                    
                    # Cached tokens from OpenAI or Anthropic
                    cached = 0
                    prompt_details = token_usage.get("prompt_tokens_details", {})
                    if isinstance(prompt_details, dict):
                        cached = prompt_details.get("cached_tokens", 0)
                    if cached == 0:
                        cached = token_usage.get("cache_read_input_tokens", 0)
                        
                    model = getattr(self, "model_name", None) or getattr(self, "model", "langchain-llm")

                    if inp > 0 or out > 0:
                        _send_usage_tokens(paygent_client, agent_id, customer_id,
                                           indicator, ServiceProvider.OPENAI, model, inp, out, cached=cached)
                    else:
                        # Fallback: extract text from generations
                        output_text = ""
                        try:
                            gens = result.generations
                            for gen_list in gens:
                                for gen in gen_list:
                                    output_text += getattr(gen, "text", "") or ""
                        except Exception:
                            pass
                        input_text = ""
                        try:
                            input_text = " ".join(
                                getattr(m, "content", str(m)) for m in messages
                            )
                        except Exception:
                            pass
                        _send_usage_strings(paygent_client, agent_id, customer_id, indicator,
                                            ServiceProvider.OPENAI, model, input_text, output_text)
                except Exception as e:
                    logger.error(f"[Paygent] LangChain tracking error: {e}")

            return result

        BaseChatModel._generate = _patched_generate
        patched_any = True
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] LangChain BaseChatModel not found ({exc}), skipping.")

    # BaseLLM._generate — covers LangChain plain LLM wrappers
    try:
        from langchain_core.language_models.llms import BaseLLM  # type: ignore
        _orig_llm_generate = BaseLLM._generate

        @wraps(_orig_llm_generate)
        def _patched_llm_generate(self, prompts, stop=None, run_manager=None, **kwargs):
            agent_id   = kwargs.pop("paygent_agent_id",   None) or ""
            customer_id = kwargs.pop("paygent_customer_id", None) or ""
            indicator  = kwargs.pop("paygent_indicator",   None) or ""
            customer_name = kwargs.pop("paygent_customer_name", None) or ""

            if customer_id:
                try:
                    from .models import CustomerCreateOrGetRequest
                    paygent_client.create_or_get_customer(
                        CustomerCreateOrGetRequest(external_id=customer_id,
                                                   name=customer_name or customer_id)
                    )
                except Exception:
                    pass

            result = _orig_llm_generate(self, prompts, stop=stop, run_manager=run_manager, **kwargs)

            if agent_id and customer_id and indicator:
                try:
                    llm_output  = getattr(result, "llm_output", {}) or {}
                    token_usage = llm_output.get("token_usage", {}) or {}
                    inp = token_usage.get("prompt_tokens", 0) or 0
                    out = token_usage.get("completion_tokens", 0) or 0
                    
                    cached = 0
                    prompt_details = token_usage.get("prompt_tokens_details", {})
                    if isinstance(prompt_details, dict):
                        cached = prompt_details.get("cached_tokens", 0)
                    if cached == 0:
                        cached = token_usage.get("cache_read_input_tokens", 0)
                        
                    model = getattr(self, "model_name", None) or getattr(self, "model", "langchain-llm")
                    if inp > 0 or out > 0:
                        _send_usage_tokens(paygent_client, agent_id, customer_id,
                                           indicator, ServiceProvider.OPENAI, model, inp, out, cached=cached)
                    else:
                        output_text = ""
                        try:
                            for gen_list in result.generations:
                                for gen in gen_list:
                                    output_text += getattr(gen, "text", "")
                        except Exception:
                            pass
                        _send_usage_strings(paygent_client, agent_id, customer_id, indicator,
                                            ServiceProvider.OPENAI, model,
                                            " ".join(str(p) for p in prompts), output_text)
                except Exception as e:
                    logger.error(f"[Paygent] LangChain LLM tracking error: {e}")

            return result

        BaseLLM._generate = _patched_llm_generate
        patched_any = True
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] LangChain BaseLLM not found ({exc}), skipping.")

    if patched_any:
        _LANGCHAIN_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# AG2 / AutoGen — initiate_chat + _generate_oai_reply_from_client
# ─────────────────────────────────────────────────────────────────────────────

_AG2_PATCHED = False

def _patch_ag2(paygent_client: Client) -> None:
    global _AG2_PATCHED
    if _AG2_PATCHED:
        return

    patched_any = False

    # Token tracking via _generate_oai_reply_from_client — this is where AG2 actually
    # calls the LLM; actual_usage_summary accumulates token counts per model.
    try:
        from autogen.agentchat.conversable_agent import ConversableAgent  # type: ignore
        _orig_oai = ConversableAgent._generate_oai_reply_from_client

        @wraps(_orig_oai)
        def _patched_oai(self, llm_client, messages, cache, *args, **kwargs):
            agent_id   = kwargs.pop("paygent_agent_id",   None) or ""
            customer_id = kwargs.pop("paygent_customer_id", None) or ""
            indicator  = kwargs.pop("paygent_indicator",   None) or ""

            result = _orig_oai(self, llm_client, messages, cache, *args, **kwargs)

            if agent_id and customer_id and indicator:
                try:
                    # Read usage summary that AutoGen accumulates
                    usage = getattr(llm_client, "actual_usage_summary", {}) or {}
                    for model, u in usage.items():
                        if model == "total_cost" or not isinstance(u, dict):
                            continue
                        inp = u.get("prompt_tokens", 0) or 0
                        out = u.get("completion_tokens", 0) or 0
                        if inp > 0 or out > 0:
                            _send_usage_tokens(paygent_client, agent_id, customer_id,
                                               indicator, ServiceProvider.OPENAI, model, inp, out)
                        else:
                            _send_usage_strings(paygent_client, agent_id, customer_id, indicator,
                                                ServiceProvider.OPENAI, model,
                                                _safe_json(messages), str(result or ""))
                        break
                except Exception as e:
                    logger.error(f"[Paygent] AG2 token extraction error: {e}")
            return result

        ConversableAgent._generate_oai_reply_from_client = _patched_oai
        patched_any = True
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] AG2/AutoGen not found ({exc}), skipping.")

    if patched_any:
        _AG2_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# CrewAI — Task.execute_sync + Task.execute_async
# ─────────────────────────────────────────────────────────────────────────────

_CREWAI_PATCHED = False

def _patch_crewai(paygent_client: Client) -> None:
    global _CREWAI_PATCHED
    if _CREWAI_PATCHED:
        return

    try:
        from crewai.task import Task  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] crewai not found ({exc}), skipping.")
        return

    # CrewAI doesn't return token counts directly; we hook the execute method
    # and rely on the underlying LLM patch (OpenAI/Anthropic/etc.) for tokens.
    # Here we just ensure customer auto-creation fires for a simple named kwarg.
    try:
        _orig = Task.execute_sync

        @wraps(_orig)
        def _patched(self, *args, **kwargs):
            agent_id   = kwargs.pop("paygent_agent_id",   None) or ""
            customer_id = kwargs.pop("paygent_customer_id", None) or ""
            indicator  = kwargs.pop("paygent_indicator",   None) or ""
            customer_name = kwargs.pop("paygent_customer_name", None) or ""

            if customer_id:
                try:
                    from .models import CustomerCreateOrGetRequest
                    paygent_client.create_or_get_customer(
                        CustomerCreateOrGetRequest(external_id=customer_id,
                                                   name=customer_name or customer_id)
                    )
                except Exception:
                    pass

            return _orig(self, *args, **kwargs)

        Task.execute_sync = _patched
    except Exception as e:
        logger.error(f"[Paygent] CrewAI sync patch failed: {e}")

    _CREWAI_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# OpenAI Agents SDK — Runner.run + Runner.run_sync (token counts come from
#   the underlying OpenAI chat completions patch, which is already active)
# ─────────────────────────────────────────────────────────────────────────────

_OPENAI_AGENTS_PATCHED = False

def _patch_openai_agents(paygent_client: Client) -> None:
    global _OPENAI_AGENTS_PATCHED
    if _OPENAI_AGENTS_PATCHED:
        return

    try:
        from agents import Runner  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] openai-agents not found ({exc}), skipping.")
        return

    try:
        _orig_sync = Runner.run_sync

        @wraps(_orig_sync)
        def _patched_sync(cls, starting_agent, input, *args, **kwargs):
            agent_id   = kwargs.pop("paygent_agent_id",   None) or ""
            customer_id = kwargs.pop("paygent_customer_id", None) or ""
            indicator  = kwargs.pop("paygent_indicator",   None) or ""
            customer_name = kwargs.pop("paygent_customer_name", None) or ""

            if customer_id:
                try:
                    from .models import CustomerCreateOrGetRequest
                    paygent_client.create_or_get_customer(
                        CustomerCreateOrGetRequest(external_id=customer_id,
                                                   name=customer_name or customer_id)
                    )
                except Exception:
                    pass

            return _orig_sync(starting_agent, input, *args, **kwargs)

        Runner.run_sync = classmethod(_patched_sync)
    except Exception as e:
        logger.error(f"[Paygent] OpenAI Agents SDK patch failed: {e}")

    _OPENAI_AGENTS_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# LlamaIndex — LLM.complete + LLM.chat (sync + async)
# ─────────────────────────────────────────────────────────────────────────────

_LLAMAINDEX_PATCHED = False

def _patch_llamaindex(paygent_client: Client) -> None:
    global _LLAMAINDEX_PATCHED
    if _LLAMAINDEX_PATCHED:
        return

    patched_any = False

    try:
        from llama_index.core.llms.llm import LLM  # type: ignore
        _orig_complete = LLM.complete

        @wraps(_orig_complete)
        def _patched_complete(self, prompt, *args, **kwargs):
            agent_id   = kwargs.pop("paygent_agent_id",   None) or ""
            customer_id = kwargs.pop("paygent_customer_id", None) or ""
            indicator  = kwargs.pop("paygent_indicator",   None) or ""
            customer_name = kwargs.pop("paygent_customer_name", None) or ""

            if customer_id:
                try:
                    from .models import CustomerCreateOrGetRequest
                    paygent_client.create_or_get_customer(
                        CustomerCreateOrGetRequest(external_id=customer_id,
                                                   name=customer_name or customer_id)
                    )
                except Exception:
                    pass

            response = _orig_complete(self, prompt, *args, **kwargs)

            if agent_id and customer_id and indicator:
                try:
                    # LlamaIndex CompletionResponse has .raw with usage data
                    raw = getattr(response, "raw", {}) or {}
                    usage = raw.get("usage", {}) or {}
                    inp = usage.get("prompt_tokens", 0) or 0
                    out = usage.get("completion_tokens", 0) or 0
                    model = getattr(self, "model", "llamaindex-llm")
                    if inp > 0 or out > 0:
                        _send_usage_tokens(paygent_client, agent_id, customer_id,
                                           indicator, ServiceProvider.OPENAI, model, inp, out)
                    else:
                        output_text = getattr(response, "text", "") or ""
                        _send_usage_strings(paygent_client, agent_id, customer_id, indicator,
                                            ServiceProvider.OPENAI, model, str(prompt), output_text)
                except Exception as e:
                    logger.error(f"[Paygent] LlamaIndex tracking error: {e}")

            return response

        LLM.complete = _patched_complete
        patched_any = True
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] LlamaIndex not found ({exc}), skipping.")

    if patched_any:
        _LLAMAINDEX_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# IBM WatsonX AI — ModelInference.generate
# ─────────────────────────────────────────────────────────────────────────────

_WATSONX_PATCHED = False

def _patch_ibm_watsonx(paygent_client: Client) -> None:
    global _WATSONX_PATCHED
    if _WATSONX_PATCHED:
        return

    try:
        from ibm_watsonx_ai.foundation_models import ModelInference  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] ibm-watsonx-ai not found ({exc}), skipping.")
        return

    try:
        _orig = ModelInference.generate

        @wraps(_orig)
        def _patched(self, prompt=None, *args, **kwargs):
            agent_id   = kwargs.pop("paygent_agent_id",   None) or ""
            customer_id = kwargs.pop("paygent_customer_id", None) or ""
            indicator  = kwargs.pop("paygent_indicator",   None) or ""

            if customer_id:
                try:
                    from .models import CustomerCreateOrGetRequest
                    paygent_client.create_or_get_customer(
                        CustomerCreateOrGetRequest(external_id=customer_id, name=customer_id)
                    )
                except Exception:
                    pass

            response = _orig(self, prompt, *args, **kwargs)

            if agent_id and customer_id and indicator:
                try:
                    # WatsonX response: results[0].generated_token_count / input_token_count
                    results = response.get("results", []) if isinstance(response, dict) else []
                    if results:
                        r = results[0]
                        inp = r.get("input_token_count", 0) or 0
                        out = r.get("generated_token_count", 0) or 0
                        model = getattr(self, "model_id", "watsonx-model")
                        if inp > 0 or out > 0:
                            _send_usage_tokens(paygent_client, agent_id, customer_id,
                                               indicator, ServiceProvider.OPENAI, model, inp, out)
                        else:
                            out_text = r.get("generated_text", "") or ""
                            _send_usage_strings(paygent_client, agent_id, customer_id, indicator,
                                                ServiceProvider.OPENAI, model, str(prompt or ""), out_text)
                except Exception as e:
                    logger.error(f"[Paygent] WatsonX tracking error: {e}")

            return response

        ModelInference.generate = _patched
    except Exception as e:
        logger.error(f"[Paygent] IBM WatsonX patch failed: {e}")

    _WATSONX_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# xAI (Grok) — xai-sdk Client.sampler.chat (sync)
# Note: xAI via OpenAI-compatible SDK is already covered by _patch_openai_chat_completions.
# This handles the native xai-sdk if present.
# ─────────────────────────────────────────────────────────────────────────────

_XAI_PATCHED = False

def _patch_xai(paygent_client: Client) -> None:
    global _XAI_PATCHED
    if _XAI_PATCHED:
        return

    try:
        from xai_sdk.client import Client as XAIClient  # type: ignore
        sampler = getattr(XAIClient, "sampler", None)
        if sampler is None:
            raise AttributeError("xai_sdk.Client.sampler not found")
        chat_fn = getattr(sampler, "chat", None)
        if chat_fn is None:
            raise AttributeError("xai_sdk.Client.sampler.chat not found")
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] xai-sdk not found ({exc}), skipping.")
        return

    # xAI via OpenAI-compatible endpoint is already captured by the OpenAI patch.
    # Mark as patched so we don't waste cycles.
    _XAI_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# Agno — Agent.run (sync + async)
# ─────────────────────────────────────────────────────────────────────────────

_AGNO_PATCHED = False

def _patch_agno(paygent_client: Client) -> None:
    global _AGNO_PATCHED
    if _AGNO_PATCHED:
        return

    patched_any = False

    try:
        from agno.agent import Agent  # type: ignore
        _orig_run = Agent.run

        @wraps(_orig_run)
        def _patched_run(self, message=None, *args, **kwargs):
            agent_id   = kwargs.pop("paygent_agent_id",   None) or ""
            customer_id = kwargs.pop("paygent_customer_id", None) or ""
            indicator  = kwargs.pop("paygent_indicator",   None) or ""
            customer_name = kwargs.pop("paygent_customer_name", None) or ""

            if customer_id:
                try:
                    from .models import CustomerCreateOrGetRequest
                    paygent_client.create_or_get_customer(
                        CustomerCreateOrGetRequest(external_id=customer_id,
                                                   name=customer_name or customer_id)
                    )
                except Exception:
                    pass

            result = _orig_run(self, message, *args, **kwargs)

            if agent_id and customer_id and indicator:
                try:
                    # Extract from RunResponse.metrics if available
                    metrics = getattr(result, "metrics", {}) or {}
                    inp = metrics.get("input_tokens", [0])
                    out = metrics.get("output_tokens", [0])
                    inp = inp[0] if isinstance(inp, list) else inp
                    out = out[0] if isinstance(out, list) else out
                    model = getattr(getattr(self, "model", None), "id", "agno-agent")
                    if inp > 0 or out > 0:
                        _send_usage_tokens(paygent_client, agent_id, customer_id,
                                           indicator, ServiceProvider.OPENAI, model, inp, out)
                    else:
                        content = getattr(result, "content", "") or ""
                        _send_usage_strings(paygent_client, agent_id, customer_id, indicator,
                                            ServiceProvider.OPENAI, model,
                                            str(message or ""), str(content))
                except Exception as e:
                    logger.error(f"[Paygent] Agno tracking error: {e}")

            return result

        Agent.run = _patched_run
        patched_any = True
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] Agno not found ({exc}), skipping.")

    if patched_any:
        _AGNO_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# Haystack — Pipeline.run (sync)
# ─────────────────────────────────────────────────────────────────────────────

_HAYSTACK_PATCHED = False

def _patch_haystack(paygent_client: Client) -> None:
    global _HAYSTACK_PATCHED
    if _HAYSTACK_PATCHED:
        return

    try:
        from haystack import Pipeline  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] haystack not found ({exc}), skipping.")
        return

    try:
        _orig = Pipeline.run

        @wraps(_orig)
        def _patched(self, data, *args, **kwargs):
            agent_id   = kwargs.pop("paygent_agent_id",   None) or ""
            customer_id = kwargs.pop("paygent_customer_id", None) or ""
            indicator  = kwargs.pop("paygent_indicator",   None) or ""
            customer_name = kwargs.pop("paygent_customer_name", None) or ""

            if customer_id:
                try:
                    from .models import CustomerCreateOrGetRequest
                    paygent_client.create_or_get_customer(
                        CustomerCreateOrGetRequest(external_id=customer_id,
                                                   name=customer_name or customer_id)
                    )
                except Exception:
                    pass

            # Token counts come from the underlying LLM patches (OpenAI/Anthropic/etc.)
            return _orig(self, data, *args, **kwargs)

        Pipeline.run = _patched
    except Exception as e:
        logger.error(f"[Paygent] Haystack patch failed: {e}")

    _HAYSTACK_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# Google ADK — Runner.run_async
# ─────────────────────────────────────────────────────────────────────────────

_GOOGLE_ADK_PATCHED = False

def _patch_google_adk(paygent_client: Client) -> None:
    global _GOOGLE_ADK_PATCHED
    if _GOOGLE_ADK_PATCHED:
        return

    try:
        from google.adk.runners import Runner  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] google-adk not found ({exc}), skipping.")
        return

    try:
        _orig = Runner.run_async

        @wraps(_orig)
        async def _patched(self, *args, **kwargs):
            agent_id   = kwargs.pop("paygent_agent_id",   None) or ""
            customer_id = kwargs.pop("paygent_customer_id", None) or ""
            indicator  = kwargs.pop("paygent_indicator",   None) or ""
            customer_name = kwargs.pop("paygent_customer_name", None) or ""

            if customer_id:
                try:
                    from .models import CustomerCreateOrGetRequest
                    paygent_client.create_or_get_customer(
                        CustomerCreateOrGetRequest(external_id=customer_id,
                                                   name=customer_name or customer_id)
                    )
                except Exception:
                    pass

            # Underlying Gemini calls are already tracked via _patch_gemini_generate_content
            async for event in _orig(self, *args, **kwargs):
                yield event

        Runner.run_async = _patched
    except Exception as e:
        logger.error(f"[Paygent] Google ADK patch failed: {e}")

    _GOOGLE_ADK_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# Smolagents — MultiStepAgent.run
# ─────────────────────────────────────────────────────────────────────────────

_SMOLAGENTS_PATCHED = False

def _patch_smolagents(paygent_client: Client) -> None:
    global _SMOLAGENTS_PATCHED
    if _SMOLAGENTS_PATCHED:
        return

    try:
        from smolagents import MultiStepAgent  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] smolagents not found ({exc}), skipping.")
        return

    try:
        _orig = MultiStepAgent.run

        @wraps(_orig)
        def _patched(self, task, *args, **kwargs):
            agent_id   = kwargs.pop("paygent_agent_id",   None) or ""
            customer_id = kwargs.pop("paygent_customer_id", None) or ""
            indicator  = kwargs.pop("paygent_indicator",   None) or ""
            customer_name = kwargs.pop("paygent_customer_name", None) or ""

            if customer_id:
                try:
                    from .models import CustomerCreateOrGetRequest
                    paygent_client.create_or_get_customer(
                        CustomerCreateOrGetRequest(external_id=customer_id,
                                                   name=customer_name or customer_id)
                    )
                except Exception:
                    pass

            # Token counts come from the underlying LLM patches
            return _orig(self, task, *args, **kwargs)

        MultiStepAgent.run = _patched
    except Exception as e:
        logger.error(f"[Paygent] Smolagents patch failed: {e}")

    _SMOLAGENTS_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# Mem0 — Memory.add / Memory.search
# ─────────────────────────────────────────────────────────────────────────────

_MEM0_PATCHED = False

def _patch_mem0(paygent_client: Client) -> None:
    global _MEM0_PATCHED
    if _MEM0_PATCHED:
        return

    try:
        from mem0 import Memory  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] mem0 not found ({exc}), skipping.")
        return

    try:
        _orig_add = Memory.add

        @wraps(_orig_add)
        def _patched_add(self, messages, *args, **kwargs):
            agent_id   = kwargs.pop("paygent_agent_id",   None) or ""
            customer_id = kwargs.pop("paygent_customer_id", None) or ""
            indicator  = kwargs.pop("paygent_indicator",   None) or ""
            customer_name = kwargs.pop("paygent_customer_name", None) or ""

            if customer_id:
                try:
                    from .models import CustomerCreateOrGetRequest
                    paygent_client.create_or_get_customer(
                        CustomerCreateOrGetRequest(external_id=customer_id,
                                                   name=customer_name or customer_id)
                    )
                except Exception:
                    pass

            # Underlying LLM calls (OpenAI/Anthropic) are already tracked
            return _orig_add(self, messages, *args, **kwargs)

        Memory.add = _patched_add
    except Exception as e:
        logger.error(f"[Paygent] Mem0 patch failed: {e}")

    _MEM0_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# Xpander AI — XpanderClient (underlying LLM calls captured by existing patches)
# We patch at the agent-run level for customer registration.
# ─────────────────────────────────────────────────────────────────────────────

_XPANDER_PATCHED = False

def _patch_xpander(paygent_client: Client) -> None:
    global _XPANDER_PATCHED
    if _XPANDER_PATCHED:
        return

    try:
        from xpander_sdk import XpanderClient  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] xpander-sdk not found ({exc}), skipping.")
        return

    try:
        _orig = XpanderClient.run

        @wraps(_orig)
        def _patched(self, *args, **kwargs):
            agent_id   = kwargs.pop("paygent_agent_id",   None) or ""
            customer_id = kwargs.pop("paygent_customer_id", None) or ""
            indicator  = kwargs.pop("paygent_indicator",   None) or ""
            customer_name = kwargs.pop("paygent_customer_name", None) or ""

            if customer_id:
                try:
                    from .models import CustomerCreateOrGetRequest
                    paygent_client.create_or_get_customer(
                        CustomerCreateOrGetRequest(external_id=customer_id,
                                                   name=customer_name or customer_id)
                    )
                except Exception:
                    pass

            return _orig(self, *args, **kwargs)

        XpanderClient.run = _patched
    except Exception as e:
        logger.error(f"[Paygent] Xpander patch failed: {e}")

    _XPANDER_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# OpenAI — Videos (sync + async)
# ─────────────────────────────────────────────────────────────────────────────

_OPENAI_VIDEOS_PATCHED = False

def _patch_openai_videos(paygent_client: Client) -> None:
    global _OPENAI_VIDEOS_PATCHED
    if _OPENAI_VIDEOS_PATCHED:
        return
    try:
        from openai.resources.videos import Videos, AsyncVideos  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] OpenAI Videos not found ({exc}), skipping.")
        return

    # ── Sync ──────────────────────────────────────────────────────────────────
    try:
        _orig_sync = Videos.create

        @wraps(_orig_sync)
        def _patched_sync(self, *args, **kwargs):
            model     = kwargs.get("model", "unknown")
            seconds   = float(kwargs.get("seconds", 5.0)) # Default to 5s if not specified
            size      = kwargs.get("size", "default")
            
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = _orig_sync(self, *args, **kwargs)
            try:
                if agent_id and customer_id and indicator:
                    _send_usage_video(paygent_client, agent_id, customer_id, indicator,
                                      ServiceProvider.OPENAI, model, seconds, size, None)
            except Exception as e:
                logger.error(f"[Paygent] OpenAI Videos sync tracking error: {e}")
            return response

        Videos.create = _patched_sync
    except Exception as e:
        logger.error(f"[Paygent] OpenAI Videos sync patch failed: {e}")

    # ── Async ─────────────────────────────────────────────────────────────────
    try:
        _orig_async = AsyncVideos.create

        @wraps(_orig_async)
        async def _patched_async(self, *args, **kwargs):
            model     = kwargs.get("model", "unknown")
            seconds   = float(kwargs.get("seconds", 5.0))
            size      = kwargs.get("size", "default")
            
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = await _orig_async(self, *args, **kwargs)
            try:
                if agent_id and customer_id and indicator:
                    _send_usage_video(paygent_client, agent_id, customer_id, indicator,
                                      ServiceProvider.OPENAI, model, seconds, size, None)
            except Exception as e:
                logger.error(f"[Paygent] OpenAI Videos async tracking error: {e}")
            return response

        AsyncVideos.create = _patched_async
    except Exception as e:
        logger.error(f"[Paygent] OpenAI Videos async patch failed: {e}")

    _OPENAI_VIDEOS_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# Google GenAI — Models.generate_videos + AsyncModels.generate_videos
# ─────────────────────────────────────────────────────────────────────────────

_GEMINI_VIDEOS_PATCHED = False

def _patch_gemini_videos(paygent_client: Client) -> None:
    global _GEMINI_VIDEOS_PATCHED
    if _GEMINI_VIDEOS_PATCHED:
        return
    try:
        from google.genai.models import Models, AsyncModels  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] google.genai.models not found ({exc}), skipping.")
        return

    # ── Sync ──────────────────────────────────────────────────────────────────
    try:
        _orig_sync = getattr(Models, "generate_videos", None)
        if _orig_sync:
            @wraps(_orig_sync)
            def _patched_sync(self, *args, **kwargs):
                model    = kwargs.get("model", "veo-3.1-generate-preview")
                
                # Extract duration from config or direct kwargs
                config = kwargs.get("config", {})
                if isinstance(config, dict):
                    seconds = float(config.get("duration_seconds", kwargs.get("seconds", 5.0)))
                else:
                    seconds = float(getattr(config, "duration_seconds", kwargs.get("seconds", 5.0)))

                agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
                response = _orig_sync(self, *args, **kwargs)
                try:
                    if agent_id and customer_id and indicator:
                        _send_usage_video(paygent_client, agent_id, customer_id, indicator,
                                          ServiceProvider.GOOGLE, model, seconds, None, None)
                except Exception as e:
                    logger.error(f"[Paygent] Gemini Videos sync tracking error: {e}")
                return response

            Models.generate_videos = _patched_sync
    except Exception as e:
        logger.error(f"[Paygent] Gemini Videos sync patch failed: {e}")

    # ── Async ─────────────────────────────────────────────────────────────────
    try:
        _orig_async = getattr(AsyncModels, "generate_videos", None)
        if _orig_async:
            @wraps(_orig_async)
            async def _patched_async(self, *args, **kwargs):
                model    = kwargs.get("model", "veo-3.1-generate-preview")
                
                # Extract duration from config or direct kwargs
                config = kwargs.get("config", {})
                if isinstance(config, dict):
                    seconds = float(config.get("duration_seconds", kwargs.get("seconds", 5.0)))
                else:
                    seconds = float(getattr(config, "duration_seconds", kwargs.get("seconds", 5.0)))

                agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
                response = await _orig_async(self, *args, **kwargs)
                try:
                    if agent_id and customer_id and indicator:
                        _send_usage_video(paygent_client, agent_id, customer_id, indicator,
                                          ServiceProvider.GOOGLE, model, seconds, None, None)
                except Exception as e:
                    logger.error(f"[Paygent] Gemini Videos async tracking error: {e}")
                return response

            AsyncModels.generate_videos = _patched_async
    except Exception as e:
        logger.error(f"[Paygent] Gemini Videos async patch failed: {e}")

    _GEMINI_VIDEOS_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# RunwayML — image_to_video + video_to_video (Official SDK)
# ─────────────────────────────────────────────────────────────────────────────

_RUNWAYML_PATCHED = False

def _patch_runwayml(paygent_client: Client) -> None:
    global _RUNWAYML_PATCHED
    if _RUNWAYML_PATCHED:
        return
    try:
        from runwayml.resources.image_to_video import ImageToVideo, AsyncImageToVideo  # type: ignore
    except (ImportError, AttributeError) as exc:
        logger.debug(f"[Paygent] runwayml not found ({exc}), skipping.")
        return

    # ── Sync ──────────────────────────────────────────────────────────────────
    try:
        _orig_sync_itv = ImageToVideo.create

        @wraps(_orig_sync_itv)
        def _patched_sync_itv(self, *args, **kwargs):
            model     = kwargs.get("model", "gen3a_turbo")
            seconds   = float(kwargs.get("seconds", 5.0)) 
            
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = _orig_sync_itv(self, *args, **kwargs)
            try:
                if agent_id and customer_id and indicator:
                    _send_usage_video(paygent_client, agent_id, customer_id, indicator,
                                      ServiceProvider.RUNWAY, model, seconds, None, None)
            except Exception as e:
                logger.error(f"[Paygent] RunwayML sync tracking error: {e}")
            return response

        ImageToVideo.create = _patched_sync_itv
    except Exception as e:
        logger.error(f"[Paygent] RunwayML sync patch failed: {e}")

    # ── Async ─────────────────────────────────────────────────────────────────
    try:
        _orig_async_itv = AsyncImageToVideo.create

        @wraps(_orig_async_itv)
        async def _patched_async_itv(self, *args, **kwargs):
            model     = kwargs.get("model", "gen3a_turbo")
            seconds   = float(kwargs.get("seconds", 5.0))
            
            agent_id, customer_id, indicator = _extract_paygent_kwargs(paygent_client, kwargs)
            response = await _orig_async_itv(self, *args, **kwargs)
            try:
                if agent_id and customer_id and indicator:
                    _send_usage_video(paygent_client, agent_id, customer_id, indicator,
                                      ServiceProvider.RUNWAY, model, seconds, None, None)
            except Exception as e:
                logger.error(f"[Paygent] RunwayML async tracking error: {e}")
            return response

        AsyncImageToVideo.create = _patched_async_itv
    except Exception as e:
        logger.error(f"[Paygent] RunwayML async patch failed: {e}")

    _RUNWAYML_PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# Orchestrator
# ─────────────────────────────────────────────────────────────────────────────

def init_patching(paygent_client: Client) -> None:
    """
    Apply ALL available patches.  Each is individually guarded so a failure
    in one provider NEVER affects others or the user's application.
    """
    for patch_fn in (
        # ── LLM Providers ───────────────────────────────────────────────────
        _patch_openai_chat_completions,     # OpenAI Chat Completions (sync + async)
        _patch_openai_responses,            # OpenAI Responses API (o1/o3, sync + async)
        _patch_anthropic_messages,          # Anthropic Messages (sync + async)
        _patch_gemini_generate_content,     # Google Gemini (sync + async)
        _patch_mistral_chat,                # Mistral AI
        _patch_cohere_chat,                 # Cohere
        _patch_groq,                        # Groq (OpenAI-compatible, sync + async)
        _patch_litellm,                     # LiteLLM (proxy for all providers)
        _patch_ollama,                      # Ollama (local models)
        _patch_ibm_watsonx,                 # IBM WatsonX AI
        _patch_xai,                         # xAI / Grok (via OpenAI-compatible)
        # ── Orchestration Frameworks ─────────────────────────────────────────
        _patch_langchain,                   # LangChain BaseChatModel + BaseLLM
        _patch_ag2,                         # AG2 / AutoGen
        _patch_crewai,                      # CrewAI
        _patch_openai_agents,               # OpenAI Agents SDK
        _patch_llamaindex,                  # LlamaIndex
        _patch_agno,                        # Agno
        _patch_haystack,                    # Haystack
        _patch_google_adk,                  # Google ADK
        _patch_smolagents,                  # Smolagents (HuggingFace)
        _patch_mem0,                        # Mem0
        _patch_xpander,                     # Xpander AI
        # ── Video Generation ──────────────────────────────────────────────────
        _patch_openai_videos,               # OpenAI Videos (Sora)
        _patch_gemini_videos,               # Google GenAI Videos (Veo)
        _patch_runwayml,                    # RunwayML
    ):
        try:
            patch_fn(paygent_client)
        except Exception as e:
            try:
                logger.error(f"[Paygent] Unexpected error in {patch_fn.__name__}: {e}")
            except Exception:
                pass
