import os
import json
import uuid
import asyncio
import inspect
import logging
import warnings
import traceback
from abc import abstractmethod
from enum import Enum
from typing import Any, Set, Union, Callable, Optional, Sequence, TypedDict
from datetime import datetime, timezone
from typing_extensions import deprecated

import nest_asyncio  # type: ignore
from wrapt import ObjectProxy  # type: ignore

from payi import Payi, AsyncPayi
from payi.types import IngestUnitsParams
from payi.lib.helpers import PayiHeaderNames
from payi.types.ingest_response import IngestResponse
from payi.types.ingest_units_params import Units
from payi.types.pay_i_common_models_api_router_header_info_param import PayICommonModelsAPIRouterHeaderInfoParam

from .helpers import PayiCategories
from .Stopwatch import Stopwatch

global _g_logger
_g_logger: logging.Logger = logging.getLogger("payi.instrument")

class _ProviderRequest:
    def __init__(self, instrumentor: '_PayiInstrumentor', category: str, streaming_type: '_StreamingType'):
        self._instrumentor: '_PayiInstrumentor' = instrumentor
        self._estimated_prompt_tokens: Optional[int] = None
        self._category: str = category
        self._ingest: IngestUnitsParams = { "category": category, "units": {} } # type: ignore
        self._streaming_type: '_StreamingType' = streaming_type

    def process_chunk(self, _chunk: Any) -> bool:
        return True

    def process_synchronous_response(self, response: Any, log_prompt_and_response: bool, kwargs: Any) -> Optional[object]:  # noqa: ARG002
        return None
    
    @abstractmethod
    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', args: Sequence[Any], kwargs: Any) -> bool:
        ...
    
    def process_request_prompt(self, prompt: 'dict[str, Any]', args: Sequence[Any], kwargs: 'dict[str, Any]') -> None:
        ...
    
    def is_bedrock(self) -> bool:
        return self._category == PayiCategories.aws_bedrock

    def is_vertex(self) -> bool:
        return self._category == PayiCategories.google_vertex 
       
    def process_exception(self, exception: Exception, kwargs: Any, ) -> bool: # noqa: ARG002
        self.exception_to_semantic_failure(exception)
        return True
    
    @property
    def streaming_type(self) -> '_StreamingType':
        return self._streaming_type

    def exception_to_semantic_failure(self, e: Exception) -> None:
        exception_str = f"{type(e).__name__}"
    
        fields: list[str] = []
    
        for attr in dir(e):
            if not attr.startswith("__"):
                try:
                    value = getattr(e, attr)
                    if value and not inspect.ismethod(value) and not inspect.isfunction(value) and not callable(value):
                        fields.append(f"{attr}={value}")
                except Exception as _ex:
                    pass
 
        properties: 'dict[str,str]' = { "system.failure": exception_str }
        if fields:
            failure_description = ",".join(fields)
            properties["system.failure.description"] = failure_description[:128]
        self._ingest["properties"] = properties

        if "http_status_code" not in self._ingest:
            # use a non existent http status code so when presented to the user, the origin is clear
            self._ingest["http_status_code"] = 299

class PayiInstrumentConfig(TypedDict, total=False):
    proxy: bool
    global_instrumentation: bool
    limit_ids: Optional["list[str]"]
    experience_name: Optional[str] = deprecated("experience_name is deprecated, use use_case_name instead") # type: ignore
    experience_id: Optional[str] = deprecated("experience_id is deprecated, use use_case_id instead") # type: ignore
    use_case_name: Optional[str]
    use_case_id: Optional[str]
    use_case_version: Optional[int]
    user_id: Optional[str]

class _Context(TypedDict, total=False):
    proxy: Optional[bool]
    experience_name: Optional[str]
    experience_id: Optional[str]
    use_case_name: Optional[str]
    use_case_id: Optional[str]
    use_case_version: Optional[int]
    use_case_step: Optional[str]
    limit_ids: Optional['list[str]']
    user_id: Optional[str]
    request_tags: Optional["list[str]"]
    price_as_category: Optional[str]
    price_as_resource: Optional[str]
    resource_scope: Optional[str]

class _IsStreaming(Enum):
    false = 0
    true = 1 
    kwargs = 2

class _StreamingType(Enum):
    generator = 0
    iterator = 1
    stream_manager = 2

class _TrackContext:
    def __init__(
        self,
        context: _Context,
    ) -> None:
        self._context = context

    def __enter__(self) -> Any:
        if not _instrumentor:
            return self
        
        _instrumentor.__enter__()
        _instrumentor._init_current_context(**self._context)
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        if _instrumentor:
            _instrumentor.__exit__(exc_type, exc_val, exc_tb)

class _PayiInstrumentor:
    def __init__(
        self,
        payi: Optional[Payi],
        apayi: Optional[AsyncPayi],
        instruments: Union[Set[str], None] = None,
        log_prompt_and_response: bool = True,
        logger: Optional[logging.Logger] = None,
        prompt_and_response_logger: Optional[
            Callable[[str, "dict[str, str]"], None]
        ] = None,  # (request id, dict of data to store) -> None
        global_config: PayiInstrumentConfig = {},
        caller_filename: str = ""
    ):
        global _g_logger
        self._logger: logging.Logger = logger if logger else _g_logger

        self._payi: Optional[Payi] = payi
        self._apayi: Optional[AsyncPayi] = apayi

        if self._payi:
            _g_logger.debug(f"Pay-i instrumentor initialized with Payi instance: {self._payi}")
        if self._apayi:
            _g_logger.debug(f"Pay-i instrumentor initialized with AsyncPayi instance: {self._apayi}")            

        self._context_stack: list[_Context] = []  # Stack of context dictionaries
        self._log_prompt_and_response: bool = log_prompt_and_response
        self._prompt_and_response_logger: Optional[Callable[[str, dict[str, str]], None]] = prompt_and_response_logger

        self._blocked_limits: set[str] = set()
        self._exceeded_limits: set[str] = set()

        # default is instrument and ingest metrics
        self._proxy_default: bool = global_config.get("proxy", False)

        global_instrumentation = global_config.pop("global_instrumentation", True)

        if instruments is None or "*" in instruments:
            self._instrument_all()
        else:
            self._instrument_specific(instruments)

        if global_instrumentation:
            if "proxy" not in global_config:
                global_config["proxy"] = self._proxy_default

            # Use default clients if not provided for global ingest instrumentation
            if not self._payi and not self._apayi:
                self._payi = Payi()
                self._apayi = AsyncPayi()

            if "use_case_name" not in global_config and caller_filename:
                description = f"Default use case for {caller_filename}.py"
                try:
                    if self._payi:
                        self._payi.use_cases.definitions.create(name=caller_filename, description=description)
                    elif self._apayi:
                        self._call_async_use_case_definition_create(use_case_name=caller_filename, use_case_description=description)
                    global_config["use_case_name"] = caller_filename
                except Exception as e:
                    self._logger.error(f"Error creating default use case definition based on file name {caller_filename}: {e}")

            self.__enter__()
            # _init_current_context will update the currrent context stack location
            self._init_current_context(**global_config) # type: ignore

    def _instrument_all(self) -> None:
        self._instrument_openai()
        self._instrument_anthropic()
        self._instrument_aws_bedrock()
        self._instrument_google_vertex()
        self._instrument_google_genai()

    def _instrument_specific(self, instruments: Set[str]) -> None:
        if PayiCategories.openai in instruments or PayiCategories.azure_openai in instruments:
            self._instrument_openai()
        if PayiCategories.anthropic in instruments:
            self._instrument_anthropic()
        if PayiCategories.aws_bedrock in instruments:
            self._instrument_aws_bedrock()
        if PayiCategories.google_vertex in instruments:
            self._instrument_google_vertex()
            self._instrument_google_genai()

    def _instrument_openai(self) -> None:
        from .OpenAIInstrumentor import OpenAiInstrumentor

        try:
            OpenAiInstrumentor.instrument(self)

        except Exception as e:
            self._logger.error(f"Error instrumenting OpenAI: {e}")

    def _instrument_anthropic(self) -> None:
        from .AnthropicInstrumentor import AnthropicInstrumentor

        try:
            AnthropicInstrumentor.instrument(self)

        except Exception as e:
            self._logger.error(f"Error instrumenting Anthropic: {e}")

    def _instrument_aws_bedrock(self) -> None:
        from .BedrockInstrumentor import BedrockInstrumentor

        try:
            BedrockInstrumentor.instrument(self)

        except Exception as e:
            self._logger.error(f"Error instrumenting AWS bedrock: {e}")

    def _instrument_google_vertex(self) -> None:
        from .VertexInstrumentor import VertexInstrumentor

        try:
            VertexInstrumentor.instrument(self)

        except Exception as e:
            self._logger.error(f"Error instrumenting Google Vertex: {e}")

    def _instrument_google_genai(self) -> None:
        from .GoogleGenAiInstrumentor import GoogleGenAiInstrumentor

        try:
            GoogleGenAiInstrumentor.instrument(self)

        except Exception as e:
            self._logger.error(f"Error instrumenting Google GenAi: {e}")

    def _process_ingest_units(self, ingest_units: IngestUnitsParams, log_data: 'dict[str, str]') -> bool:
        if int(ingest_units.get("http_status_code") or 0) < 400:
            units = ingest_units.get("units", {})
            if not units or all(unit.get("input", 0) == 0 and unit.get("output", 0) == 0 for unit in units.values()):
                self._logger.error('No units to ingest!')
                return False

        if self._log_prompt_and_response and self._prompt_and_response_logger:
            response_json = ingest_units.pop("provider_response_json", None)
            request_json = ingest_units.pop("provider_request_json", None)
            stack_trace = ingest_units.get("properties", {}).pop("system.stack_trace", None)  # type: ignore

            if response_json is not None:
                # response_json is a list of strings, convert a single json string
                log_data["provider_response_json"] = json.dumps(response_json)
            if request_json is not None:
                log_data["provider_request_json"] = request_json
            if stack_trace is not None:
                log_data["stack_trace"] = stack_trace

        return True

    def _process_ingest_units_response(self, ingest_response: IngestResponse) -> None:
        if ingest_response.xproxy_result.limits:
            for limit_id, state in ingest_response.xproxy_result.limits.items():
                removeBlockedId: bool = False

                if state.state == "blocked":
                    self._blocked_limits.add(limit_id)
                elif state.state == "exceeded":
                    self._exceeded_limits.add(limit_id)
                    removeBlockedId = True
                elif state.state == "ok":
                    removeBlockedId = True

                # opportunistically remove blocked limits
                if removeBlockedId:
                    self._blocked_limits.discard(limit_id)

    async def _aingest_units(self, ingest_units: IngestUnitsParams) -> Optional[IngestResponse]:
        ingest_response: Optional[IngestResponse] = None
    
        self._logger.debug(f"_aingest_units")

        # return early if there are no units to ingest and on a successul ingest request
        log_data: 'dict[str,str]' = {}
        if not self._process_ingest_units(ingest_units, log_data):
            self._logger.debug(f"_aingest_units: exit early")
            return None

        try:
            if self._apayi:    
                ingest_response = await self._apayi.ingest.units(**ingest_units)
            elif self._payi:
                ingest_response = self._payi.ingest.units(**ingest_units)
            else:
                self._logger.error("No payi instance to ingest units")
                return None         

            self._logger.debug(f"_aingest_units: success ({ingest_response})")

            if ingest_response:
                self._process_ingest_units_response(ingest_response)

                if ingest_response and self._log_prompt_and_response and self._prompt_and_response_logger:
                    request_id = ingest_response.xproxy_result.request_id
                    self._prompt_and_response_logger(request_id, log_data)  # type: ignore

            return ingest_response
        except Exception as e:
            self._logger.error(f"Error Pay-i ingesting request: {e}")
    
        return None         
    
    def _call_async_use_case_definition_create(self, use_case_name: str, use_case_description: str) -> None:
        if not self._apayi:
            return

        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = None
        
        try:
            if loop and loop.is_running():
                nest_asyncio.apply(loop) # type: ignore
                asyncio.run(self._apayi.use_cases.definitions.create(name=use_case_name, description=use_case_description))                
            else:
                # When there's no running loop, create a new one
                asyncio.run(self._apayi.use_cases.definitions.create(name=use_case_name, description=use_case_description))
        except Exception as e:
            self._logger.error(f"Error calling async use_cases.definitions.create synchronously: {e}")

    def _call_aingest_sync(self, ingest_units: IngestUnitsParams) -> Optional[IngestResponse]:
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = None
        
        try:
            if loop and loop.is_running():
                nest_asyncio.apply(loop) # type: ignore
                return asyncio.run(self._aingest_units(ingest_units))                
            else:
                # When there's no running loop, create a new one
                return asyncio.run(self._aingest_units(ingest_units))
        except Exception as e:
            self._logger.error(f"Error calling aingest_units synchronously: {e}")
        return None
        
    def _ingest_units(self, ingest_units: IngestUnitsParams) -> Optional[IngestResponse]:
        ingest_response: Optional[IngestResponse] = None
        
        self._logger.debug(f"_ingest_units")

        # return early if there are no units to ingest and on a successul ingest request
        log_data: 'dict[str,str]' = {}
        if not self._process_ingest_units(ingest_units, log_data):
            self._logger.debug(f"_ingest_units: exit early")
            return None

        try:
            if self._payi:
                ingest_response = self._payi.ingest.units(**ingest_units)
                self._logger.debug(f"_ingest_units: success ({ingest_response})")

                self._process_ingest_units_response(ingest_response)

                if self._log_prompt_and_response and self._prompt_and_response_logger:
                    request_id = ingest_response.xproxy_result.request_id
                    self._prompt_and_response_logger(request_id, log_data)  # type: ignore

                return ingest_response
            elif self._apayi:
                # task runs async. aingest_units will invoke the callback and post process
                ingest_response = self._call_aingest_sync(ingest_units)
                self._logger.debug(f"_ingest_units: apayi success ({ingest_response})")
                return ingest_response
            else:
                self._logger.error("No payi instance to ingest units")

        except Exception as e:
            self._logger.error(f"Error Pay-i ingesting request: {e}")
        
        return None

    def _setup_call_func(
        self
        ) -> _Context:

        if len(self._context_stack) > 0:
            # copy current context into the upcoming context
            return self._context_stack[-1].copy()

        return {}

    def _init_current_context(
        self,
        proxy: Optional[bool] = None,
        limit_ids: Optional["list[str]"] = None,
        experience_name: Optional[str] = None,
        experience_id: Optional[str] = None,
        use_case_name: Optional[str]= None,
        use_case_id: Optional[str]= None,
        use_case_version: Optional[int]= None,
        use_case_step: Optional[str]= None,
        user_id: Optional[str]= None,
        request_tags: Optional["list[str]"] = None,
        price_as_category: Optional[str] = None,
        price_as_resource: Optional[str] = None,
        resource_scope: Optional[str] = None,
        ) -> None:

        # there will always be a current context
        context: _Context = self.get_context() # type: ignore
        parent_context: _Context = self._context_stack[-2] if len(self._context_stack) > 1 else {}

        parent_proxy = parent_context.get("proxy", self._proxy_default)
        context["proxy"] = proxy if proxy else parent_proxy

        parent_experience_name = parent_context.get("experience_name", None)
        parent_experience_id = parent_context.get("experience_id", None)

        if experience_name is None:
            # If no experience_name specified, use previous values
            context["experience_name"] = parent_experience_name
            context["experience_id"] = parent_experience_id
        elif len(experience_name) == 0:
            # Empty string explicitly blocks inheriting from the parent state
            context["experience_name"] = None
            context["experience_id"] = None
        else:
            # Check if experience_name is the same as the previous one
            if experience_name == parent_experience_name:
                # Same experience name, use previous ID unless new one specified
                context["experience_name"] = experience_name
                context["experience_id"] = experience_id if experience_id else parent_experience_id
            else:
                # Different experience name, use specified ID or generate one
                context["experience_name"] = experience_name
                context["experience_id"] = experience_id if experience_id else str(uuid.uuid4())

        parent_use_case_name = parent_context.get("use_case_name", None)
        parent_use_case_id = parent_context.get("use_case_id", None)
        parent_use_case_version = parent_context.get("use_case_version", None)

        if use_case_name is None:
            # If no use_case_name specified, use previous values
            context["use_case_name"] = parent_use_case_name
            context["use_case_id"] = parent_use_case_id
            context["use_case_version"] = parent_use_case_version
        elif len(use_case_name) == 0:
            # Empty string explicitly blocks inheriting from the parent state
            context["use_case_name"] = None
            context["use_case_id"] = None
            context["use_case_version"] = None
            context["use_case_step"] = None
        else:
            if use_case_name == parent_use_case_name:
                # Same use case name, use previous ID unless new one specified
                context["use_case_name"] = use_case_name
                context["use_case_id"] = use_case_id if use_case_id else parent_use_case_id
                context["use_case_version"] = use_case_version if use_case_version else parent_use_case_version
            else:
                # Different use case name, use specified ID or generate one
                context["use_case_name"] = use_case_name
                context["use_case_id"] = use_case_id if use_case_id else str(uuid.uuid4())
                context["use_case_version"] = use_case_version if use_case_version else None

        parent_limit_ids = parent_context.get("limit_ids", None)
        if limit_ids is None:
            # use the parent limit_ids if it exists
            context["limit_ids"] = parent_limit_ids
        elif len(limit_ids) == 0:
            # caller passing an empty array explicitly blocks inheriting from the parent state
            context["limit_ids"] = None
        else:
            # union of new and parent lists if the parent context contains limit ids
            context["limit_ids"] = list(set(limit_ids) | set(parent_limit_ids)) if parent_limit_ids else limit_ids

        parent_user_id = parent_context.get("user_id", None)
        if user_id is None:
            # use the parent user_id if it exists
            context["user_id"] = parent_user_id
        elif len(user_id) == 0:
            # caller passing an empty string explicitly blocks inheriting from the parent state
            context["user_id"] = None
        else:
            context["user_id"] = user_id

        if use_case_step and (context["use_case_name"] or context["use_case_id"]):
            context["use_case_step"] = use_case_step
        if request_tags:
            context["request_tags"] = request_tags
        if price_as_category:
            context["price_as_category"] = price_as_category
        if price_as_resource:
            context["price_as_resource"] = price_as_resource
        if resource_scope:
            context["resource_scope"] = resource_scope
        
    async def _acall_func(
        self,
        func: Any,
        proxy: Optional[bool],
        limit_ids: Optional["list[str]"],
        experience_name: Optional[str],
        experience_id: Optional[str],
        use_case_name: Optional[str],
        use_case_id: Optional[str],
        use_case_version: Optional[int],        
        user_id: Optional[str],
        *args: Any,
        **kwargs: Any,
    ) -> Any:
        with self:
            self._init_current_context(
                proxy=proxy,
                limit_ids=limit_ids,
                experience_name=experience_name,
                experience_id=experience_id,
                use_case_name=use_case_name,
                use_case_id=use_case_id,
                use_case_version=use_case_version,
                user_id=user_id)
            return await func(*args, **kwargs)

    def _call_func(
        self,
        func: Any,
        proxy: Optional[bool],
        limit_ids: Optional["list[str]"],
        experience_name: Optional[str],
        experience_id: Optional[str],
        use_case_name: Optional[str],
        use_case_id: Optional[str],
        use_case_version: Optional[int],        
        user_id: Optional[str],
        *args: Any,
        **kwargs: Any,
    ) -> Any:
        with self:
            self._init_current_context(
                proxy=proxy,
                limit_ids=limit_ids,
                experience_name=experience_name,
                experience_id=experience_id,
                use_case_name=use_case_name,
                use_case_id=use_case_id,
                use_case_version=use_case_version,
                user_id=user_id)
            return func(*args, **kwargs)

    def __enter__(self) -> Any:
        # Push a new context dictionary onto the stack
        self._context_stack.append({})
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        # Pop the current context off the stack
        if self._context_stack:
            self._context_stack.pop()

    def get_context(self) -> Optional[_Context]:
        # Return the current top of the stack
        return self._context_stack[-1] if self._context_stack else None

    def get_context_safe(self) -> _Context:
        # Return the current top of the stack
        return self.get_context() or {}

    def _prepare_ingest(
        self,
        request: _ProviderRequest,
        ingest_extra_headers: "dict[str, str]", # do not coflict potential kwargs["extra_headers"]
        args: Sequence[Any],
        kwargs: 'dict[str, Any]',
    ) -> None:

        limit_ids = ingest_extra_headers.pop(PayiHeaderNames.limit_ids, None)
        request_tags = ingest_extra_headers.pop(PayiHeaderNames.request_tags, None)

        experience_name = ingest_extra_headers.pop(PayiHeaderNames.experience_name, None)
        experience_id = ingest_extra_headers.pop(PayiHeaderNames.experience_id, None)

        use_case_name = ingest_extra_headers.pop(PayiHeaderNames.use_case_name, None)
        use_case_id = ingest_extra_headers.pop(PayiHeaderNames.use_case_id, None)
        use_case_version = ingest_extra_headers.pop(PayiHeaderNames.use_case_version, None)
        use_case_step = ingest_extra_headers.pop(PayiHeaderNames.use_case_step, None)

        user_id = ingest_extra_headers.pop(PayiHeaderNames.user_id, None)

        if limit_ids:
            request._ingest["limit_ids"] = limit_ids.split(",")
        if request_tags:
            request._ingest["request_tags"] = request_tags.split(",")
        if experience_name:
            request._ingest["experience_name"] = experience_name
        if experience_id:
            request._ingest["experience_id"] = experience_id
        if use_case_name:
            request._ingest["use_case_name"] = use_case_name
        if use_case_id:
            request._ingest["use_case_id"] = use_case_id
        if use_case_version:
            request._ingest["use_case_version"] = int(use_case_version)
        if use_case_step:
            request._ingest["use_case_step"] = use_case_step
        if user_id:
            request._ingest["user_id"] = user_id

        if len(ingest_extra_headers) > 0:
            request._ingest["provider_request_headers"] = [PayICommonModelsAPIRouterHeaderInfoParam(name=k, value=v) for k, v in ingest_extra_headers.items()]

        provider_prompt: "dict[str, Any]" = {}
        for k, v in kwargs.items():
            if k == "messages":
                provider_prompt[k] = [m.model_dump() if hasattr(m, "model_dump") else m for m in v]
            elif k in ["extra_headers", "extra_query"]:
                pass
            else:
                try:
                    if hasattr(v, "to_dict"):
                        provider_prompt[k] = v.to_dict()
                    else:
                        json.dumps(v)
                        provider_prompt[k] = v
                except Exception as _e:
                    pass

        request.process_request_prompt(provider_prompt, args, kwargs)

        if self._log_prompt_and_response:
            request._ingest["provider_request_json"] = json.dumps(provider_prompt)

        request._ingest["event_timestamp"] = datetime.now(timezone.utc)
        
    async def async_invoke_wrapper(
        self,
        request: _ProviderRequest,
        is_streaming: _IsStreaming,
        wrapped: Any,
        instance: Any,
        args: Sequence[Any],
        kwargs: 'dict[str, Any]',
    ) -> Any:
        self._logger.debug(f"async_invoke_wrapper: instance {instance}, category {request._category}")

        context = self.get_context()

        # Bedrock client does not have an async method

        if not context:   
            self._logger.debug(f"async_invoke_wrapper: no instrumentation context, exit early")

            # wrapped function invoked outside of decorator scope
            return await wrapped(*args, **kwargs)

        # after _udpate_headers, all metadata to add to ingest is in extra_headers, keyed by the xproxy-xxx header name
        extra_headers: Optional[dict[str, str]] = kwargs.get("extra_headers")
        if extra_headers is None:
            extra_headers = {}
        self._update_extra_headers(context, extra_headers)

        if context.get("proxy", self._proxy_default):
            if "extra_headers" not in kwargs and extra_headers:
                kwargs["extra_headers"] = extra_headers

            self._logger.debug(f"async_invoke_wrapper: sending proxy request")

            return await wrapped(*args, **kwargs)

        current_frame = inspect.currentframe()
        # f_back excludes the current frame, strip() cleans up whitespace and newlines
        stack = [frame.strip() for frame in traceback.format_stack(current_frame.f_back)]  # type: ignore

        request._ingest['properties'] = { 'system.stack_trace': json.dumps(stack) }

        if request.process_request(instance, extra_headers, args, kwargs) is False:
            self._logger.debug(f"async_invoke_wrapper: calling wrapped instance")
            return await wrapped(*args, **kwargs)

        sw = Stopwatch()
        stream: bool = False
        
        if is_streaming == _IsStreaming.kwargs:
            stream = kwargs.get("stream", False)
        elif is_streaming == _IsStreaming.true:
            stream = True
        else:
            stream = False

        try:
            self._prepare_ingest(request, extra_headers, args, kwargs)
            self._logger.debug(f"async_invoke_wrapper: calling wrapped instance (stream={stream})")

            sw.start()
            response = await wrapped(*args, **kwargs)

        except Exception as e:  # pylint: disable=broad-except
            sw.stop()
            duration = sw.elapsed_ms_int()

            self._logger.debug(f"invoke_wrapper: calling wrapped instance exception {e}")

            if request.process_exception(e, kwargs):
                request._ingest["end_to_end_latency_ms"] = duration
                await self._aingest_units(request._ingest)

            raise e

        if stream:
            if request.streaming_type == _StreamingType.generator:
                return _GeneratorWrapper(
                    generator=response,
                    instance=instance,
                    instrumentor=self,
                    stopwatch=sw,
                    request=request,
                    )
            elif request.streaming_type == _StreamingType.stream_manager:
                return _StreamManagerWrapper(
                    stream_manager=response,
                    instance=instance,
                    instrumentor=self,
                    stopwatch=sw,
                    request=request,
                )
            else:
                return _StreamIteratorWrapper(
                    response=response,
                    instance=instance,
                    instrumentor=self,
                    stopwatch=sw,
                    request=request,
                )

        sw.stop()
        duration = sw.elapsed_ms_int()
        request._ingest["end_to_end_latency_ms"] = duration
        request._ingest["http_status_code"] = 200

        return_result: Any = request.process_synchronous_response(
            response=response,
            log_prompt_and_response=self._log_prompt_and_response,
            kwargs=kwargs)

        if return_result:
            self._logger.debug(f"async_invoke_wrapper: process sync response return")
            return return_result

        await self._aingest_units(request._ingest)

        self._logger.debug(f"async_invoke_wrapper: finished")
        return response

    def invoke_wrapper(
        self,
        request: _ProviderRequest,
        is_streaming: _IsStreaming,
        wrapped: Any,
        instance: Any,
        args: Sequence[Any],
        kwargs: 'dict[str, Any]',
    ) -> Any:
        self._logger.debug(f"invoke_wrapper: instance {instance}, category {request._category}")

        context = self.get_context()

        if not context:
            if request.is_bedrock():
                # boto3 doesn't allow extra_headers
                kwargs.pop("extra_headers", None)

            self._logger.debug(f"invoke_wrapper: no instrumentation context, exit early")

            # wrapped function invoked outside of decorator scope
            return wrapped(*args, **kwargs)

        # after _udpate_headers, all metadata to add to ingest is in extra_headers, keyed by the xproxy-xxx header name
        extra_headers: Optional[dict[str, str]] = kwargs.get("extra_headers")
        if extra_headers is None:
            extra_headers = {}
        self._update_extra_headers(context, extra_headers)

        if context.get("proxy", self._proxy_default):
            if request.is_bedrock():
                # boto3 doesn't allow extra_headers
                kwargs.pop("extra_headers", None)
            elif "extra_headers" not in kwargs and extra_headers:
                # assumes anthropic and openai clients
                kwargs["extra_headers"] = extra_headers

            self._logger.debug(f"invoke_wrapper: sending proxy request")

            return wrapped(*args, **kwargs)
        
        current_frame = inspect.currentframe()
        # f_back excludes the current frame, strip() cleans up whitespace and newlines
        stack = [frame.strip() for frame in traceback.format_stack(current_frame.f_back)]  # type: ignore

        request._ingest['properties'] = { 'system.stack_trace': json.dumps(stack) }

        if request.process_request(instance, extra_headers, args, kwargs) is False:
            self._logger.debug(f"invoke_wrapper: calling wrapped instance")
            return wrapped(*args, **kwargs)

        sw = Stopwatch()
        stream: bool = False
        
        if is_streaming == _IsStreaming.kwargs:
            stream = kwargs.get("stream", False)
        elif is_streaming == _IsStreaming.true:
            stream = True
        else:
            stream = False

        try:
            self._prepare_ingest(request, extra_headers, args, kwargs)
            self._logger.debug(f"invoke_wrapper: calling wrapped instance (stream={stream})")

            sw.start()
            response = wrapped(*args, **kwargs)
            
        except Exception as e:  # pylint: disable=broad-except
            sw.stop()
            duration = sw.elapsed_ms_int()

            self._logger.debug(f"invoke_wrapper: calling wrapped instance exception {e}")

            if request.process_exception(e, kwargs):
                request._ingest["end_to_end_latency_ms"] = duration
                self._ingest_units(request._ingest)

            raise e

        if stream:
            if request.streaming_type == _StreamingType.generator:
                return _GeneratorWrapper(
                    generator=response,
                    instance=instance,
                    instrumentor=self,
                    stopwatch=sw,
                    request=request,
                )
            elif request.streaming_type == _StreamingType.stream_manager:
                return _StreamManagerWrapper(
                    stream_manager=response,
                    instance=instance,
                    instrumentor=self,
                    stopwatch=sw,
                    request=request,
                )
            else:
                # request.streaming_type == _StreamingType.iterator
                stream_result = _StreamIteratorWrapper(
                    response=response,
                    instance=instance,
                    instrumentor=self,
                    stopwatch=sw,
                    request=request,
                )

                if request.is_bedrock():
                    if "body" in response:
                        response["body"] = stream_result
                    else:
                        response["stream"] = stream_result
                    return response
                
                return stream_result

        sw.stop()
        duration = sw.elapsed_ms_int()
        request._ingest["end_to_end_latency_ms"] = duration
        request._ingest["http_status_code"] = 200

        return_result: Any = request.process_synchronous_response(
            response=response,
            log_prompt_and_response=self._log_prompt_and_response,
            kwargs=kwargs)
        if return_result:
            self._logger.debug(f"invoke_wrapper: process sync response return")
            return return_result

        self._ingest_units(request._ingest)

        self._logger.debug(f"invoke_wrapper: finished")
        return response

    def _create_extra_headers(
        self
    ) -> 'dict[str, str]':
        extra_headers: dict[str, str] = {}
        context = self.get_context()
        if context:
            self._update_extra_headers(context, extra_headers)

        return extra_headers

    @staticmethod
    def _update_extra_headers(
        context: _Context,
        extra_headers: "dict[str, str]",
    ) -> None:
        context_limit_ids: Optional[list[str]] = context.get("limit_ids")

        context_experience_name: Optional[str] = context.get("experience_name")
        context_experience_id: Optional[str] = context.get("experience_id")

        context_use_case_name: Optional[str] = context.get("use_case_name")
        context_use_case_id: Optional[str] = context.get("use_case_id")
        context_use_case_version: Optional[int] = context.get("use_case_version")
        context_use_case_step: Optional[str] = context.get("use_case_step")

        context_user_id: Optional[str] = context.get("user_id")
        context_request_tags: Optional[list[str]] = context.get("request_tags")

        context_price_as_category: Optional[str] = context.get("price_as_category")
        context_price_as_resource: Optional[str] = context.get("price_as_resource")
        context_resource_scope: Optional[str] = context.get("resource_scope")

        # headers_limit_ids = extra_headers.get(PayiHeaderNames.limit_ids, None)
    
        # If the caller specifies limit_ids in extra_headers, it takes precedence over the decorator
        if PayiHeaderNames.limit_ids in extra_headers:
            headers_limit_ids = extra_headers.get(PayiHeaderNames.limit_ids)

            if headers_limit_ids is None or len(headers_limit_ids) == 0:
                # headers_limit_ids is empty, remove it from extra_headers
                extra_headers.pop(PayiHeaderNames.limit_ids, None)
            else:   
                # leave the value in extra_headers
                ...
        elif context_limit_ids:
            extra_headers[PayiHeaderNames.limit_ids] = ",".join(context_limit_ids)

        if PayiHeaderNames.user_id in extra_headers:
            headers_user_id = extra_headers.get(PayiHeaderNames.user_id, None)
            if headers_user_id is None or len(headers_user_id) == 0:
                # headers_user_id is empty, remove it from extra_headers
                extra_headers.pop(PayiHeaderNames.user_id, None)
            else:
                # leave the value in extra_headers
                ...
        elif context_user_id:
            extra_headers[PayiHeaderNames.user_id] = context_user_id

        if PayiHeaderNames.use_case_name in extra_headers:
            headers_use_case_name = extra_headers.get(PayiHeaderNames.use_case_name, None)
            if headers_use_case_name is None or len(headers_use_case_name) == 0:
                # headers_use_case_name is empty, remove all use case related headers
                extra_headers.pop(PayiHeaderNames.use_case_name, None)
                extra_headers.pop(PayiHeaderNames.use_case_id, None)
                extra_headers.pop(PayiHeaderNames.use_case_version, None)
                extra_headers.pop(PayiHeaderNames.use_case_step, None)
            else:
                # leave the value in extra_headers
                ...
        elif context_use_case_name:
            extra_headers[PayiHeaderNames.use_case_name] = context_use_case_name
            if context_use_case_id is not None:
                extra_headers[PayiHeaderNames.use_case_id] = context_use_case_id
            if context_use_case_version is not None:
                extra_headers[PayiHeaderNames.use_case_version] = str(context_use_case_version)
            if context_use_case_step is not None:
                extra_headers[PayiHeaderNames.use_case_step] = str(context_use_case_step)

        if PayiHeaderNames.experience_name in extra_headers:
            headers_experience_name = extra_headers.get(PayiHeaderNames.experience_name, None)
            if headers_experience_name is None or len(headers_experience_name) == 0:
                # headers_experience_name is empty, remove all experience related headers
                extra_headers.pop(PayiHeaderNames.experience_name, None)
                extra_headers.pop(PayiHeaderNames.experience_id, None)
            else:
                # leave the value in extra_headers
                ...
        elif context_experience_name is not None:
            extra_headers[PayiHeaderNames.experience_name] = context_experience_name
            if context_experience_id is not None:
                extra_headers[PayiHeaderNames.experience_id] = context_experience_id

        if PayiHeaderNames.request_tags not in extra_headers and context_request_tags:
            extra_headers[PayiHeaderNames.request_tags] = ",".join(context_request_tags)

        if PayiHeaderNames.price_as_category not in extra_headers and context_price_as_category:
            extra_headers[PayiHeaderNames.price_as_category] = context_price_as_category

        if PayiHeaderNames.price_as_resource not in extra_headers and context_price_as_resource:
            extra_headers[PayiHeaderNames.price_as_resource] = context_price_as_resource

        if PayiHeaderNames.resource_scope not in extra_headers and context_resource_scope:
            extra_headers[PayiHeaderNames.resource_scope] = context_resource_scope

    @staticmethod
    def update_for_vision(input: int, units: 'dict[str, Units]', estimated_prompt_tokens: Optional[int]) -> int:
        if estimated_prompt_tokens:
            vision = input - estimated_prompt_tokens
            if (vision > 0):
                units["vision"] = Units(input=vision, output=0)
                input = estimated_prompt_tokens
        
        return input

    @staticmethod
    def payi_wrapper(func: Any) -> Any:
        def _payi_wrapper(o: Any) -> Any:
            def wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any:
                return func(
                    o,
                    wrapped,
                    instance,
                    *args,
                    **kwargs,
                )

            return wrapper

        return _payi_wrapper

    @staticmethod
    def payi_awrapper(func: Any) -> Any:
        def _payi_awrapper(o: Any) -> Any:
            async def wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any:
                return await func(
                    o,
                    wrapped,
                    instance,
                    *args,
                    **kwargs,
                )

            return wrapper

        return _payi_awrapper

class _StreamIteratorWrapper(ObjectProxy):  # type: ignore
    def __init__(
        self,
        response: Any,
        instance: Any,
        instrumentor: _PayiInstrumentor,
        stopwatch: Stopwatch,
        request: _ProviderRequest,
    ) -> None:

        instrumentor._logger.debug(f"StreamIteratorWrapper: instance {instance}, category {request._category}")

        bedrock_from_stream: bool = False
        if request.is_bedrock():
            request._ingest["provider_response_id"] = response["ResponseMetadata"]["RequestId"]
            stream = response.get("stream", None)

            if stream:
                response = stream
                bedrock_from_stream = True
            else:
                response = response.get("body")
                bedrock_from_stream = False

        super().__init__(response)  # type: ignore

        self._response = response
        self._instance = instance

        self._instrumentor = instrumentor
        self._stopwatch: Stopwatch = stopwatch
        self._responses: list[str] = []

        self._request: _ProviderRequest = request

        self._first_token: bool = True
        self._is_bedrock: bool = request.is_bedrock()
        self._bedrock_from_stream: bool = bedrock_from_stream

    def __enter__(self) -> Any:
        self._instrumentor._logger.debug(f"StreamIteratorWrapper: __enter__")
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: 
        self._instrumentor._logger.debug(f"StreamIteratorWrapper: __exit__")
        self.__wrapped__.__exit__(exc_type, exc_val, exc_tb)  # type: ignore

    async def __aenter__(self) -> Any:
        self._instrumentor._logger.debug(f"StreamIteratorWrapper: __aenter__")
        return self

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        self._instrumentor._logger.debug(f"StreamIteratorWrapper: __aexit__")
        await self.__wrapped__.__aexit__(exc_type, exc_val, exc_tb)  # type: ignore

    def __iter__(self) -> Any:  
        if self._is_bedrock:
            # MUST reside in a separate function so that the yield statement (e.g. the generator) doesn't implicitly return its own iterator and overriding self
            self._instrumentor._logger.debug(f"StreamIteratorWrapper: bedrock __iter__")
            return self._iter_bedrock()

        self._instrumentor._logger.debug(f"StreamIteratorWrapper: __iter__")
        return self

    def _iter_bedrock(self) -> Any:
        # botocore EventStream doesn't have a __next__ method so iterate over the wrapped object in place
        for event in self.__wrapped__: # type: ignore
            if (self._bedrock_from_stream):
                self._evaluate_chunk(event)
            else:
                chunk = event.get('chunk') # type: ignore
                if chunk:
                    decode = chunk.get('bytes').decode() # type: ignore
                    self._evaluate_chunk(decode)
            yield event

        self._instrumentor._logger.debug(f"StreamIteratorWrapper: bedrock iter finished")

        self._stop_iteration()

    def __aiter__(self) -> Any:
        self._instrumentor._logger.debug(f"StreamIteratorWrapper: __aiter__")
        return self

    def __next__(self) -> object:
        try:
            chunk: object = self.__wrapped__.__next__()  # type: ignore
        except Exception as e:
            if isinstance(e, StopIteration):
                self._stop_iteration()
            else:
                self._instrumentor._logger.debug(f"StreamIteratorWrapper: __next__ exception {e}")
            raise e
        else:
            if self._evaluate_chunk(chunk) == False:
                return self.__next__()
            
            return chunk # type: ignore

    async def __anext__(self) -> object:
        try:
            chunk: object = await self.__wrapped__.__anext__()  # type: ignore
        except Exception as e:
            if isinstance(e, StopAsyncIteration):
                await self._astop_iteration()
            else:
                self._instrumentor._logger.debug(f"StreamIteratorWrapper: __anext__ exception {e}")
            raise e
        else:
            if self._evaluate_chunk(chunk) == False:
                return await self.__anext__()

            return chunk # type: ignore

    def _evaluate_chunk(self, chunk: Any) -> bool:
        if self._first_token:
            self._request._ingest["time_to_first_token_ms"] = self._stopwatch.elapsed_ms_int()
            self._first_token = False

        if self._instrumentor._log_prompt_and_response:
            self._responses.append(self.chunk_to_json(chunk))

        return self._request.process_chunk(chunk)

    def _process_stop_iteration(self) -> None:
        self._instrumentor._logger.debug(f"StreamIteratorWrapper: stop iteration")

        self._stopwatch.stop()
        self._request._ingest["end_to_end_latency_ms"] = self._stopwatch.elapsed_ms_int()
        self._request._ingest["http_status_code"] = 200

        if self._instrumentor._log_prompt_and_response:
            self._request._ingest["provider_response_json"] = self._responses

    async def _astop_iteration(self) -> None:
        self._process_stop_iteration()
        await self._instrumentor._aingest_units(self._request._ingest)

    def _stop_iteration(self) -> None:
        self._process_stop_iteration()
        self._instrumentor._ingest_units(self._request._ingest)

    @staticmethod
    def chunk_to_json(chunk: Any) -> str:
        if hasattr(chunk, "to_json"):
            return str(chunk.to_json())
        elif isinstance(chunk, bytes):
            return chunk.decode()
        elif isinstance(chunk, str):
            return chunk
        else:
            # assume dict
            return json.dumps(chunk)

class _StreamManagerWrapper(ObjectProxy):  # type: ignore
    def __init__(
        self,
        stream_manager: Any,  # type: ignore
        instance: Any,
        instrumentor: _PayiInstrumentor, 
        stopwatch: Stopwatch,
        request: _ProviderRequest,
    ) -> None:
        instrumentor._logger.debug(f"StreamManagerWrapper: instance {instance}, category {request._category}")

        super().__init__(stream_manager)  # type: ignore

        self._stream_manager = stream_manager  
        self._instance = instance
        self._instrumentor = instrumentor
        self._stopwatch: Stopwatch = stopwatch
        self._responses: list[str] = []
        self._request: _ProviderRequest = request
        self._first_token: bool = True
        self._done: bool = False

    def __enter__(self) -> _StreamIteratorWrapper:
        self._instrumentor._logger.debug(f"_StreamManagerWrapper: __enter__")

        return _StreamIteratorWrapper(
            response=self.__wrapped__.__enter__(),  # type: ignore
            instance=self._instance,
            instrumentor=self._instrumentor,
            stopwatch=self._stopwatch,
            request=self._request,
        )

class _GeneratorWrapper:  # type: ignore
    def __init__(
        self,
        generator: Any,
        instance: Any,
        instrumentor: _PayiInstrumentor, 
        stopwatch: Stopwatch,
        request: _ProviderRequest,
    ) -> None:
        instrumentor._logger.debug(f"GeneratorWrapper: instance {instance}, category {request._category}")

        super().__init__()  # type: ignore
        
        self._generator = generator
        self._instance = instance
        self._instrumentor = instrumentor
        self._stopwatch: Stopwatch = stopwatch
        self._log_prompt_and_response: bool = instrumentor._log_prompt_and_response
        self._responses: list[str] = []
        self._request: _ProviderRequest = request
        self._first_token: bool = True
        self._done: bool = False
        
    def __iter__(self) -> Any:
        self._instrumentor._logger.debug(f"GeneratorWrapper: __iter__")
        return self
        
    def __aiter__(self) -> Any:
        self._instrumentor._logger.debug(f"GeneratorWrapper: __aiter__")
        return self
        
    def __next__(self) -> Any:
        if self._done:
            raise StopIteration
            
        try:
            chunk = next(self._generator)
            return self._process_chunk(chunk)

        except Exception as e:
            if isinstance(e, StopIteration):
                self._process_stop_iteration()
            else:
                self._instrumentor._logger.debug(f"GeneratorWrapper: __next__ exception {e}")            
            raise e

    async def __anext__(self) -> Any:
        if self._done:
            raise StopAsyncIteration
            
        try:
            chunk = await anext(self._generator) # type: ignore
            return self._process_chunk(chunk)

        except Exception as e:
            if isinstance(e, StopAsyncIteration):
                await self._process_async_stop_iteration()
            else:
                self._instrumentor._logger.debug(f"GeneratorWrapper: __anext__ exception {e}")
            raise e

    @staticmethod
    def _chunk_to_dict(chunk: Any) -> 'dict[str, object]':
        if hasattr(chunk, "to_json"):
            return chunk.to_json() # type: ignore
        elif hasattr(chunk, "to_json_dict"):  
            return chunk.to_json_dict() # type: ignore
        else:
            return {}

    def _process_chunk(self, chunk: Any) -> Any:
        if self._first_token:
            self._request._ingest["time_to_first_token_ms"] = self._stopwatch.elapsed_ms_int()
            self._first_token = False
            
        if self._log_prompt_and_response:
            dict = self._chunk_to_dict(chunk) 
            self._responses.append(json.dumps(dict))
                
        self._request.process_chunk(chunk)
        return chunk

    def _process_stop_iteration(self) -> None:
        self._instrumentor._logger.debug(f"GeneratorWrapper: stop iteration")

        self._stopwatch.stop()
        self._request._ingest["end_to_end_latency_ms"] = self._stopwatch.elapsed_ms_int()
        self._request._ingest["http_status_code"] = 200
            
        if self._log_prompt_and_response:
            self._request._ingest["provider_response_json"] = self._responses
                
        self._instrumentor._ingest_units(self._request._ingest)
        self._done = True

    async def _process_async_stop_iteration(self) -> None:
        self._instrumentor._logger.debug(f"GeneratorWrapper: async stop iteration")
        
        self._stopwatch.stop() 
        self._request._ingest["end_to_end_latency_ms"] = self._stopwatch.elapsed_ms_int()
        self._request._ingest["http_status_code"] = 200
            
        if self._log_prompt_and_response:
            self._request._ingest["provider_response_json"] = self._responses
                
        await self._instrumentor._aingest_units(self._request._ingest)
        self._done = True

global _instrumentor
_instrumentor: Optional[_PayiInstrumentor] = None

def payi_instrument(
    payi: Optional[Union[Payi, AsyncPayi, 'list[Union[Payi, AsyncPayi]]']] = None,
    instruments: Optional[Set[str]] = None,
    log_prompt_and_response: bool = True,
    prompt_and_response_logger: Optional[Callable[[str, "dict[str, str]"], None]] = None,
    config: Optional[PayiInstrumentConfig] = None,
    logger: Optional[logging.Logger] = None,
) -> None:
    global _instrumentor
    if (_instrumentor):
        return
    
    payi_param: Optional[Payi] = None
    apayi_param: Optional[AsyncPayi] = None

    if isinstance(payi, Payi):
        payi_param = payi
    elif isinstance(payi, AsyncPayi):
        apayi_param = payi
    elif isinstance(payi, list):
        for p in payi:
            if isinstance(p, Payi):
                payi_param = p
            elif isinstance(p, AsyncPayi): # type: ignore
                apayi_param = p
    frameinfo = inspect.stack()[1]
    caller_filename = os.path.basename(frameinfo.filename).replace(' ', '_').lower()
    if caller_filename.endswith('.py'):
        caller_filename = caller_filename[:-3]

    # allow for both payi and apayi to be None for the @proxy case
    _instrumentor = _PayiInstrumentor(
        payi=payi_param,
        apayi=apayi_param,
        instruments=instruments,
        log_prompt_and_response=log_prompt_and_response,
        logger=logger,
        prompt_and_response_logger=prompt_and_response_logger,
        global_config=config if config else PayiInstrumentConfig(),
        caller_filename=caller_filename
    )

def track(
    limit_ids: Optional["list[str]"] = None,
    use_case_name: Optional[str] = None,
    use_case_id: Optional[str] = None,
    use_case_version: Optional[int] = None,
    user_id: Optional[str] = None,
    proxy: Optional[bool] = None,
) -> Any:

    def _track(func: Any) -> Any:
        import asyncio
        if asyncio.iscoroutinefunction(func):
            async def awrapper(*args: Any, **kwargs: Any) -> Any:
                if not _instrumentor:
                    _g_logger.debug(f"track: no instrumentor!")
                    return await func(*args, **kwargs)

                _instrumentor._logger.debug(f"track: call async function (proxy={proxy}, limit_ids={limit_ids}, use_case_name={use_case_name}, use_case_id={use_case_id}, use_case_version={use_case_version}, user_id={user_id})")
                # Call the instrumentor's _call_func for async functions
                return await _instrumentor._acall_func(
                    func,
                    proxy,
                    limit_ids,
                    None, # experience_name,
                    None, #experience_id,
                    use_case_name,
                    use_case_id,
                    use_case_version,
                    user_id,
                    *args,
                    **kwargs,
                )
            return awrapper
        else:
            def wrapper(*args: Any, **kwargs: Any) -> Any:
                if not _instrumentor:
                    _g_logger.debug(f"track: no instrumentor!")
                    return func(*args, **kwargs)

                _instrumentor._logger.debug(f"track: call sync function (proxy={proxy}, limit_ids={limit_ids}, use_case_name={use_case_name}, use_case_id={use_case_id}, use_case_version={use_case_version}, user_id={user_id})")

                return _instrumentor._call_func(
                    func,
                    proxy,
                    limit_ids,
                    None, # experience_name,
                    None, # experience_id,
                    use_case_name,
                    use_case_id,
                    use_case_version,
                    user_id,
                    *args,
                    **kwargs,
                )
            return wrapper
    return _track

def track_context(
    limit_ids: Optional["list[str]"] = None,
    use_case_name: Optional[str] = None,
    use_case_id: Optional[str] = None,
    use_case_version: Optional[int] = None,
    use_case_step: Optional[str] = None,
    user_id: Optional[str] = None,
    request_tags: Optional["list[str]"] = None,
    price_as_category: Optional[str] = None,
    price_as_resource: Optional[str] = None,
    resource_scope: Optional[str] = None,
    proxy: Optional[bool] = None,
) -> _TrackContext:
    # Create a new context for tracking
    context: _Context = {}

    context["proxy"] = proxy

    context["limit_ids"] = limit_ids

    context["use_case_name"] = use_case_name
    context["use_case_id"] = use_case_id
    context["use_case_version"] = use_case_version
    context["use_case_step"] = use_case_step

    context["user_id"] = user_id
    context["request_tags"] = request_tags

    context["price_as_category"] = price_as_category
    context["price_as_resource"] = price_as_resource
    context["resource_scope"] = resource_scope

    return _TrackContext(context)

@deprecated("@ingest() is deprecated. Use @track() instead")
def ingest(
    limit_ids: Optional["list[str]"] = None,
    experience_name: Optional[str] = None,
    experience_id: Optional[str] = None,
    use_case_name: Optional[str] = None,
    use_case_id: Optional[str] = None,
    use_case_version: Optional[int] = None,
    user_id: Optional[str] = None,
) -> Any:
    warnings.warn(
        "@ingest is deprecated and will be removed in a future version. Use @track instead.",
        DeprecationWarning,
        stacklevel=2

    )

    def _ingest(func: Any) -> Any:
        import asyncio
        if asyncio.iscoroutinefunction(func):
            async def awrapper(*args: Any, **kwargs: Any) -> Any:
                if not _instrumentor:
                    _g_logger.debug(f"ingest: call no instrumentor!")
                    return await func(*args, **kwargs)

                _instrumentor._logger.debug(f"ingest: call async function (limit_ids={limit_ids}, experience_name={experience_name}, experience_id={experience_id}, use_case_name={use_case_name}, use_case_id={use_case_id}, use_case_version={use_case_version}, user_id={user_id})")

                # Call the instrumentor's _call_func for async functions
                return await _instrumentor._acall_func(
                    func,
                    False,
                    limit_ids,
                    experience_name,
                    experience_id,
                    use_case_name,
                    use_case_id,
                    use_case_version,
                    user_id,
                    *args,
                    **kwargs,
                )
            return awrapper
        else:
            def wrapper(*args: Any, **kwargs: Any) -> Any:
                if not _instrumentor:
                    _g_logger.debug(f"ingest: call no instrumentor!")
                    return func(*args, **kwargs)

                _instrumentor._logger.debug(f"ingest: call sync function (limit_ids={limit_ids}, experience_name={experience_name}, experience_id={experience_id}, use_case_name={use_case_name}, use_case_id={use_case_id}, use_case_version={use_case_version}, user_id={user_id})")

                return _instrumentor._call_func(
                    func,
                    False,
                    limit_ids,
                    experience_name,
                    experience_id,
                    use_case_name,
                    use_case_id,
                    use_case_version,
                    user_id,
                    *args,
                    **kwargs,
                )
            return wrapper

    return _ingest

@deprecated("@proxy() is deprecated. Use @track() instead")
def proxy(
    limit_ids: Optional["list[str]"] = None,
    experience_name: Optional[str] = None,
    experience_id: Optional[str] = None,
    use_case_id: Optional[str] = None,
    use_case_name: Optional[str] = None,
    use_case_version: Optional[int] = None,
    user_id: Optional[str] = None,
) -> Any:
    warnings.warn(
        "@proxy is deprecated and will be removed in a future version. Use @track instead.",
        DeprecationWarning,
        stacklevel=2
    )

    def _proxy(func: Any) -> Any:
        import asyncio
        if asyncio.iscoroutinefunction(func):
            async def _proxy_awrapper(*args: Any, **kwargs: Any) -> Any:
                if not _instrumentor:
                    _g_logger.debug(f"proxy: call no instrumentor!")
                    return await func(*args, **kwargs)

                _instrumentor._logger.debug(f"proxy: call async function (limit_ids={limit_ids}, experience_name={experience_name}, experience_id={experience_id}, use_case_name={use_case_name}, use_case_id={use_case_id}, use_case_version={use_case_version}, user_id={user_id})")

                return await _instrumentor._call_func(
                    func,
                    True,
                    limit_ids,
                    experience_name,
                    experience_id,
                    use_case_name,
                    use_case_id,
                    use_case_version,
                    user_id,
                    *args,
                    **kwargs
                )

            return _proxy_awrapper
        else:
            def _proxy_wrapper(*args: Any, **kwargs: Any) -> Any:
                if not _instrumentor:
                    _g_logger.debug(f"proxy: call no instrumentor!")
                    return func(*args, **kwargs)

                _instrumentor._logger.debug(f"proxy: call sync function (limit_ids={limit_ids}, experience_name={experience_name}, experience_id={experience_id}, use_case_name={use_case_name}, use_case_id={use_case_id}, use_case_version={use_case_version}, user_id={user_id})")

                return _instrumentor._call_func(
                    func,
                    True,
                    limit_ids,
                    experience_name,
                    experience_id,
                    use_case_name,
                    use_case_id,
                    use_case_version,
                    user_id,
                    *args,
                    **kwargs
                )

            return _proxy_wrapper

    return _proxy
