"""
AWS Bedrock provider integration for Visibe.

Tracks Bedrock converse() and converse_stream() calls by wrapping
the boto3 bedrock-runtime client. Supports all models available via
the Bedrock Converse API (Claude, Llama, Mistral, Titan, Nova, etc.).

Two modes:
1. instrument(client) — wrap once, every call auto-sends its own trace
2. track(client, name) — context manager to group multiple calls into one trace
"""
import io
import json
import logging
import time
import uuid
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Optional, Tuple

from .base import TraceSummary
from ..utils import calculate_cost

logger = logging.getLogger("visibe.bedrock")


# ------------------------------------------------------------------
# Response parsing helpers
# ------------------------------------------------------------------

def _extract_bedrock_prompt(kwargs: dict) -> str:
    """Extract user prompt from Bedrock converse() kwargs.

    Checks both `messages` (for user content) and the top-level `system`
    parameter (Bedrock uses a separate kwarg, not part of messages).
    """
    parts = []

    # Capture system prompt if present (top-level kwarg, not in messages)
    system = kwargs.get('system', [])
    if system:
        for block in system:
            if isinstance(block, dict) and 'text' in block:
                parts.append(f"[system] {block['text'][:200]}")
                break

    # Capture last user message
    messages = kwargs.get('messages', [])
    for msg in reversed(messages):
        if isinstance(msg, dict) and msg.get('role') == 'user':
            for block in msg.get('content', []):
                if isinstance(block, dict) and 'text' in block:
                    parts.append(block['text'][:500])
                    break
            break

    return ' | '.join(parts) if parts else ""


def _extract_bedrock_output(response: dict) -> str:
    """Extract output text from converse() response."""
    parts = []
    message = response.get('output', {}).get('message', {})
    for block in message.get('content', []):
        if isinstance(block, dict) and 'text' in block:
            parts.append(block['text'])
    return ''.join(parts)


def _extract_bedrock_usage(response: dict) -> tuple:
    """Extract (input_tokens, output_tokens) from response."""
    usage = response.get('usage', {})
    return (usage.get('inputTokens', 0), usage.get('outputTokens', 0))


def _extract_bedrock_tool_calls(response: dict) -> list:
    """Extract tool use blocks from response content."""
    tools = []
    message = response.get('output', {}).get('message', {})
    for block in message.get('content', []):
        if isinstance(block, dict) and 'toolUse' in block:
            tu = block['toolUse']
            tools.append({
                'tool_name': tu.get('name', 'unknown'),
                'agent_name': 'bedrock',
                'status': 'success',
                'input': str(tu.get('input', ''))[:500],
            })
    return tools


def _extract_bedrock_model(kwargs: dict) -> str:
    """Extract modelId from converse() kwargs.

    Handles both simple IDs and ARN-based cross-account IDs.
    """
    model_id = kwargs.get('modelId', 'unknown')
    # ARN format: arn:aws:bedrock:region:account:inference-profile/model-id
    if model_id.startswith('arn:'):
        parts = model_id.split('/')
        if len(parts) > 1:
            return parts[-1]
    return model_id


def _detect_provider_from_model(model_id: str) -> str:
    """Detect the LLM provider from a Bedrock model ID prefix."""
    provider_map = {
        'anthropic': 'anthropic',
        'meta': 'meta',
        'mistral': 'mistral',
        'amazon': 'amazon',
        'cohere': 'cohere',
        'ai21': 'ai21',
    }
    prefix = model_id.split('.')[0] if '.' in model_id else ''
    return provider_map.get(prefix, 'bedrock')


# ------------------------------------------------------------------
# invoke_model() provider-specific parsers
# Each returns (prompt_text, output_text, input_tokens, output_tokens)
# ------------------------------------------------------------------

def _parse_invoke_anthropic(response_json: dict, request_json: dict) -> Tuple[str, str, int, int]:
    """Anthropic Claude via invoke_model."""
    # Prompt from request messages
    prompt = ""
    system = request_json.get('system', '')
    if system:
        prompt = f"[system] {str(system)[:200]} | "
    msgs = request_json.get('messages', [])
    for msg in reversed(msgs):
        if msg.get('role') == 'user':
            content = msg.get('content', '')
            if isinstance(content, list):
                for block in content:
                    if isinstance(block, dict) and block.get('type') == 'text':
                        prompt += block.get('text', '')[:500]
                        break
            elif isinstance(content, str):
                prompt += content[:500]
            break

    # Output
    output = ""
    for block in response_json.get('content', []):
        if isinstance(block, dict) and block.get('type') == 'text':
            output += block.get('text', '')

    # Tokens
    usage = response_json.get('usage', {})
    return (prompt, output, usage.get('input_tokens', 0), usage.get('output_tokens', 0))


def _parse_invoke_meta(response_json: dict, request_json: dict) -> Tuple[str, str, int, int]:
    """Meta Llama via invoke_model."""
    prompt = request_json.get('prompt', '')[:500]
    output = response_json.get('generation', '')
    return (prompt, output,
            response_json.get('prompt_token_count', 0),
            response_json.get('generation_token_count', 0))


def _parse_invoke_mistral(response_json: dict, request_json: dict) -> Tuple[str, str, int, int]:
    """Mistral via invoke_model (text completion or chat). No token counts available."""
    # Prompt
    prompt = request_json.get('prompt', '')[:500]
    if not prompt:
        msgs = request_json.get('messages', [])
        for msg in reversed(msgs):
            if msg.get('role') == 'user':
                prompt = str(msg.get('content', ''))[:500]
                break

    # Output — text completion format
    outputs = response_json.get('outputs', [])
    if outputs:
        return (prompt, outputs[0].get('text', ''), 0, 0)

    # Output — chat completion format
    choices = response_json.get('choices', [])
    if choices:
        msg = choices[0].get('message', {})
        return (prompt, msg.get('content', ''), 0, 0)

    return (prompt, str(response_json)[:500], 0, 0)


def _parse_invoke_titan(response_json: dict, request_json: dict) -> Tuple[str, str, int, int]:
    """Amazon Titan Text via invoke_model."""
    prompt = request_json.get('inputText', '')[:500]
    results = response_json.get('results', [{}])
    output = results[0].get('outputText', '') if results else ''
    input_tokens = response_json.get('inputTextTokenCount', 0)
    output_tokens = results[0].get('tokenCount', 0) if results else 0
    return (prompt, output, input_tokens, output_tokens)


def _parse_invoke_nova(response_json: dict, request_json: dict) -> Tuple[str, str, int, int]:
    """Amazon Nova via invoke_model."""
    # Prompt — Nova uses messages format like converse
    prompt = ""
    system = request_json.get('system', [])
    if system:
        for block in system:
            if isinstance(block, dict) and 'text' in block:
                prompt = f"[system] {block['text'][:200]} | "
                break
    msgs = request_json.get('messages', [])
    for msg in reversed(msgs):
        if msg.get('role') == 'user':
            for block in msg.get('content', []):
                if isinstance(block, dict) and 'text' in block:
                    prompt += block['text'][:500]
                    break
            break

    # Output
    output = ""
    message = response_json.get('output', {}).get('message', {})
    for block in message.get('content', []):
        if isinstance(block, dict) and 'text' in block:
            output += block['text']

    # Tokens
    usage = response_json.get('usage', {})
    return (prompt, output, usage.get('inputTokens', 0), usage.get('outputTokens', 0))


def _parse_invoke_cohere(response_json: dict, request_json: dict) -> Tuple[str, str, int, int]:
    """Cohere Command / Command R via invoke_model."""
    # Prompt
    prompt = request_json.get('prompt', '') or request_json.get('message', '')
    prompt = str(prompt)[:500]

    # Command R format
    if 'text' in response_json:
        meta = response_json.get('meta', {}).get('billed_units', {})
        return (prompt, response_json['text'],
                meta.get('input_tokens', 0), meta.get('output_tokens', 0))

    # Original Command format — no token counts
    gens = response_json.get('generations', [])
    if gens:
        return (prompt, gens[0].get('text', ''), 0, 0)

    return (prompt, str(response_json)[:500], 0, 0)


def _parse_invoke_ai21(response_json: dict, request_json: dict) -> Tuple[str, str, int, int]:
    """AI21 Jamba via invoke_model (OpenAI-compatible format)."""
    # Prompt
    prompt = ""
    msgs = request_json.get('messages', [])
    for msg in reversed(msgs):
        if msg.get('role') == 'user':
            prompt = str(msg.get('content', ''))[:500]
            break

    # Output
    choices = response_json.get('choices', [])
    output = choices[0].get('message', {}).get('content', '') if choices else ''

    # Tokens
    usage = response_json.get('usage', {})
    return (prompt, output, usage.get('prompt_tokens', 0), usage.get('completion_tokens', 0))


def _parse_invoke_unknown(response_json: dict, request_json: dict) -> Tuple[str, str, int, int]:
    """Fallback parser for unrecognized models."""
    return ("", str(response_json)[:500], 0, 0)


def _get_invoke_parser(model_id: str):
    """Return the right parser function based on model ID prefix."""
    model_id = model_id.lower()
    if model_id.startswith('anthropic.'):
        return _parse_invoke_anthropic
    if model_id.startswith('meta.'):
        return _parse_invoke_meta
    if model_id.startswith('mistral.'):
        return _parse_invoke_mistral
    if model_id.startswith('amazon.titan'):
        return _parse_invoke_titan
    if model_id.startswith('amazon.nova'):
        return _parse_invoke_nova
    if model_id.startswith('cohere.'):
        return _parse_invoke_cohere
    if model_id.startswith('ai21.'):
        return _parse_invoke_ai21
    return _parse_invoke_unknown


class _StreamingBodyWrapper:
    """Wraps a consumed StreamingBody so the caller can still .read() it."""

    def __init__(self, data: bytes):
        self._stream = io.BytesIO(data)

    def read(self, amt=None):
        if amt is not None:
            return self._stream.read(amt)
        return self._stream.read()

    def close(self):
        self._stream.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()


# ------------------------------------------------------------------
# Main integration class
# ------------------------------------------------------------------

class BedrockIntegration:
    """
    AWS Bedrock provider integration for Visibe observability.

    Tracks every converse() and converse_stream() call with token usage,
    cost, model info, and tool calls.

    Usage (instrument mode — each call = one trace):
        from visibe import Visibe
        from visibe.integrations.bedrock import BedrockIntegration
        import boto3

        obs = Visibe(api_key="sk_live_abc123")
        tracker = BedrockIntegration(obs)

        bedrock = boto3.client('bedrock-runtime', region_name='us-east-1')
        tracker.instrument(bedrock)

        response = bedrock.converse(
            modelId='anthropic.claude-3-haiku-20240307-v1:0',
            messages=[{"role": "user", "content": [{"text": "Hello!"}]}],
        )

        tracker.uninstrument(bedrock)

    Usage (track mode — group multiple calls into one trace):
        with tracker.track(client=bedrock, name="multi-step-chat"):
            r1 = bedrock.converse(...)
            r2 = bedrock.converse(...)
        # Both calls sent as one trace
    """

    def __init__(self, client):
        """
        Initialize Bedrock integration.

        Args:
            client: Visibe client instance
        """
        self.client = client
        self._instrumented_clients = {}

    def instrument(self, bedrock_client, name=None):
        """
        Instrument a Bedrock client so every converse(), converse_stream(),
        invoke_model(), and invoke_model_with_response_stream() call
        automatically sends its own trace.

        Args:
            bedrock_client: A boto3 bedrock-runtime client
            name: Optional trace name shown in the Visibe dashboard
        """
        client_id = id(bedrock_client)
        if client_id in self._instrumented_clients:
            logger.debug("Bedrock client is already instrumented")
            return

        original_converse = bedrock_client.converse
        original_converse_stream = bedrock_client.converse_stream
        original_invoke = bedrock_client.invoke_model
        original_invoke_stream = bedrock_client.invoke_model_with_response_stream

        self._instrumented_clients[client_id] = {
            'converse': original_converse,
            'converse_stream': original_converse_stream,
            'invoke_model': original_invoke,
            'invoke_model_with_response_stream': original_invoke_stream,
        }

        visibe_client = self.client
        integration = self
        trace_name = name

        def wrapped_converse(*args, **kwargs):
            return integration._execute_and_trace(
                original_converse, visibe_client, None, args, kwargs,
                name=trace_name,
            )

        def wrapped_converse_stream(*args, **kwargs):
            return integration._execute_and_trace_stream(
                original_converse_stream, visibe_client, None, args, kwargs,
                name=trace_name,
            )

        def wrapped_invoke(*args, **kwargs):
            return integration._execute_and_trace_invoke(
                original_invoke, visibe_client, None, args, kwargs,
                name=trace_name,
            )

        def wrapped_invoke_stream(*args, **kwargs):
            return integration._execute_and_trace_invoke_stream(
                original_invoke_stream, visibe_client, None, args, kwargs,
                name=trace_name,
            )

        bedrock_client.converse = wrapped_converse
        bedrock_client.converse_stream = wrapped_converse_stream
        bedrock_client.invoke_model = wrapped_invoke
        bedrock_client.invoke_model_with_response_stream = wrapped_invoke_stream
        logger.debug("Instrumented Bedrock client")

    def uninstrument(self, bedrock_client):
        """
        Restore all original methods on the Bedrock client.

        Args:
            bedrock_client: The previously instrumented boto3 bedrock-runtime client
        """
        client_id = id(bedrock_client)
        original = self._instrumented_clients.pop(client_id, None)
        if original:
            bedrock_client.converse = original['converse']
            bedrock_client.converse_stream = original['converse_stream']
            bedrock_client.invoke_model = original['invoke_model']
            bedrock_client.invoke_model_with_response_stream = original['invoke_model_with_response_stream']
            logger.debug("Uninstrumented Bedrock client")
        else:
            logger.warning("Bedrock client was not instrumented")

    @contextmanager
    def track(self, client, name: str = "bedrock-chat"):
        """
        Track multiple Bedrock calls grouped into a single trace.

        Args:
            client: A boto3 bedrock-runtime client
            name: Human-readable name for this trace

        Yields:
            None

        Example:
            with tracker.track(client=bedrock, name="conversation"):
                r1 = bedrock.converse(...)
                r2 = bedrock.converse(...)
        """
        tracker = _BedrockGroupTracker(self.client, client, name, self)
        tracker.start()
        try:
            yield
        except Exception as e:
            tracker.record_error(e)
            raise
        finally:
            tracker.stop()

    @staticmethod
    def _execute_and_trace(original_converse, visibe_client, group_tracker,
                           args, kwargs, name=None):
        """
        Execute a Bedrock converse() call and track it.

        If group_tracker is provided, sends span to the group's trace.
        Otherwise, creates its own trace (create -> span -> complete).
        """
        call_start = time.time()
        call_timestamp = datetime.now(timezone.utc)

        # Extract request data
        model = _extract_bedrock_model(kwargs)
        prompt_text = _extract_bedrock_prompt(kwargs)

        # If a LangChain/LangGraph trace is already active, skip standalone trace
        try:
            from visibe.integrations.langchain import _active_root_callback
            if _active_root_callback is not None:
                return original_converse(*args, **kwargs)
        except ImportError:
            pass

        # If a CrewAI trace is already active, skip standalone trace
        try:
            from visibe.integrations.crewai import _active_crew_trace
            if _active_crew_trace:
                return original_converse(*args, **kwargs)
        except ImportError:
            pass

        # For instrument mode (no group_tracker), create trace header upfront
        trace_id = None
        if not group_tracker:
            trace_id = str(uuid.uuid4())
            trace_label = name or f"bedrock-{model}"
            trace_header = {
                'trace_id': trace_id,
                'name': trace_label,
                'framework': 'bedrock',
                'model': model,
                'started_at': call_timestamp.isoformat(),
            }
            visibe_client.create_trace(trace_header)

        try:
            response = original_converse(*args, **kwargs)

            duration_sec = time.time() - call_start
            duration_ms = int(duration_sec * 1000)

            # Extract usage
            input_tokens, output_tokens = _extract_bedrock_usage(response)

            # Extract output text
            output_text = _extract_bedrock_output(response)

            # Extract tool calls
            tools_used = _extract_bedrock_tool_calls(response)

            # If tool call response with no text, summarize tool calls
            if not output_text and tools_used:
                parts = [f"{t['tool_name']}(...)" for t in tools_used]
                output_text = "; ".join(parts)

            # Calculate cost
            cost = calculate_cost(model, input_tokens, output_tokens)

            call_data = {
                'model': model,
                'input_tokens': input_tokens,
                'output_tokens': output_tokens,
                'cost': cost,
                'duration_ms': duration_ms,
                'prompt': prompt_text,
                'output': output_text,
                'tools_used': tools_used,
                'timestamp': call_timestamp,
            }

            if group_tracker:
                group_tracker.add_call(call_data)
            else:
                _send_call_as_trace(visibe_client, trace_id, call_data,
                                    model, name=name)

            return response

        except Exception as e:
            duration_ms = int((time.time() - call_start) * 1000)

            if group_tracker:
                group_tracker.record_error(e)
            else:
                _send_error_as_trace(
                    visibe_client, trace_id, model, prompt_text,
                    e, call_timestamp, duration_ms, name=name,
                )

            raise

    @staticmethod
    def _execute_and_trace_stream(original_converse_stream, visibe_client,
                                  group_tracker, args, kwargs, name=None):
        """
        Execute a Bedrock converse_stream() call and track it.

        Returns a _BedrockStreamWrapper that captures events as the
        user iterates, then finalizes the trace when the stream ends.
        """
        call_start = time.time()
        call_timestamp = datetime.now(timezone.utc)

        model = _extract_bedrock_model(kwargs)
        prompt_text = _extract_bedrock_prompt(kwargs)

        # If a LangChain/LangGraph trace is already active, skip standalone trace
        try:
            from visibe.integrations.langchain import _active_root_callback
            if _active_root_callback is not None:
                return original_converse_stream(*args, **kwargs)
        except ImportError:
            pass

        # If a CrewAI trace is already active, skip standalone trace
        try:
            from visibe.integrations.crewai import _active_crew_trace
            if _active_crew_trace:
                return original_converse_stream(*args, **kwargs)
        except ImportError:
            pass

        # For instrument mode, create trace header upfront
        trace_id = None
        if not group_tracker:
            trace_id = str(uuid.uuid4())
            trace_label = name or f"bedrock-{model}"
            trace_header = {
                'trace_id': trace_id,
                'name': trace_label,
                'framework': 'bedrock',
                'model': model,
                'started_at': call_timestamp.isoformat(),
            }
            visibe_client.create_trace(trace_header)

        try:
            response = original_converse_stream(*args, **kwargs)

            # converse_stream() returns a dict with 'stream' key
            event_stream = response.get('stream')
            if event_stream is None:
                # Unexpected — return raw response
                return response

            return _BedrockStreamWrapper(
                event_stream, response, visibe_client, group_tracker,
                trace_id, model, prompt_text, call_start, call_timestamp,
                name=name,
            )

        except Exception as e:
            duration_ms = int((time.time() - call_start) * 1000)

            if group_tracker:
                group_tracker.record_error(e)
            else:
                _send_error_as_trace(
                    visibe_client, trace_id, model, prompt_text,
                    e, call_timestamp, duration_ms, name=name,
                )

            raise

    @staticmethod
    def _execute_and_trace_invoke(original_invoke, visibe_client,
                                  group_tracker, args, kwargs, name=None):
        """
        Execute a Bedrock invoke_model() call and track it.

        invoke_model() uses provider-specific JSON bodies (unlike converse
        which has a unified format). We parse Body from kwargs, call the
        original, read the response body, parse with the right provider
        parser, then replace response['body'] so the caller can still read it.
        """
        call_start = time.time()
        call_timestamp = datetime.now(timezone.utc)

        model = kwargs.get('modelId', args[0] if args else 'unknown')
        if model.startswith('arn:'):
            parts = model.split('/')
            if len(parts) > 1:
                model = parts[-1]

        # Parse the request body
        request_body_raw = kwargs.get('body') or kwargs.get('Body', b'{}')
        if isinstance(request_body_raw, (str, bytes)):
            try:
                request_json = json.loads(request_body_raw)
            except (json.JSONDecodeError, TypeError):
                request_json = {}
        else:
            request_json = {}

        # Skip if another framework's trace is active
        try:
            from visibe.integrations.langchain import _active_root_callback
            if _active_root_callback is not None:
                return original_invoke(*args, **kwargs)
        except ImportError:
            pass
        try:
            from visibe.integrations.crewai import _active_crew_trace
            if _active_crew_trace:
                return original_invoke(*args, **kwargs)
        except ImportError:
            pass

        # For instrument mode, create trace header upfront
        trace_id = None
        if not group_tracker:
            trace_id = str(uuid.uuid4())
            trace_label = name or f"bedrock-{model}"
            trace_header = {
                'trace_id': trace_id,
                'name': trace_label,
                'framework': 'bedrock',
                'model': model,
                'started_at': call_timestamp.isoformat(),
            }
            visibe_client.create_trace(trace_header)

        try:
            response = original_invoke(*args, **kwargs)

            duration_ms = int((time.time() - call_start) * 1000)

            # Read the response body (StreamingBody) and replace it
            raw_bytes = response['body'].read()
            response['body'] = _StreamingBodyWrapper(raw_bytes)

            try:
                response_json = json.loads(raw_bytes)
            except (json.JSONDecodeError, TypeError):
                response_json = {}

            # Use provider-specific parser
            parser = _get_invoke_parser(model)
            prompt_text, output_text, input_tokens, output_tokens = parser(
                response_json, request_json
            )

            cost = calculate_cost(model, input_tokens, output_tokens)

            call_data = {
                'model': model,
                'input_tokens': input_tokens,
                'output_tokens': output_tokens,
                'cost': cost,
                'duration_ms': duration_ms,
                'prompt': prompt_text,
                'output': output_text,
                'tools_used': [],
                'timestamp': call_timestamp,
            }

            if group_tracker:
                group_tracker.add_call(call_data)
            else:
                _send_call_as_trace(visibe_client, trace_id, call_data,
                                    model, name=name)

            return response

        except Exception as e:
            duration_ms = int((time.time() - call_start) * 1000)
            prompt_text = str(request_json)[:500] if request_json else ''

            if group_tracker:
                group_tracker.record_error(e)
            else:
                _send_error_as_trace(
                    visibe_client, trace_id, model, prompt_text,
                    e, call_timestamp, duration_ms, name=name,
                )
            raise

    @staticmethod
    def _execute_and_trace_invoke_stream(original_invoke_stream, visibe_client,
                                         group_tracker, args, kwargs, name=None):
        """
        Execute invoke_model_with_response_stream() and track it.

        Returns a _BedrockInvokeStreamWrapper that collects chunks as the
        user iterates, then finalizes the trace when the stream ends.
        """
        call_start = time.time()
        call_timestamp = datetime.now(timezone.utc)

        model = kwargs.get('modelId', args[0] if args else 'unknown')
        if model.startswith('arn:'):
            parts = model.split('/')
            if len(parts) > 1:
                model = parts[-1]

        # Parse the request body
        request_body_raw = kwargs.get('body') or kwargs.get('Body', b'{}')
        if isinstance(request_body_raw, (str, bytes)):
            try:
                request_json = json.loads(request_body_raw)
            except (json.JSONDecodeError, TypeError):
                request_json = {}
        else:
            request_json = {}

        # Skip if another framework's trace is active
        try:
            from visibe.integrations.langchain import _active_root_callback
            if _active_root_callback is not None:
                return original_invoke_stream(*args, **kwargs)
        except ImportError:
            pass
        try:
            from visibe.integrations.crewai import _active_crew_trace
            if _active_crew_trace:
                return original_invoke_stream(*args, **kwargs)
        except ImportError:
            pass

        # For instrument mode, create trace header upfront
        trace_id = None
        if not group_tracker:
            trace_id = str(uuid.uuid4())
            trace_label = name or f"bedrock-{model}"
            trace_header = {
                'trace_id': trace_id,
                'name': trace_label,
                'framework': 'bedrock',
                'model': model,
                'started_at': call_timestamp.isoformat(),
            }
            visibe_client.create_trace(trace_header)

        try:
            response = original_invoke_stream(*args, **kwargs)

            event_stream = response.get('body')
            if event_stream is None:
                return response

            return _BedrockInvokeStreamWrapper(
                event_stream, response, visibe_client, group_tracker,
                trace_id, model, request_json, call_start, call_timestamp,
                name=name,
            )

        except Exception as e:
            duration_ms = int((time.time() - call_start) * 1000)
            prompt_text = str(request_json)[:500] if request_json else ''

            if group_tracker:
                group_tracker.record_error(e)
            else:
                _send_error_as_trace(
                    visibe_client, trace_id, model, prompt_text,
                    e, call_timestamp, duration_ms, name=name,
                )
            raise


# ------------------------------------------------------------------
# Stream wrapper
# ------------------------------------------------------------------

class _BedrockStreamWrapper:
    """Wraps a Bedrock converse_stream() EventStream to capture usage.

    Events include:
    - {'contentBlockDelta': {'delta': {'text': '...'}}}
    - {'metadata': {'usage': {'inputTokens': N, 'outputTokens': N}}}
    - {'messageStop': {'stopReason': 'end_turn'}}
    """

    def __init__(self, event_stream, raw_response, visibe_client,
                 group_tracker, trace_id, model, prompt_text,
                 call_start, call_timestamp, name=None):
        # EventStream is iterable but not an iterator — convert it
        self._stream = iter(event_stream)
        self._raw_response = raw_response
        self._visibe_client = visibe_client
        self._group_tracker = group_tracker
        self._trace_id = trace_id
        self._model = model
        self._prompt_text = prompt_text
        self._call_start = call_start
        self._call_timestamp = call_timestamp
        self._name = name
        self._collected_content = []
        self._input_tokens = 0
        self._output_tokens = 0
        self._tools_used = []
        self._finalized = False

    def __iter__(self):
        return self

    def __next__(self):
        try:
            event = next(self._stream)
            self._process_event(event)
            return event
        except StopIteration:
            self._finalize()
            raise

    def close(self):
        """Allow early termination — finalize trace even if stream not fully consumed."""
        if not self._finalized:
            self._finalize()

    def __del__(self):
        """Safety net: finalize if stream was abandoned without close()."""
        if not self._finalized:
            try:
                self._finalize()
            except Exception:
                pass  # Best-effort in __del__

    def _process_event(self, event):
        """Process a single stream event."""
        if 'contentBlockDelta' in event:
            delta = event['contentBlockDelta'].get('delta', {})
            if 'text' in delta:
                self._collected_content.append(delta['text'])

        if 'contentBlockStop' in event:
            # Could track tool use blocks here in the future
            pass

        if 'metadata' in event:
            usage = event['metadata'].get('usage', {})
            self._input_tokens = usage.get('inputTokens', 0)
            self._output_tokens = usage.get('outputTokens', 0)

    def _finalize(self):
        """Send trace data after stream completes or is closed."""
        if self._finalized:
            return
        self._finalized = True

        duration_sec = time.time() - self._call_start
        duration_ms = int(duration_sec * 1000)
        output_text = ''.join(self._collected_content)

        cost = calculate_cost(
            self._model, self._input_tokens, self._output_tokens
        )

        call_data = {
            'model': self._model,
            'input_tokens': self._input_tokens,
            'output_tokens': self._output_tokens,
            'cost': cost,
            'duration_ms': duration_ms,
            'prompt': self._prompt_text,
            'output': output_text,
            'tools_used': self._tools_used,
            'timestamp': self._call_timestamp,
        }

        if self._group_tracker:
            self._group_tracker.add_call(call_data)
        else:
            _send_call_as_trace(self._visibe_client, self._trace_id,
                                call_data, self._model, name=self._name)


# ------------------------------------------------------------------
# invoke_model stream wrapper
# ------------------------------------------------------------------

class _BedrockInvokeStreamWrapper:
    """Wraps invoke_model_with_response_stream() EventStream.

    Each event is {'chunk': {'bytes': b'<json>'}}.
    We collect all chunk bytes, then on finalization parse the accumulated
    response JSON using the provider-specific parser.
    """

    def __init__(self, event_stream, raw_response, visibe_client,
                 group_tracker, trace_id, model, request_json,
                 call_start, call_timestamp, name=None):
        self._stream = iter(event_stream)
        self._raw_response = raw_response
        self._visibe_client = visibe_client
        self._group_tracker = group_tracker
        self._trace_id = trace_id
        self._model = model
        self._request_json = request_json
        self._call_start = call_start
        self._call_timestamp = call_timestamp
        self._name = name
        self._collected_chunks = []
        self._finalized = False

    def __iter__(self):
        return self

    def __next__(self):
        try:
            event = next(self._stream)
            # Collect raw chunk bytes for later parsing
            chunk = event.get('chunk', {})
            if 'bytes' in chunk:
                self._collected_chunks.append(chunk['bytes'])
            return event
        except StopIteration:
            self._finalize()
            raise

    def close(self):
        if not self._finalized:
            self._finalize()

    def __del__(self):
        if not self._finalized:
            try:
                self._finalize()
            except Exception:
                pass

    def _finalize(self):
        if self._finalized:
            return
        self._finalized = True

        duration_ms = int((time.time() - self._call_start) * 1000)

        # Concatenate all chunk bytes and parse as JSON
        all_bytes = b''.join(self._collected_chunks)
        try:
            response_json = json.loads(all_bytes)
        except (json.JSONDecodeError, TypeError):
            # Some providers send multiple JSON objects per chunk — try
            # parsing each chunk individually and merge
            response_json = self._parse_chunks_individually()

        # If we got accumulated data from chunk-by-chunk parsing, use it directly
        if '_accumulated_text' in response_json:
            # Extract prompt from request using the parser
            parser = _get_invoke_parser(self._model)
            prompt_text, _, _, _ = parser({}, self._request_json)
            output_text = response_json['_accumulated_text']
            input_tokens = response_json['_input_tokens']
            output_tokens = response_json['_output_tokens']
        else:
            parser = _get_invoke_parser(self._model)
            prompt_text, output_text, input_tokens, output_tokens = parser(
                response_json, self._request_json
            )

        cost = calculate_cost(self._model, input_tokens, output_tokens)

        call_data = {
            'model': self._model,
            'input_tokens': input_tokens,
            'output_tokens': output_tokens,
            'cost': cost,
            'duration_ms': duration_ms,
            'prompt': prompt_text,
            'output': output_text,
            'tools_used': [],
            'timestamp': self._call_timestamp,
        }

        if self._group_tracker:
            self._group_tracker.add_call(call_data)
        else:
            _send_call_as_trace(self._visibe_client, self._trace_id,
                                call_data, self._model, name=self._name)

    def _parse_chunks_individually(self) -> dict:
        """Fallback: parse each chunk as separate JSON, accumulate text and usage."""
        accumulated_text = []
        input_tokens = 0
        output_tokens = 0

        for chunk_bytes in self._collected_chunks:
            try:
                chunk_json = json.loads(chunk_bytes)
            except (json.JSONDecodeError, TypeError):
                continue

            # Try common streaming text locations
            text = ""
            # Nova / Converse-style streaming: contentBlockDelta events
            if 'contentBlockDelta' in chunk_json:
                delta = chunk_json['contentBlockDelta'].get('delta', {})
                text = delta.get('text', '')
            elif 'output' in chunk_json:
                msg = chunk_json.get('output', {}).get('message', {})
                for block in msg.get('content', []):
                    if isinstance(block, dict) and 'text' in block:
                        text += block['text']
            elif 'generation' in chunk_json:
                text = chunk_json.get('generation', '')
            elif 'outputText' in chunk_json:
                text = chunk_json.get('outputText', '')
            elif 'content' in chunk_json and isinstance(chunk_json['content'], list):
                for block in chunk_json['content']:
                    if isinstance(block, dict) and block.get('type') == 'text':
                        text += block.get('text', '')
            elif 'choices' in chunk_json:
                choices = chunk_json.get('choices', [])
                if choices:
                    delta = choices[0].get('delta', choices[0].get('message', {}))
                    text = delta.get('content', '') or ''

            if text:
                accumulated_text.append(text)

            # Accumulate token usage from any chunk that has it
            # Nova streams usage inside 'metadata' key
            usage = chunk_json.get('usage', {}) or chunk_json.get('metadata', {}).get('usage', {})
            if usage:
                input_tokens = usage.get('inputTokens', 0) or usage.get('input_tokens', 0) or usage.get('prompt_tokens', 0) or input_tokens
                output_tokens = usage.get('outputTokens', 0) or usage.get('output_tokens', 0) or usage.get('completion_tokens', 0) or output_tokens

            # Bedrock invocation metrics (present in final chunk for all providers)
            metrics = chunk_json.get('amazon-bedrock-invocationMetrics', {})
            if metrics:
                input_tokens = metrics.get('inputTokenCount', input_tokens)
                output_tokens = metrics.get('outputTokenCount', output_tokens)

            # Titan style
            if 'inputTextTokenCount' in chunk_json:
                input_tokens = chunk_json['inputTextTokenCount']
            if 'totalOutputTextTokenCount' in chunk_json:
                output_tokens = chunk_json['totalOutputTextTokenCount']

            # Meta style
            if 'prompt_token_count' in chunk_json:
                input_tokens = chunk_json['prompt_token_count']
            if 'generation_token_count' in chunk_json:
                output_tokens = chunk_json['generation_token_count']

        # Build a synthetic response_json the parser can understand
        return {
            '_accumulated_text': ''.join(accumulated_text),
            '_input_tokens': input_tokens,
            '_output_tokens': output_tokens,
        }


# ------------------------------------------------------------------
# Group tracker (for track() context manager)
# ------------------------------------------------------------------

class _BedrockGroupTracker:
    """Internal tracker for grouping multiple Bedrock calls into one trace."""

    def __init__(self, visibe_client, bedrock_client, name: str,
                 integration: 'BedrockIntegration'):
        self.visibe_client = visibe_client
        self.bedrock_client = bedrock_client
        self.name = name
        self._integration = integration
        self.trace_id = str(uuid.uuid4())

        self.start_time = None
        self.end_time = None
        self.calls: List[Dict] = []
        self.errors: List[Dict] = []
        self._step_counter = 0

        self._saved = {}    # method_name -> saved wrapper/original
        self._originals = {}  # method_name -> true original

    def _next_step_id(self) -> str:
        self._step_counter += 1
        return f"step_{self._step_counter}"

    def start(self):
        """Patch the client and start tracking."""
        self.start_time = datetime.now(timezone.utc)

        # Create trace header on backend
        self.visibe_client.create_trace({
            'trace_id': self.trace_id,
            'name': self.name,
            'framework': 'bedrock',
            'started_at': self.start_time.isoformat(),
        })

        methods = ['converse', 'converse_stream',
                    'invoke_model', 'invoke_model_with_response_stream']

        # Save whatever is currently on the client (may be instrument wrapper)
        for m in methods:
            self._saved[m] = getattr(self.bedrock_client, m)

        # If client is already instrumented, use the real original to avoid
        # double-tracking (instrument + track both sending traces)
        client_id = id(self.bedrock_client)
        if client_id in self._integration._instrumented_clients:
            entry = self._integration._instrumented_clients[client_id]
            for m in methods:
                self._originals[m] = entry.get(m, getattr(self.bedrock_client, m))
        else:
            for m in methods:
                self._originals[m] = getattr(self.bedrock_client, m)

        tracker = self
        orig_converse = self._originals['converse']
        orig_converse_stream = self._originals['converse_stream']
        orig_invoke = self._originals['invoke_model']
        orig_invoke_stream = self._originals['invoke_model_with_response_stream']

        def wrapped_converse(*args, **kwargs):
            return BedrockIntegration._execute_and_trace(
                orig_converse, tracker.visibe_client, tracker,
                args, kwargs,
            )

        def wrapped_converse_stream(*args, **kwargs):
            return BedrockIntegration._execute_and_trace_stream(
                orig_converse_stream, tracker.visibe_client, tracker,
                args, kwargs,
            )

        def wrapped_invoke(*args, **kwargs):
            return BedrockIntegration._execute_and_trace_invoke(
                orig_invoke, tracker.visibe_client, tracker,
                args, kwargs,
            )

        def wrapped_invoke_stream(*args, **kwargs):
            return BedrockIntegration._execute_and_trace_invoke_stream(
                orig_invoke_stream, tracker.visibe_client, tracker,
                args, kwargs,
            )

        self.bedrock_client.converse = wrapped_converse
        self.bedrock_client.converse_stream = wrapped_converse_stream
        self.bedrock_client.invoke_model = wrapped_invoke
        self.bedrock_client.invoke_model_with_response_stream = wrapped_invoke_stream

    def stop(self):
        """Restore original methods and complete the trace."""
        self.end_time = datetime.now(timezone.utc)

        # Restore what was on the client before track() started
        for m, fn in self._saved.items():
            setattr(self.bedrock_client, m, fn)

        self._complete_trace()

    def add_call(self, call_data: Dict):
        """Add a completed call — sends span to backend immediately."""
        self.calls.append(call_data)

        provider = _detect_provider_from_model(call_data['model'])

        span = {
            'span_id': self._next_step_id(),
            'type': 'llm_call',
            'timestamp': call_data['timestamp'].isoformat(),
            'agent_name': 'bedrock',
            'model': call_data['model'],
            'provider': provider,
            'status': 'success',
            'description': f"LLM Call using {call_data['model']}",
            'input_tokens': call_data['input_tokens'],
            'output_tokens': call_data['output_tokens'],
            'cost': call_data['cost'],
            'duration_ms': call_data['duration_ms'],
            'input_text': (call_data.get('prompt') or '')[:1000],
            'output_text': (call_data.get('output') or '')[:1000],
        }
        self.visibe_client.queue_span(self.trace_id, span)

    def record_error(self, exc: Exception):
        """Record an error — sends error span to backend."""
        self.errors.append({
            'timestamp': datetime.now(timezone.utc).isoformat(),
            'error_type': type(exc).__name__,
            'message': str(exc)[:500],
        })

        span = {
            'span_id': self._next_step_id(),
            'type': 'error',
            'timestamp': datetime.now(timezone.utc).isoformat(),
            'description': f"Error: {type(exc).__name__}",
            'error_type': type(exc).__name__,
            'error_message': str(exc)[:500],
        }
        self.visibe_client.queue_span(self.trace_id, span)

    def _complete_trace(self):
        """Complete the trace with summary aggregates."""
        duration_ms = int(
            (self.end_time - self.start_time).total_seconds() * 1000
        )

        total_input = sum(c['input_tokens'] for c in self.calls)
        total_output = sum(c['output_tokens'] for c in self.calls)
        total_cost = round(sum(c['cost'] for c in self.calls), 6)

        prompt = self.calls[0]['prompt'] if self.calls else self.name
        model = self.calls[0]['model'] if self.calls else None
        status = 'failed' if self.errors else 'completed'

        summary = {
            'status': status,
            'ended_at': self.end_time.isoformat(),
            'duration_ms': duration_ms,
            'prompt': prompt,
            'model': model,
            'total_cost': total_cost,
            'total_tokens': total_input + total_output,
            'total_input_tokens': total_input,
            'total_output_tokens': total_output,
        }

        success = self.visibe_client.complete_trace(self.trace_id, summary)
        if success:
            logger.info(f"Bedrock trace completed: {self.trace_id}")
        else:
            logger.error(f"Failed to complete Bedrock trace: {self.trace_id}")

        # Print TraceSummary to stdout for developer feedback
        total_tool_calls = sum(
            len(c.get('tools_used', [])) for c in self.calls
        )
        trace_summary = TraceSummary(
            name=self.name,
            status=status,
            llm_calls=len(self.calls),
            tool_calls=total_tool_calls,
            total_tokens=total_input + total_output,
            total_cost=total_cost,
            duration_s=round(duration_ms / 1000, 1),
            agents=['bedrock'],
            sent=success,
        )
        print(f"[Visibe] {trace_summary}")


# ------------------------------------------------------------------
# Helper functions (instrument mode — single-call traces)
# ------------------------------------------------------------------

def _send_call_as_trace(visibe_client, trace_id: str, call_data: dict,
                        model: str, name: str = None):
    """Send a single-call trace: span + complete (instrument mode)."""
    now = call_data['timestamp']
    end_time = now + timedelta(milliseconds=call_data['duration_ms'])

    provider = _detect_provider_from_model(model)

    # Send LLM span
    span = {
        'span_id': 'step_1',
        'type': 'llm_call',
        'timestamp': now.isoformat(),
        'agent_name': 'bedrock',
        'model': call_data['model'],
        'provider': provider,
        'status': 'success',
        'description': f"LLM Call using {call_data['model']}",
        'input_tokens': call_data['input_tokens'],
        'output_tokens': call_data['output_tokens'],
        'cost': call_data['cost'],
        'duration_ms': call_data['duration_ms'],
        'input_text': (call_data.get('prompt') or '')[:1000],
        'output_text': (call_data.get('output') or '')[:1000],
    }
    visibe_client.queue_span(trace_id, span)

    # Send tool spans if any
    for i, tool in enumerate(call_data.get('tools_used', []), start=2):
        tool_span = {
            'span_id': f'step_{i}',
            'type': 'tool_call',
            'timestamp': now.isoformat(),
            'tool_name': tool['tool_name'],
            'agent_name': tool.get('agent_name', 'bedrock'),
            'status': tool.get('status', 'success'),
            'input_text': tool.get('input', ''),
            'output_text': tool.get('output', ''),
        }
        visibe_client.queue_span(trace_id, tool_span)

    # Complete trace
    summary = {
        'status': 'completed',
        'ended_at': end_time.isoformat(),
        'duration_ms': call_data['duration_ms'],
        'prompt': call_data.get('prompt') or 'No prompt captured',
        'model': call_data['model'],
        'total_cost': call_data['cost'],
        'total_tokens': call_data['input_tokens'] + call_data['output_tokens'],
        'total_input_tokens': call_data['input_tokens'],
        'total_output_tokens': call_data['output_tokens'],
    }
    success = visibe_client.complete_trace(trace_id, summary)

    # Print TraceSummary to stdout for developer feedback
    trace_summary = TraceSummary(
        name=name or call_data['model'],
        status='completed',
        llm_calls=1,
        tool_calls=len(call_data.get('tools_used', [])),
        total_tokens=call_data['input_tokens'] + call_data['output_tokens'],
        total_cost=call_data['cost'],
        duration_s=round(call_data['duration_ms'] / 1000, 1),
        agents=['bedrock'],
        sent=success,
    )
    print(f"[Visibe] {trace_summary}")


def _send_error_as_trace(visibe_client, trace_id: str, model: str,
                         prompt: str, exc: Exception,
                         timestamp: datetime, duration_ms: int,
                         name: str = None):
    """Send a failed trace: error span + complete as failed (instrument mode)."""
    end_time = timestamp + timedelta(milliseconds=duration_ms)

    # Send error span
    span = {
        'span_id': 'step_1',
        'type': 'error',
        'timestamp': timestamp.isoformat(),
        'description': f"Error: {type(exc).__name__}",
        'error_type': type(exc).__name__,
        'error_message': str(exc)[:500],
        'duration_ms': duration_ms,
    }
    visibe_client.queue_span(trace_id, span)

    # Complete trace as failed
    summary = {
        'status': 'failed',
        'ended_at': end_time.isoformat(),
        'duration_ms': duration_ms,
        'prompt': prompt or 'No prompt captured',
        'total_cost': 0.0,
        'total_tokens': 0,
        'total_input_tokens': 0,
        'total_output_tokens': 0,
    }
    visibe_client.complete_trace(trace_id, summary)
