"""
Amazon Bedrock SDK instrumentation.

Patches botocore's BaseClient._make_api_call to capture:
- Converse (sync)
- ConverseStream (sync)
- InvokeModel (sync)
- InvokeModelWithResponseStream (sync)

Bedrock uses boto3/botocore with dynamic client generation, so we patch at
the botocore level and filter by service name and operation.

Based on ARCHITECTURE_V2.md Section 6.3.
"""

from __future__ import annotations

import logging
import re
import time
from typing import Any, AsyncIterator, Callable, Dict, Iterator, Optional, TypeVar

import wrapt
from risicare.integrations._base import scrub_sensitive

logger = logging.getLogger(__name__)

T = TypeVar("T")

_instrumented = False

_INSTRUMENTED_OPERATIONS = frozenset({
    "Converse",
    "ConverseStream",
    "InvokeModel",
    "InvokeModelWithResponseStream",
})


def instrument_bedrock(module: Any) -> None:
    """
    Apply instrumentation to botocore module for Bedrock runtime.

    This patches:
    - botocore.client.BaseClient._make_api_call

    The wrapper filters to only instrument bedrock-runtime service calls
    for known operations (Converse, ConverseStream, InvokeModel,
    InvokeModelWithResponseStream).
    """
    global _instrumented
    if _instrumented:
        return

    try:
        # Patch the universal API call entry point
        wrapt.wrap_function_wrapper(
            module,
            "client.BaseClient._make_api_call",
            _wrap_make_api_call,
        )

        _instrumented = True
        logger.debug("Instrumented Amazon Bedrock (botocore)")

    except Exception as e:
        logger.warning(f"Failed to instrument Bedrock (botocore): {e}")


def _get_tracer() -> Optional[Any]:
    """Get the Risicare tracer if available."""
    try:
        from risicare.tracer import get_tracer

        return get_tracer()
    except ImportError:
        return None


def _get_config() -> Optional[Any]:
    """Get the Risicare client config if available."""
    try:
        from risicare.client import get_client

        client = get_client()
        return client.config if client else None
    except ImportError:
        return None


def _should_trace_content() -> bool:
    """Check if content should be traced."""
    config = _get_config()
    return config.trace_content if config else False


def _shorten_model_id(model_id: str) -> str:
    """Shorten Bedrock model ID for span names.

    'us.anthropic.claude-3-5-sonnet-20241022-v2:0' -> 'claude-3-5-sonnet'
    'anthropic.claude-3-haiku-20240307-v1:0' -> 'claude-3-haiku'
    'amazon.titan-text-express-v1' -> 'titan-text-express'
    """
    # Strip region prefix (e.g., "us." or "eu.")
    parts = model_id.split(".")
    if len(parts) >= 3:
        # Has region prefix: "us.anthropic.claude-3..."
        model_part = ".".join(parts[2:])
    elif len(parts) >= 2:
        # No region: "anthropic.claude-3..."
        model_part = ".".join(parts[1:])
    else:
        model_part = model_id

    # Strip version suffix (e.g., "-20241022-v2:0", "-v1:0", "-v1")
    # Remove date-version patterns
    model_part = re.sub(r"-\d{8}-v\d+:\d+$", "", model_part)
    model_part = re.sub(r"-v\d+:\d+$", "", model_part)
    model_part = re.sub(r"-v\d+$", "", model_part)

    return model_part or model_id


def _capture_converse_input(
    span: Any,
    api_params: Dict[str, Any],
    trace_content: bool,
) -> None:
    """Capture input messages from a Converse/ConverseStream request."""
    messages = api_params.get("messages", [])
    if not trace_content:
        span.set_attribute("gen_ai.prompt.count", len(messages))
        return

    for i, msg in enumerate(messages[:10]):  # Limit to 10 messages
        span.set_attribute(f"gen_ai.prompt.{i}.role", msg.get("role", ""))
        content_blocks = msg.get("content", [])
        if isinstance(content_blocks, list):
            text_parts = [
                b.get("text", "")
                for b in content_blocks
                if isinstance(b, dict) and "text" in b
            ]
            combined = "\n".join(text_parts)
            if combined:
                span.set_attribute(
                    f"gen_ai.prompt.{i}.content",
                    combined[:10000] if len(combined) > 10000 else combined,
                )

    # System prompt
    system = api_params.get("system", [])
    if system and trace_content:
        if isinstance(system, list):
            sys_text = " ".join(
                b.get("text", "") for b in system if isinstance(b, dict)
            )
        elif isinstance(system, str):
            sys_text = system
        else:
            sys_text = ""
        if sys_text:
            span.set_attribute(
                "gen_ai.system_prompt",
                sys_text[:5000] if len(sys_text) > 5000 else sys_text,
            )


def _capture_converse_response(
    span: Any,
    response: Dict[str, Any],
    latency_ms: float,
    trace_content: bool,
) -> None:
    """Capture response data from a Converse API call."""
    span.set_attribute("gen_ai.latency_ms", latency_ms)
    span.set_attribute("gen_ai.response.stop_reason", response.get("stopReason", ""))

    # Usage
    usage = response.get("usage", {})
    span.set_attribute("gen_ai.usage.prompt_tokens", usage.get("inputTokens", 0))
    span.set_attribute("gen_ai.usage.completion_tokens", usage.get("outputTokens", 0))
    span.set_attribute("gen_ai.usage.total_tokens", usage.get("totalTokens", 0))

    # Metrics
    metrics = response.get("metrics", {})
    if metrics.get("latencyMs"):
        span.set_attribute("gen_ai.response.bedrock.latency_ms", metrics["latencyMs"])

    # Content
    if trace_content:
        output = response.get("output", {})
        message = output.get("message", {})
        content_blocks = message.get("content", [])
        if isinstance(content_blocks, list):
            text_parts = [
                b.get("text", "")
                for b in content_blocks
                if isinstance(b, dict) and "text" in b
            ]
            combined = "\n".join(text_parts)
            if combined:
                span.set_attribute(
                    "gen_ai.completion.content",
                    combined[:10000] if len(combined) > 10000 else combined,
                )

            # Tool use
            tool_uses = [
                b for b in content_blocks
                if isinstance(b, dict) and b.get("toolUse")
            ]
            if tool_uses:
                span.set_attribute("gen_ai.completion.tool_uses", len(tool_uses))
                for j, tu in enumerate(tool_uses[:5]):  # Limit to 5
                    tool = tu.get("toolUse", {})
                    span.set_attribute(
                        f"gen_ai.completion.tool_use.{j}.name",
                        tool.get("name", ""),
                    )


def _capture_invoke_response(
    span: Any,
    response: Dict[str, Any],
    latency_ms: float,
    trace_content: bool,
) -> None:
    """Capture response data from an InvokeModel API call.

    InvokeModel returns a raw bytes body that varies by model family.
    We capture metadata but do not parse the body — it is a StreamingBody
    that must be read once, and parsing it would break user code.
    """
    span.set_attribute("gen_ai.latency_ms", latency_ms)

    # Response metadata
    metadata = response.get("ResponseMetadata", {})
    status_code = metadata.get("HTTPStatusCode", 0)
    span.set_attribute("gen_ai.response.status_code", status_code)

    # Content type tells us the response format
    content_type = response.get("contentType", "")
    span.set_attribute("gen_ai.response.content_type", content_type)


def _wrap_make_api_call(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for botocore.client.BaseClient._make_api_call.

    This wraps ALL botocore API calls, so we filter aggressively:
    1. Only instrument bedrock-runtime service
    2. Only instrument known operations (Converse, InvokeModel, etc.)
    3. Standard tracer and dedup guards
    """
    # Guard 1: Only instrument bedrock-runtime service
    service_model = getattr(instance, "_service_model", None)
    service_name = getattr(service_model, "service_name", "") if service_model else ""
    if service_name != "bedrock-runtime":
        return wrapped(*args, **kwargs)

    # Guard 2: Only instrument known operations
    operation_name = args[0] if args else ""
    if operation_name not in _INSTRUMENTED_OPERATIONS:
        return wrapped(*args, **kwargs)

    # Guard 3: Tracer check
    tracer = _get_tracer()
    if tracer is None or not tracer.is_enabled:
        return wrapped(*args, **kwargs)

    # Guard 4: Dedup check
    from risicare.integrations._dedup import is_provider_suppressed

    if is_provider_suppressed():
        return wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    # Extract model ID from api_params
    api_params = args[1] if len(args) > 1 else kwargs.get("api_params", {})
    model_id = api_params.get("modelId", "unknown")
    model_short = _shorten_model_id(model_id)

    with tracer.start_span(
        name=f"bedrock.{operation_name}/{model_short}",
        kind=SpanKind.LLM_CALL,
        attributes={
            "gen_ai.system": "bedrock",
            "gen_ai.request.model": model_id,
            "gen_ai.request.bedrock.operation": operation_name,
        },
    ) as span:
        trace_content = _should_trace_content()

        # Capture input for Converse operations
        if operation_name in ("Converse", "ConverseStream"):
            _capture_converse_input(span, api_params, trace_content)

        try:
            start_time = time.perf_counter()
            response = wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000

            if operation_name == "Converse":
                _capture_converse_response(span, response, latency_ms, trace_content)
            elif operation_name == "ConverseStream":
                span.set_attribute("gen_ai.response.stream", True)
                span.set_attribute("gen_ai.latency_ms", latency_ms)
                # Stream events are consumed by the user; we can't wrap the
                # EventStream without breaking boto3's response handling.
                # Record what we can.
            elif operation_name == "InvokeModel":
                _capture_invoke_response(span, response, latency_ms, trace_content)
            else:
                # InvokeModelWithResponseStream
                span.set_attribute("gen_ai.response.stream", True)
                span.set_attribute("gen_ai.latency_ms", latency_ms)

            return response

        except Exception as e:
            span.record_exception(e)
            span.set_attribute("error", True)
            span.set_attribute("error.type", type(e).__name__)
            span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
            raise
