"""OpenTelemetry Bedrock instrumentation"""

import json
import logging
import os
import time
import traceback
import warnings
from functools import partial, wraps
from typing import Collection, Optional

from opentelemetry import context as context_api
from opentelemetry._logs import get_logger
from opentelemetry.instrumentation.bedrock.config import Config
from opentelemetry.instrumentation.bedrock.event_emitter import (
    emit_choice_events,
    emit_input_events_converse,
    emit_message_events,
    emit_response_event_converse,
    emit_streaming_converse_response_event,
    emit_streaming_response_event,
)
from opentelemetry.instrumentation.bedrock.guardrail import (
    guardrail_converse,
    guardrail_handling,
)
from opentelemetry.instrumentation.bedrock.prompt_caching import prompt_caching_handling
from opentelemetry.instrumentation.bedrock.reusable_streaming_body import (
    BufferedAsyncBody,
    ReusableStreamingBody,
)
from opentelemetry.instrumentation.bedrock.span_utils import (
    converse_usage_record,
    set_converse_input_prompt_span_attributes,
    _set_converse_finish_reasons,
    _set_finish_reasons_unconditionally,
    set_converse_model_span_attributes,
    set_converse_response_span_attributes,
    set_converse_streaming_response_span_attributes,
    set_model_choice_span_attributes,
    set_model_message_span_attributes,
    set_model_span_attributes,
)
from opentelemetry.instrumentation.bedrock.streaming_wrapper import (
    AsyncStreamingWrapper,
    StreamingWrapper,
)
from opentelemetry.instrumentation.bedrock.utils import (
    dont_throw,
    should_emit_events,
)
from opentelemetry.instrumentation.bedrock.version import __version__
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
from opentelemetry.instrumentation.utils import (
    _SUPPRESS_INSTRUMENTATION_KEY,
    unwrap,
)
from opentelemetry.metrics import Counter, Histogram, Meter, get_meter
from opentelemetry.semconv._incubating.attributes import (
    gen_ai_attributes as GenAIAttributes,
)
from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import (
    GenAiOperationNameValues,
    GenAiSystemValues,
)
from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
from opentelemetry.semconv_ai import (
    SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY,
    Meters,
)
from opentelemetry.trace import Span, SpanKind, get_tracer
from opentelemetry.trace.status import Status, StatusCode
from wrapt import wrap_function_wrapper


class MetricParams:
    def __init__(
        self,
        token_histogram: Histogram,
        choice_counter: Counter,
        duration_histogram: Histogram,
        exception_counter: Counter,
        guardrail_activation: Counter,
        guardrail_latency_histogram: Histogram,
        guardrail_coverage: Counter,
        guardrail_sensitive_info: Counter,
        guardrail_topic: Counter,
        guardrail_content: Counter,
        guardrail_words: Counter,
        prompt_caching: Counter,
    ):
        self.vendor = ""
        self.model = ""
        self.is_stream = False
        self.token_histogram = token_histogram
        self.choice_counter = choice_counter
        self.duration_histogram = duration_histogram
        self.exception_counter = exception_counter
        self.guardrail_activation = guardrail_activation
        self.guardrail_latency_histogram = guardrail_latency_histogram
        self.guardrail_coverage = guardrail_coverage
        self.guardrail_sensitive_info = guardrail_sensitive_info
        self.guardrail_topic = guardrail_topic
        self.guardrail_content = guardrail_content
        self.guardrail_words = guardrail_words
        self.prompt_caching = prompt_caching
        self.start_time = time.time()


logger = logging.getLogger(__name__)


_instruments = ("boto3 >= 1.28.57",)

WRAPPED_METHODS = [
    {
        "package": "botocore.client",
        "object": "ClientCreator",
        "method": "create_client",
    },
    {"package": "botocore.session", "object": "Session", "method": "create_client"},
    {
        "package": "aiobotocore.session",
        "object": "AioSession",
        "method": "create_client",
        "async": True,
    },
]

def _span_name(operation_name, model):
    """Build span name per OTel semconv: '{operation_name} {model}'."""
    return f"{operation_name} {model}" if model else operation_name


def is_metrics_enabled() -> bool:
    return (os.getenv("TRACELOOP_METRICS_ENABLED") or "true").lower() == "true"


def _with_tracer_wrapper(func):
    """Helper for providing tracer for wrapper functions."""

    def _with_tracer(
        tracer,
        metric_params,
        event_logger,
        to_wrap,
    ):
        def wrapper(wrapped, instance, args, kwargs):
            return func(
                tracer,
                metric_params,
                event_logger,
                to_wrap,
                wrapped,
                instance,
                args,
                kwargs,
            )

        return wrapper

    return _with_tracer


@_with_tracer_wrapper
def _wrap(
    tracer,
    metric_params,
    event_logger,
    to_wrap,
    wrapped,
    instance,
    args,
    kwargs,
):
    """Instruments and calls every function defined in TO_WRAP."""
    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
        return wrapped(*args, **kwargs)

    if kwargs.get("service_name") == "bedrock-runtime":
        try:
            start_time = time.time()
            metric_params.start_time = time.time()
            client = wrapped(*args, **kwargs)
            client.invoke_model = _instrumented_model_invoke(
                client.invoke_model, tracer, metric_params, event_logger
            )
            client.invoke_model_with_response_stream = (
                _instrumented_model_invoke_with_response_stream(
                    client.invoke_model_with_response_stream,
                    tracer,
                    metric_params,
                    event_logger,
                )
            )
            client.converse = _instrumented_converse(
                client.converse, tracer, metric_params, event_logger
            )
            client.converse_stream = _instrumented_converse_stream(
                client.converse_stream, tracer, metric_params, event_logger
            )
            return client
        except Exception as e:
            end_time = time.time()
            duration = end_time - start_time if "start_time" in locals() else 0

            attributes = {
                "error.type": e.__class__.__name__,
            }

            if duration > 0 and metric_params.duration_histogram:
                metric_params.duration_histogram.record(duration, attributes=attributes)
            if metric_params.exception_counter:
                metric_params.exception_counter.add(1, attributes=attributes)

            raise

    return wrapped(*args, **kwargs)


def _instrumented_model_invoke(fn, tracer, metric_params, event_logger):
    @wraps(fn)
    def with_instrumentation(*args, **kwargs):
        if context_api.get_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY):
            return fn(*args, **kwargs)

        _buffer_request_body(kwargs)
        (provider, _model_vendor, _model) = _get_vendor_model(kwargs.get("modelId"))
        operation_name = _derive_operation_name(kwargs)
        span_attributes = {
            GenAIAttributes.GEN_AI_PROVIDER_NAME: provider,
            GenAIAttributes.GEN_AI_OPERATION_NAME: operation_name,
            GenAIAttributes.GEN_AI_REQUEST_MODEL: _model,
        }
        with tracer.start_as_current_span(
            _span_name(operation_name, _model), kind=SpanKind.CLIENT, attributes=span_attributes,
            record_exception=False, set_status_on_exception=False,
        ) as span:
            try:
                response = fn(*args, **kwargs)
            except Exception as e:
                span.set_attribute(ERROR_TYPE, e.__class__.__name__)
                span.record_exception(e)
                span.set_status(Status(StatusCode.ERROR, str(e)))
                raise
            _handle_call(span, kwargs, response, metric_params, event_logger)
            return response

    return with_instrumentation


def _instrumented_model_invoke_with_response_stream(
    fn, tracer, metric_params, event_logger
):
    @wraps(fn)
    def with_instrumentation(*args, **kwargs):
        if context_api.get_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY):
            return fn(*args, **kwargs)

        _buffer_request_body(kwargs)
        (provider, _model_vendor, _model) = _get_vendor_model(kwargs.get("modelId"))
        operation_name = _derive_operation_name(kwargs)
        span_attributes = {
            GenAIAttributes.GEN_AI_PROVIDER_NAME: provider,
            GenAIAttributes.GEN_AI_OPERATION_NAME: operation_name,
            GenAIAttributes.GEN_AI_REQUEST_MODEL: _model,
        }
        span = tracer.start_span(
            _span_name(operation_name, _model),
            kind=SpanKind.CLIENT,
            attributes=span_attributes,
        )

        try:
            response = fn(*args, **kwargs)
        except Exception as e:
            span.set_attribute(ERROR_TYPE, e.__class__.__name__)
            span.record_exception(e)
            span.set_status(Status(StatusCode.ERROR, str(e)))
            span.end()
            raise
        _handle_stream_call(span, kwargs, response, metric_params, event_logger)

        return response

    return with_instrumentation


def _instrumented_converse(fn, tracer, metric_params, event_logger):
    # see
    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime/client/converse.html
    # for the request/response format
    @wraps(fn)
    def with_instrumentation(*args, **kwargs):
        if context_api.get_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY):
            return fn(*args, **kwargs)

        (provider, _model_vendor, _model) = _get_vendor_model(kwargs.get("modelId"))
        span_attributes = {
            GenAIAttributes.GEN_AI_PROVIDER_NAME: provider,
            GenAIAttributes.GEN_AI_OPERATION_NAME: GenAiOperationNameValues.CHAT.value,
            GenAIAttributes.GEN_AI_REQUEST_MODEL: _model,
        }
        with tracer.start_as_current_span(
            _span_name(GenAiOperationNameValues.CHAT.value, _model),
            kind=SpanKind.CLIENT,
            attributes=span_attributes,
            record_exception=False, set_status_on_exception=False,
        ) as span:
            try:
                response = fn(*args, **kwargs)
            except Exception as e:
                span.set_attribute(ERROR_TYPE, e.__class__.__name__)
                span.record_exception(e)
                span.set_status(Status(StatusCode.ERROR, str(e)))
                raise
            _handle_converse(span, kwargs, response, metric_params, event_logger)

            return response

    return with_instrumentation


def _instrumented_converse_stream(fn, tracer, metric_params, event_logger):
    @wraps(fn)
    def with_instrumentation(*args, **kwargs):
        if context_api.get_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY):
            return fn(*args, **kwargs)

        (provider, _model_vendor, _model) = _get_vendor_model(kwargs.get("modelId"))
        span_attributes = {
            GenAIAttributes.GEN_AI_PROVIDER_NAME: provider,
            GenAIAttributes.GEN_AI_OPERATION_NAME: GenAiOperationNameValues.CHAT.value,
            GenAIAttributes.GEN_AI_REQUEST_MODEL: _model,
        }
        span = tracer.start_span(
            _span_name(GenAiOperationNameValues.CHAT.value, _model),
            kind=SpanKind.CLIENT,
            attributes=span_attributes,
        )
        try:
            response = fn(*args, **kwargs)
        except Exception as e:
            span.set_attribute(ERROR_TYPE, e.__class__.__name__)
            span.record_exception(e)
            span.set_status(Status(StatusCode.ERROR, str(e)))
            span.end()
            raise
        if span.is_recording():
            _handle_converse_stream(span, kwargs, response, metric_params, event_logger)

        return response

    return with_instrumentation


def _instrumented_async_model_invoke(fn, tracer, metric_params, event_logger):
    @wraps(fn)
    async def with_instrumentation(*args, **kwargs):
        if context_api.get_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY):
            return await fn(*args, **kwargs)

        _buffer_request_body(kwargs)
        (provider, _model_vendor, _model) = _get_vendor_model(kwargs.get("modelId"))
        operation_name = _derive_operation_name(kwargs)
        span_attributes = {
            GenAIAttributes.GEN_AI_PROVIDER_NAME: provider,
            GenAIAttributes.GEN_AI_OPERATION_NAME: operation_name,
            GenAIAttributes.GEN_AI_REQUEST_MODEL: _model,
        }
        with tracer.start_as_current_span(
            _span_name(operation_name, _model), kind=SpanKind.CLIENT, attributes=span_attributes
        ) as span:
            response = await fn(*args, **kwargs)
            await _handle_async_call(span, kwargs, response, metric_params, event_logger)
            return response

    return with_instrumentation


def _instrumented_async_model_invoke_with_response_stream(
    fn, tracer, metric_params, event_logger
):
    @wraps(fn)
    async def with_instrumentation(*args, **kwargs):
        if context_api.get_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY):
            return await fn(*args, **kwargs)

        _buffer_request_body(kwargs)
        (provider, _model_vendor, _model) = _get_vendor_model(kwargs.get("modelId"))
        operation_name = _derive_operation_name(kwargs)
        span_attributes = {
            GenAIAttributes.GEN_AI_PROVIDER_NAME: provider,
            GenAIAttributes.GEN_AI_OPERATION_NAME: operation_name,
            GenAIAttributes.GEN_AI_REQUEST_MODEL: _model,
        }
        span = tracer.start_span(
            _span_name(operation_name, _model),
            kind=SpanKind.CLIENT,
            attributes=span_attributes,
        )

        response = await fn(*args, **kwargs)
        _handle_async_stream_call(span, kwargs, response, metric_params, event_logger)

        return response

    return with_instrumentation


def _instrumented_async_converse(fn, tracer, metric_params, event_logger):
    @wraps(fn)
    async def with_instrumentation(*args, **kwargs):
        if context_api.get_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY):
            return await fn(*args, **kwargs)

        (provider, _model_vendor, _model) = _get_vendor_model(kwargs.get("modelId"))
        span_attributes = {
            GenAIAttributes.GEN_AI_PROVIDER_NAME: provider,
            GenAIAttributes.GEN_AI_OPERATION_NAME: GenAiOperationNameValues.CHAT.value,
            GenAIAttributes.GEN_AI_REQUEST_MODEL: _model,
        }
        with tracer.start_as_current_span(
            _span_name(GenAiOperationNameValues.CHAT.value, _model),
            kind=SpanKind.CLIENT,
            attributes=span_attributes,
        ) as span:
            response = await fn(*args, **kwargs)
            _handle_converse(span, kwargs, response, metric_params, event_logger)

            return response

    return with_instrumentation


def _instrumented_async_converse_stream(fn, tracer, metric_params, event_logger):
    @wraps(fn)
    async def with_instrumentation(*args, **kwargs):
        if context_api.get_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY):
            return await fn(*args, **kwargs)

        (provider, _model_vendor, _model) = _get_vendor_model(kwargs.get("modelId"))
        span_attributes = {
            GenAIAttributes.GEN_AI_PROVIDER_NAME: provider,
            GenAIAttributes.GEN_AI_OPERATION_NAME: GenAiOperationNameValues.CHAT.value,
            GenAIAttributes.GEN_AI_REQUEST_MODEL: _model,
        }
        span = tracer.start_span(
            _span_name(GenAiOperationNameValues.CHAT.value, _model),
            kind=SpanKind.CLIENT,
            attributes=span_attributes,
        )
        response = await fn(*args, **kwargs)
        if span.is_recording():
            _handle_async_converse_stream(span, kwargs, response, metric_params, event_logger)

        return response

    return with_instrumentation


class _InstrumentedClientContext:
    """Wraps aiobotocore's ClientCreatorContext to monkey-patch Bedrock methods
    on the live async client after it's created."""

    def __init__(self, inner_ctx, tracer, metric_params, event_logger):
        self._inner_ctx = inner_ctx
        self._tracer = tracer
        self._metric_params = metric_params
        self._event_logger = event_logger

    async def __aenter__(self):
        client = await self._inner_ctx.__aenter__()
        try:
            client.invoke_model = _instrumented_async_model_invoke(
                client.invoke_model, self._tracer, self._metric_params, self._event_logger
            )
            client.invoke_model_with_response_stream = (
                _instrumented_async_model_invoke_with_response_stream(
                    client.invoke_model_with_response_stream,
                    self._tracer,
                    self._metric_params,
                    self._event_logger,
                )
            )
            client.converse = _instrumented_async_converse(
                client.converse, self._tracer, self._metric_params, self._event_logger
            )
            client.converse_stream = _instrumented_async_converse_stream(
                client.converse_stream, self._tracer, self._metric_params, self._event_logger
            )
        except Exception:
            # If monkey-patching fails (e.g., aiobotocore exposes a method as a
            # read-only property in a future version), release the underlying
            # HTTP connection rather than leaking it.
            try:
                await self._inner_ctx.__aexit__(None, None, None)
            except Exception:  # noqa: BLE001 - don't mask original error
                pass
            raise
        return client

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        return await self._inner_ctx.__aexit__(exc_type, exc_val, exc_tb)


def _wrap_async_factory(tracer, metric_params, event_logger, to_wrap):
    """Wrapper for aiobotocore's AioSession.create_client (sync, returns async ctx mgr)."""

    def wrapper(wrapped, instance, args, kwargs):
        if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
            return wrapped(*args, **kwargs)

        ctx = wrapped(*args, **kwargs)

        service_name = kwargs.get("service_name") or (args[0] if args else None)
        if service_name != "bedrock-runtime":
            return ctx

        return _InstrumentedClientContext(ctx, tracer, metric_params, event_logger)

    return wrapper


@dont_throw
def _handle_stream_call(span, kwargs, response, metric_params, event_logger):

    (provider, model_vendor, model) = _get_vendor_model(kwargs.get("modelId"))
    request_body = json.loads(kwargs.get("body"))

    headers = {}
    if "ResponseMetadata" in response:
        headers = response.get("ResponseMetadata").get("HTTPHeaders", {})

    @dont_throw
    def stream_done(response_body):

        metric_params.vendor = provider
        metric_params.model = model
        metric_params.is_stream = True

        prompt_caching_handling(headers, provider, model, metric_params)
        guardrail_handling(span, response_body, provider, model, metric_params)

        if span.is_recording():
            set_model_span_attributes(
                provider,
                model_vendor,
                model,
                span,
                request_body,
                response_body,
                headers,
                metric_params,
                kwargs,
            )
        if should_emit_events() and event_logger:
            emit_message_events(event_logger, kwargs)
            emit_streaming_response_event(response_body, event_logger)
            _set_finish_reasons_unconditionally(model_vendor, span, response_body)
        else:
            set_model_message_span_attributes(model_vendor, span, request_body)
            set_model_choice_span_attributes(model_vendor, span, response_body)

        span.end()

    response["body"] = StreamingWrapper(
        response["body"], stream_done_callback=stream_done
    )


@dont_throw
def _handle_async_stream_call(span, kwargs, response, metric_params, event_logger):
    """Async counterpart of _handle_stream_call — wraps the response body
    with AsyncStreamingWrapper so the user can `async for` over it."""
    (provider, model_vendor, model) = _get_vendor_model(kwargs.get("modelId"))
    request_body = json.loads(kwargs.get("body"))

    headers = {}
    if "ResponseMetadata" in response:
        headers = response.get("ResponseMetadata").get("HTTPHeaders", {})

    @dont_throw
    def stream_done(response_body):
        metric_params.vendor = provider
        metric_params.model = model
        metric_params.is_stream = True

        prompt_caching_handling(headers, provider, model, metric_params)
        guardrail_handling(span, response_body, provider, model, metric_params)

        if span.is_recording():
            set_model_span_attributes(
                provider,
                model_vendor,
                model,
                span,
                request_body,
                response_body,
                headers,
                metric_params,
                kwargs,
            )
        if should_emit_events() and event_logger:
            emit_message_events(event_logger, kwargs)
            emit_streaming_response_event(response_body, event_logger)
            _set_finish_reasons_unconditionally(model_vendor, span, response_body)
        else:
            set_model_message_span_attributes(model_vendor, span, request_body)
            set_model_choice_span_attributes(model_vendor, span, response_body)

        span.end()

    response["body"] = AsyncStreamingWrapper(
        response["body"], stream_done_callback=stream_done
    )


@dont_throw
def _handle_call(span: Span, kwargs, response, metric_params, event_logger):
    response["body"] = ReusableStreamingBody(
        response["body"]._raw_stream, response["body"]._content_length
    )
    request_body = json.loads(kwargs.get("body"))
    response_body = json.loads(response.get("body").read())
    headers = {}
    if "ResponseMetadata" in response:
        headers = response.get("ResponseMetadata").get("HTTPHeaders", {})

    (provider, model_vendor, model) = _get_vendor_model(kwargs.get("modelId"))
    metric_params.vendor = provider
    metric_params.model = model
    metric_params.is_stream = False

    prompt_caching_handling(headers, provider, model, metric_params)
    guardrail_handling(span, response_body, provider, model, metric_params)

    if span.is_recording():
        set_model_span_attributes(
            provider,
            model_vendor,
            model,
            span,
            request_body,
            response_body,
            headers,
            metric_params,
            kwargs,
        )

    if should_emit_events() and event_logger:
        emit_message_events(event_logger, kwargs)
        emit_choice_events(event_logger, response)
        _set_finish_reasons_unconditionally(model_vendor, span, response_body)
    else:
        set_model_message_span_attributes(model_vendor, span, request_body)
        set_model_choice_span_attributes(model_vendor, span, response_body)


async def _handle_async_call(span: Span, kwargs, response, metric_params, event_logger):
    """Async counterpart of _handle_call — reads the response body via `await
    body.read()` (aiobotocore's body is async) and replaces it with a buffered
    wrapper so the user can read it again."""
    original_body = response.get("body")
    try:
        raw_bytes = await original_body.read()
        # Close the underlying aiobotocore body to release the HTTP connection
        # back to aiohttp's pool. Without this, the connection stays held until
        # the response object is GC'd.
        close_fn = getattr(original_body, "close", None)
        if close_fn is not None:
            try:
                result = close_fn()
                if hasattr(result, "__await__"):
                    await result
            except Exception:  # noqa: BLE001 - never fail instrumentation on close
                pass

        response["body"] = BufferedAsyncBody(raw_bytes)
        request_body = json.loads(kwargs.get("body"))
        response_body = json.loads(raw_bytes)
        headers = {}
        if "ResponseMetadata" in response:
            headers = response.get("ResponseMetadata").get("HTTPHeaders", {})

        (provider, model_vendor, model) = _get_vendor_model(kwargs.get("modelId"))
        metric_params.vendor = provider
        metric_params.model = model
        metric_params.is_stream = False

        prompt_caching_handling(headers, provider, model, metric_params)
        guardrail_handling(span, response_body, provider, model, metric_params)

        if span.is_recording():
            set_model_span_attributes(
                provider,
                model_vendor,
                model,
                span,
                request_body,
                response_body,
                headers,
                metric_params,
                kwargs,
            )

        if should_emit_events() and event_logger:
            emit_message_events(event_logger, kwargs)
            # emit_choice_events does a sync `response["body"].read()` then
            # accesses other top-level response keys. Pass the real response
            # with body swapped for a sync-readable view so both work.
            response_for_emit = dict(response)
            response_for_emit["body"] = _SyncBodyView(raw_bytes)
            emit_choice_events(event_logger, response_for_emit)
            _set_finish_reasons_unconditionally(model_vendor, span, response_body)
        else:
            set_model_message_span_attributes(model_vendor, span, request_body)
            set_model_choice_span_attributes(model_vendor, span, response_body)
    except Exception as e:
        # Mirror @dont_throw behavior for async: log + route through Config.exception_logger
        logger.debug(
            "OpenLLMetry failed to trace in _handle_async_call, error: %s",
            traceback.format_exc(),
        )
        if Config.exception_logger:
            Config.exception_logger(e)


class _SyncBodyView:
    """Tiny sync-readable view of already-buffered bytes, for helpers that
    expect a botocore-style body with a sync .read()."""

    def __init__(self, raw_bytes: bytes):
        self._raw = raw_bytes

    def read(self):
        return self._raw


@dont_throw
def _handle_converse(span, kwargs, response, metric_params, event_logger):
    (provider, model_vendor, model) = _get_vendor_model(kwargs.get("modelId"))
    guardrail_converse(span, response, provider, model, metric_params)

    set_converse_model_span_attributes(span, provider, model, kwargs)

    converse_usage_record(span, response, metric_params)

    if should_emit_events() and event_logger:
        emit_input_events_converse(kwargs, event_logger)
        emit_response_event_converse(response, event_logger)
        _set_converse_finish_reasons(span, response.get("stopReason"))
    else:
        set_converse_input_prompt_span_attributes(kwargs, span)
        set_converse_response_span_attributes(response, span)


@dont_throw
def _handle_converse_stream(span, kwargs, response, metric_params, event_logger):
    (provider, model_vendor, model) = _get_vendor_model(kwargs.get("modelId"))

    set_converse_model_span_attributes(span, provider, model, kwargs)

    if should_emit_events() and event_logger:
        emit_input_events_converse(kwargs, event_logger)
    else:
        set_converse_input_prompt_span_attributes(kwargs, span)

    stream = response.get("stream")
    role = "unknown"
    if stream:

        def handler(func):
            def wrap(*args, **kwargs):
                response_msg = kwargs.pop("response_msg")
                tool_blocks = kwargs.pop("tool_blocks")
                reasoning_blocks = kwargs.pop("reasoning_blocks")
                span = kwargs.pop("span")
                event = func(*args, **kwargs)
                nonlocal role
                if "contentBlockDelta" in event:
                    delta = event["contentBlockDelta"].get("delta", {})
                    if "text" in delta:
                        response_msg.append(delta["text"])
                    if "toolUse" in delta:
                        # Merge delta input into the last tool block (created by contentBlockStart)
                        if tool_blocks:
                            tool_blocks[-1].setdefault("input", "")
                            tool_blocks[-1]["input"] += delta["toolUse"].get("input", "")
                        else:
                            tool_blocks.append(delta["toolUse"])
                    if "reasoningContent" in delta:
                        reasoning_blocks.append(delta["reasoningContent"].get("text", ""))
                elif "contentBlockStart" in event:
                    start = event["contentBlockStart"].get("start", {})
                    if "toolUse" in start:
                        tool_blocks.append(start["toolUse"])
                elif "messageStart" in event:
                    role = event["messageStart"]["role"]
                elif "metadata" in event:
                    # last message sent
                    guardrail_converse(span, event["metadata"], provider, model, metric_params)
                    converse_usage_record(span, event["metadata"], metric_params)
                    span.end()
                elif "messageStop" in event:
                    stop_reason = event.get("messageStop", {}).get("stopReason")
                    if should_emit_events() and event_logger:
                        emit_streaming_converse_response_event(
                            event_logger,
                            response_msg,
                            role,
                            stop_reason,
                        )
                        _set_converse_finish_reasons(span, stop_reason)
                    else:
                        set_converse_streaming_response_span_attributes(
                            response_msg,
                            role,
                            span,
                            finish_reason=stop_reason,
                            tool_blocks=tool_blocks,
                            reasoning_blocks=reasoning_blocks,
                        )

                return event

            return partial(wrap, response_msg=[], tool_blocks=[], reasoning_blocks=[], span=span)

        stream._parse_event = handler(stream._parse_event)


@dont_throw
def _handle_async_converse_stream(span, kwargs, response, metric_params, event_logger):
    """Async variant of _handle_converse_stream — `_parse_event` is a coroutine
    in aiobotocore, so the wrapper must await it before inspecting the event."""
    (provider, model_vendor, model) = _get_vendor_model(kwargs.get("modelId"))

    set_converse_model_span_attributes(span, provider, model, kwargs)

    if should_emit_events() and event_logger:
        emit_input_events_converse(kwargs, event_logger)
    else:
        set_converse_input_prompt_span_attributes(kwargs, span)

    stream = response.get("stream")
    role = "unknown"
    if stream:
        # Track whether span.end() has been called inside the metadata branch.
        # If the caller breaks out of `async for` early, metadata never arrives
        # and the span would otherwise leak. We close it from the iterator
        # wrapper's finally block as a fallback.
        #
        # We also keep references to the accumulators (response_msg etc.) on
        # span_state so the iterator wrapper can flush partial content into
        # the span as attributes when iteration ends without a messageStop.
        span_state = {"ended": False, "saw_message_stop": False}
        partial_state = {
            "response_msg": [],
            "tool_blocks": [],
            "reasoning_blocks": [],
        }

        def handler(func):
            async def wrap(*args, **kwargs):
                response_msg = kwargs.pop("response_msg")
                tool_blocks = kwargs.pop("tool_blocks")
                reasoning_blocks = kwargs.pop("reasoning_blocks")
                span = kwargs.pop("span")
                event = await func(*args, **kwargs)
                nonlocal role
                if "contentBlockDelta" in event:
                    delta = event["contentBlockDelta"].get("delta", {})
                    if "text" in delta:
                        response_msg.append(delta["text"])
                    if "toolUse" in delta:
                        if tool_blocks:
                            tool_blocks[-1].setdefault("input", "")
                            tool_blocks[-1]["input"] += delta["toolUse"].get("input", "")
                        else:
                            tool_blocks.append(delta["toolUse"])
                    if "reasoningContent" in delta:
                        reasoning_blocks.append(delta["reasoningContent"].get("text", ""))
                elif "contentBlockStart" in event:
                    start = event["contentBlockStart"].get("start", {})
                    if "toolUse" in start:
                        tool_blocks.append(start["toolUse"])
                elif "messageStart" in event:
                    role = event["messageStart"]["role"]
                elif "metadata" in event:
                    guardrail_converse(span, event["metadata"], provider, model, metric_params)
                    converse_usage_record(span, event["metadata"], metric_params)
                    span_state["ended"] = True
                    span.end()
                elif "messageStop" in event:
                    span_state["saw_message_stop"] = True
                    stop_reason = event.get("messageStop", {}).get("stopReason")
                    if should_emit_events() and event_logger:
                        emit_streaming_converse_response_event(
                            event_logger,
                            response_msg,
                            role,
                            stop_reason,
                        )
                        _set_converse_finish_reasons(span, stop_reason)
                    else:
                        set_converse_streaming_response_span_attributes(
                            response_msg,
                            role,
                            span,
                            finish_reason=stop_reason,
                            tool_blocks=tool_blocks,
                            reasoning_blocks=reasoning_blocks,
                        )

                return event

            return partial(
                wrap,
                response_msg=partial_state["response_msg"],
                tool_blocks=partial_state["tool_blocks"],
                reasoning_blocks=partial_state["reasoning_blocks"],
                span=span,
            )

        stream._parse_event = handler(stream._parse_event)

        def _flush_partial():
            """Called when iteration ends without seeing messageStop+metadata.
            Records partial response content so the span isn't empty."""
            if span_state["saw_message_stop"]:
                return  # already flushed by messageStop handler
            if not span.is_recording():
                return
            try:
                if should_emit_events() and event_logger:
                    emit_streaming_converse_response_event(
                        event_logger,
                        partial_state["response_msg"],
                        role,
                        None,  # no finish_reason — stream was cut short
                    )
                else:
                    set_converse_streaming_response_span_attributes(
                        partial_state["response_msg"],
                        role,
                        span,
                        finish_reason=None,
                        tool_blocks=partial_state["tool_blocks"],
                        reasoning_blocks=partial_state["reasoning_blocks"],
                    )
            except Exception:  # noqa: BLE001 — never fail in instrumentation cleanup
                pass

        # Wrap the stream's __aiter__ to guarantee span.end() fires even when
        # the caller breaks early or the underlying iterator raises.
        response["stream"] = _ConverseStreamCloser(stream, span, span_state, _flush_partial)


class _ConverseStreamCloser:
    """Wraps a converse_stream `stream` so that `span.end()` is guaranteed to
    fire exactly once: normally via the patched `_parse_event` when metadata
    arrives, or as a fallback when iteration completes without metadata (early
    `break`, exception, or empty stream).

    On the fallback path, partial response content captured up to that point
    is flushed onto the span as attributes (no finish_reason) so the dashboard
    still shows what the caller consumed. The span is ended with default
    (Unset) status — per OTel guidance, early termination is not an error."""

    def __init__(self, inner, span, span_state, flush_partial):
        self._inner = inner
        self._span = span
        self._span_state = span_state
        self._flush_partial = flush_partial

    def __getattr__(self, name):
        return getattr(self._inner, name)

    def __aiter__(self):
        return self._aiter()

    async def _aiter(self):
        try:
            async for event in self._inner:
                yield event
        finally:
            if not self._span_state["ended"]:
                self._span_state["ended"] = True
                self._flush_partial()
                self._span.end()


def _get_vendor_model(modelId):
    # Docs:
    # https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles-support.html#inference-profiles-support-system
    provider = GenAiSystemValues.AWS_BEDROCK.value
    model_vendor = "imported_model"
    model = modelId

    if modelId is not None and modelId.startswith("arn"):
        components = modelId.split(":")
        if len(components) > 5:
            inf_profile = components[5].split("/")
            if len(inf_profile) == 2:
                if "." in inf_profile[1]:
                    (model_vendor, model) = _cross_region_check(inf_profile[1])
    elif modelId is not None and "." in modelId:
        (model_vendor, model) = _cross_region_check(modelId)

    return provider, model_vendor, model


def _cross_region_check(value):
    prefixes = ["us", "us-gov", "eu", "apac"]
    if any(value.startswith(prefix + ".") for prefix in prefixes):
        parts = value.split(".")
        if len(parts) > 2:
            parts.pop(0)
        return parts[0], parts[1]
    else:
        (model_vendor, model) = value.split(".", 1)
    return model_vendor, model


def _buffer_request_body(kwargs):
    """boto3 also accepts a file-like `body`. Read it once so botocore and the
    readers that run after the call (when the stream is already consumed) see
    the same bytes, as ReusableStreamingBody does for the response."""
    body = kwargs.get("body")
    if not hasattr(body, "read"):
        return
    try:
        start = body.tell()
    except Exception:
        start = None
    try:
        kwargs["body"] = body.read()
    except Exception:
        # A failed read may have moved the cursor: rewind so the caller's
        # object is left as it was, then surface the error (botocore would
        # hit the same one) instead of tracing or sending a partial body.
        if start is not None:
            try:
                body.seek(start)
            except Exception:
                pass
        raise


def _derive_operation_name(kwargs):
    """Derive operation name for invoke_model spans prior to creation."""
    body_str = kwargs.get("body")
    if body_str:
        try:
            body = json.loads(body_str)
            if isinstance(body, dict) and "messages" in body:
                return GenAiOperationNameValues.CHAT.value
        except (json.JSONDecodeError, TypeError):
            pass
    return GenAiOperationNameValues.TEXT_COMPLETION.value


class GuardrailMeters:
    GEN_AI_BEDROCK_GUARDRAIL_ACTIVATION = "gen_ai.bedrock.guardrail.activation"
    GEN_AI_BEDROCK_GUARDRAIL_LATENCY = "gen_ai.bedrock.guardrail.latency"
    GEN_AI_BEDROCK_GUARDRAIL_COVERAGE = "gen_ai.bedrock.guardrail.coverage"
    GEN_AI_BEDROCK_GUARDRAIL_SENSITIVE = "gen_ai.bedrock.guardrail.sensitive_info"
    GEN_AI_BEDROCK_GUARDRAIL_TOPICS = "gen_ai.bedrock.guardrail.topics"
    GEN_AI_BEDROCK_GUARDRAIL_CONTENT = "gen_ai.bedrock.guardrail.content"
    GEN_AI_BEDROCK_GUARDRAIL_WORDS = "gen_ai.bedrock.guardrail.words"


class PromptCaching:
    # will be moved under the AI SemConv. Not namespaced since also OpenAI supports this.
    GEN_AI_PROMPT_CACHING = "gen_ai.prompt.caching"


def _create_metrics(meter: Meter):
    token_histogram = meter.create_histogram(
        name=Meters.LLM_TOKEN_USAGE,
        unit="token",
        description="Measures number of input and output tokens used",
    )

    choice_counter = meter.create_counter(
        name=Meters.LLM_GENERATION_CHOICES,
        unit="choice",
        description="Number of choices returned by chat completions call",
    )

    duration_histogram = meter.create_histogram(
        name=Meters.LLM_OPERATION_DURATION,
        unit="s",
        description="GenAI operation duration",
    )

    exception_counter = meter.create_counter(
        name="gen_ai.bedrock.completions.exceptions",
        unit="time",
        description="Number of exceptions occurred during chat completions",
    )

    # Guardrail metrics
    guardrail_activation = meter.create_counter(
        name=GuardrailMeters.GEN_AI_BEDROCK_GUARDRAIL_ACTIVATION,
        unit="",
        description="Number of guardrail activation",
    )

    guardrail_latency_histogram = meter.create_histogram(
        name=GuardrailMeters.GEN_AI_BEDROCK_GUARDRAIL_LATENCY,
        unit="ms",
        description="GenAI guardrail latency",
    )

    guardrail_coverage = meter.create_counter(
        name=GuardrailMeters.GEN_AI_BEDROCK_GUARDRAIL_COVERAGE,
        unit="char",
        description="GenAI guardrail coverage",
    )

    guardrail_sensitive_info = meter.create_counter(
        name=GuardrailMeters.GEN_AI_BEDROCK_GUARDRAIL_SENSITIVE,
        unit="",
        description="GenAI guardrail sensitive information protection",
    )

    guardrail_topic = meter.create_counter(
        name=GuardrailMeters.GEN_AI_BEDROCK_GUARDRAIL_TOPICS,
        unit="",
        description="GenAI guardrail topics protection",
    )

    guardrail_content = meter.create_counter(
        name=GuardrailMeters.GEN_AI_BEDROCK_GUARDRAIL_CONTENT,
        unit="",
        description="GenAI guardrail content filter protection",
    )

    guardrail_words = meter.create_counter(
        name=GuardrailMeters.GEN_AI_BEDROCK_GUARDRAIL_WORDS,
        unit="",
        description="GenAI guardrail words filter protection",
    )

    # Prompt Caching
    prompt_caching = meter.create_counter(
        name=PromptCaching.GEN_AI_PROMPT_CACHING,
        unit="",
        description="Number of cached tokens",
    )

    return (
        token_histogram,
        choice_counter,
        duration_histogram,
        exception_counter,
        guardrail_activation,
        guardrail_latency_histogram,
        guardrail_coverage,
        guardrail_sensitive_info,
        guardrail_topic,
        guardrail_content,
        guardrail_words,
        prompt_caching,
    )


class BedrockInstrumentor(BaseInstrumentor):
    """An instrumentor for Bedrock's client library."""

    def __init__(
        self,
        enrich_token_usage: bool = False,
        exception_logger=None,
        use_attributes: Optional[bool] = None,
        use_legacy_attributes: Optional[bool] = None,
    ):
        super().__init__()
        if use_attributes is not None and use_legacy_attributes is not None:
            raise TypeError(
                "Cannot pass both `use_attributes` and `use_legacy_attributes`; "
                "`use_legacy_attributes` is deprecated, use `use_attributes` instead."
            )
        if use_legacy_attributes is not None:
            warnings.warn(
                "`use_legacy_attributes` is deprecated and will be removed in a "
                "future release; use `use_attributes` instead. The current OTel "
                "GenAI spec emits prompts/completions as span attributes "
                "(`gen_ai.input.messages` / `gen_ai.output.messages`), which is "
                "what `use_attributes=True` (the default) does. "
                "`use_attributes=False` opts into the events path instead.",
                DeprecationWarning,
                stacklevel=2,
            )
            use_attributes = use_legacy_attributes
        if use_attributes is None:
            use_attributes = True
        Config.enrich_token_usage = enrich_token_usage
        Config.exception_logger = exception_logger
        Config.use_legacy_attributes = use_attributes

    def instrumentation_dependencies(self) -> Collection[str]:
        return _instruments

    def _instrument(self, **kwargs):
        tracer_provider = kwargs.get("tracer_provider")
        tracer = get_tracer(__name__, __version__, tracer_provider)

        # meter and counters are inited here
        meter_provider = kwargs.get("meter_provider")
        meter = get_meter(__name__, __version__, meter_provider)

        if is_metrics_enabled():
            (
                token_histogram,
                choice_counter,
                duration_histogram,
                exception_counter,
                guardrail_activation,
                guardrail_latency_histogram,
                guardrail_coverage,
                guardrail_sensitive_info,
                guardrail_topic,
                guardrail_content,
                guardrail_words,
                prompt_caching,
            ) = _create_metrics(meter)
        else:
            (
                token_histogram,
                choice_counter,
                duration_histogram,
                exception_counter,
                guardrail_activation,
                guardrail_latency_histogram,
                guardrail_coverage,
                guardrail_sensitive_info,
                guardrail_topic,
                guardrail_content,
                guardrail_words,
                prompt_caching,
            ) = (None, None, None, None, None, None, None, None, None, None, None, None)

        metric_params = MetricParams(
            token_histogram,
            choice_counter,
            duration_histogram,
            exception_counter,
            guardrail_activation,
            guardrail_latency_histogram,
            guardrail_coverage,
            guardrail_sensitive_info,
            guardrail_topic,
            guardrail_content,
            guardrail_words,
            prompt_caching,
        )

        event_logger = None
        if not Config.use_legacy_attributes:
            logger_provider = kwargs.get("logger_provider")
            event_logger = get_logger(
                __name__, __version__, logger_provider=logger_provider
            )

        for wrapped_method in WRAPPED_METHODS:
            wrap_package = wrapped_method.get("package")
            wrap_object = wrapped_method.get("object")
            wrap_method = wrapped_method.get("method")
            if wrapped_method.get("async"):
                wrapper_factory = _wrap_async_factory(
                    tracer, metric_params, event_logger, wrapped_method
                )
            else:
                wrapper_factory = _wrap(
                    tracer, metric_params, event_logger, wrapped_method
                )
            try:
                wrap_function_wrapper(
                    wrap_package,
                    f"{wrap_object}.{wrap_method}",
                    wrapper_factory,
                )
            except (ImportError, ModuleNotFoundError):
                pass

    def _uninstrument(self, **kwargs):
        for wrapped_method in WRAPPED_METHODS:
            wrap_package = wrapped_method.get("package")
            wrap_object = wrapped_method.get("object")
            try:
                unwrap(
                    f"{wrap_package}.{wrap_object}",
                    wrapped_method.get("method"),
                )
            except (ImportError, ModuleNotFoundError, AttributeError):
                pass