"""
AWS Bedrock Instrumentation - Event-Driven Capture

Instruments AWS Bedrock's boto3 SDK to automatically capture model calls
across all Bedrock-supported providers as events for detection and analysis.

Supports:
- Anthropic Claude models
- Meta Llama models
- Mistral models
- Amazon Titan models
- Cohere models
- And other Bedrock-hosted models

Uses the Converse API (unified interface) and InvokeModel API.
"""
from typing import Collection, Optional
import logging
import json

from wrapt import wrap_function_wrapper

from evoke.wrappers.base import BaseInstrumentor

logger = logging.getLogger(__name__)


class EvokeBedrockInstrumentor(BaseInstrumentor):
    """Instrumentor for AWS Bedrock via boto3 - covers all Bedrock-hosted models"""

    # Operations we want to capture
    INSTRUMENTED_OPERATIONS = {"Converse", "ConverseStream", "InvokeModel", "InvokeModelWithResponseStream"}

    def instrumentation_dependencies(self) -> Collection[str]:
        return ["boto3 >= 1.28.0"]

    def _instrument(self, **kwargs):
        """Instrument AWS Bedrock via botocore"""
        try:
            wrap_function_wrapper(
                "botocore.client",
                "BaseClient._make_api_call",
                self._api_call_wrapper(),
            )
            logger.debug("AWS Bedrock (boto3) instrumented")
        except Exception as e:
            logger.debug(f"Could not instrument botocore: {e}")

    def _uninstrument(self, **kwargs):
        """Uninstrument AWS Bedrock"""
        try:
            import botocore.client

            if hasattr(botocore.client.BaseClient._make_api_call, "__wrapped__"):
                botocore.client.BaseClient._make_api_call = (
                    botocore.client.BaseClient._make_api_call.__wrapped__
                )
        except Exception as e:
            logger.warning(f"Error uninstrumenting Bedrock: {e}")

    def _api_call_wrapper(self):
        """Wrapper for botocore BaseClient._make_api_call"""

        def wrapper(wrapped, instance, args, kwargs):
            # Extract operation name and parameters
            operation_name = args[0] if args else ""
            api_params = args[1] if len(args) > 1 else {}

            # Check if this is a bedrock-runtime service
            if not self._is_bedrock_runtime(instance):
                return wrapped(*args, **kwargs)

            # Only instrument specific operations
            if operation_name not in self.INSTRUMENTED_OPERATIONS:
                return wrapped(*args, **kwargs)

            # Route to appropriate handler
            if operation_name == "Converse":
                return self._handle_converse(wrapped, args, kwargs, api_params)
            elif operation_name == "InvokeModel":
                return self._handle_invoke_model(wrapped, args, kwargs, api_params)
            else:
                # For streaming operations, just pass through for now
                return wrapped(*args, **kwargs)

        return wrapper

    def _handle_converse(self, wrapped, args, kwargs, api_params):
        """Handle Converse API calls"""
        model_id = api_params.get("modelId", "unknown")
        messages = api_params.get("messages", [])
        system = api_params.get("system", [])

        # Extract provider and normalize model name
        provider = self._extract_provider(model_id)
        model_name = self._normalize_model_name(model_id)

        # Format input messages
        input_messages = self._format_converse_messages(messages, system)

        try:
            result = wrapped(*args, **kwargs)

            # Extract response
            output_messages = self._extract_converse_response(result)
            tokens_in, tokens_out = self._extract_converse_tokens(result)

            # Build metadata
            metadata = {
                "bedrock_model_id": model_id,
                "stop_reason": result.get("stopReason"),
            }

            # Deduplicate
            response_id = result.get("ResponseMetadata", {}).get("RequestId")
            if not self._should_capture_global(response_id):
                return result

            # Emit event
            self.capture_model_event(
                input_messages=input_messages,
                output_messages=output_messages,
                model_name=model_name,
                provider=provider,
                tokens_in=tokens_in,
                tokens_out=tokens_out,
                metadata=metadata,
            )

            return result

        except Exception as e:
            self.capture_error_event(
                error_message=str(e),
                error_type=type(e).__name__,
                metadata={"model": model_id, "provider": provider},
            )
            raise

    def _handle_invoke_model(self, wrapped, args, kwargs, api_params):
        """Handle InvokeModel API calls (model-specific formats)"""
        model_id = api_params.get("modelId", "unknown")
        body = api_params.get("body", b"")

        # Extract provider and normalize model name
        provider = self._extract_provider(model_id)
        model_type = self._get_model_type(model_id)
        model_name = self._normalize_model_name(model_id)

        # Try to parse body as JSON
        input_messages = []
        try:
            if isinstance(body, bytes):
                body_dict = json.loads(body.decode("utf-8"))
            elif isinstance(body, str):
                body_dict = json.loads(body)
            else:
                body_dict = {}

            # Extract messages based on model type format
            input_messages = self._extract_invoke_model_input(body_dict, model_type)
        except Exception:
            input_messages = [{"role": "user", "content": "[binary/unparseable input]"}]

        try:
            result = wrapped(*args, **kwargs)

            # Parse response body
            output_messages = []
            tokens_in = None
            tokens_out = None

            try:
                response_body = result.get("body")
                if response_body:
                    response_dict = json.loads(response_body.read())
                    output_messages, tokens_in, tokens_out = self._extract_invoke_model_output(
                        response_dict, model_type
                    )
            except Exception:
                output_messages = [{"role": "assistant", "content": "[unparseable response]"}]

            # Build metadata
            metadata = {
                "bedrock_model_id": model_id,
                "content_type": api_params.get("contentType", "application/json"),
            }

            # Deduplicate
            response_id = result.get("ResponseMetadata", {}).get("RequestId")
            if not self._should_capture_global(response_id):
                return result

            # Emit event
            self.capture_model_event(
                input_messages=input_messages,
                output_messages=output_messages,
                model_name=model_name,
                provider=provider,
                tokens_in=tokens_in,
                tokens_out=tokens_out,
                metadata=metadata,
            )

            return result

        except Exception as e:
            self.capture_error_event(
                error_message=str(e),
                error_type=type(e).__name__,
                metadata={"model": model_id, "provider": provider},
            )
            raise

    # === Helper Methods ===

    @staticmethod
    def _is_bedrock_runtime(instance) -> bool:
        """Check if the client is a bedrock-runtime client"""
        try:
            service_model = getattr(instance, "_service_model", None)
            if service_model:
                service_name = getattr(service_model, "service_name", "")
                return service_name == "bedrock-runtime"
        except Exception:
            pass
        return False

    @staticmethod
    def _extract_provider(model_id: str) -> str:
        """Return provider name for Bedrock (used in event capture)"""
        return "aws-bedrock"

    @staticmethod
    def _get_model_type(model_id: str) -> str:
        """Get model type for parsing InvokeModel request/response formats"""
        model_lower = model_id.lower()
        if model_lower.startswith("anthropic."):
            return "anthropic"
        elif model_lower.startswith("meta."):
            return "meta"
        elif model_lower.startswith("mistral."):
            return "mistral"
        elif model_lower.startswith("amazon."):
            return "amazon"
        elif model_lower.startswith("cohere."):
            return "cohere"
        elif model_lower.startswith("ai21."):
            return "ai21"
        elif model_lower.startswith("stability."):
            return "stability"
        return "unknown"

    @staticmethod
    def _normalize_model_name(model_id: str) -> str:
        """Strip provider prefix from model ID"""
        prefixes = ["anthropic.", "meta.", "mistral.", "amazon.", "cohere.", "ai21.", "stability."]
        for prefix in prefixes:
            if model_id.lower().startswith(prefix):
                return model_id[len(prefix):]
        return model_id

    @staticmethod
    def _format_converse_messages(messages: list, system: list = None) -> list:
        """Format Converse API messages as structured list"""
        formatted = []

        # Add system messages
        if system:
            for sys_msg in system:
                if isinstance(sys_msg, dict) and "text" in sys_msg:
                    formatted.append({"role": "system", "content": sys_msg["text"]})

        # Format conversation messages
        for msg in messages:
            if isinstance(msg, dict):
                role = msg.get("role", "user")
                content_blocks = msg.get("content", [])

                # Extract text from content blocks
                text_parts = []
                for block in content_blocks:
                    if isinstance(block, dict) and "text" in block:
                        text_parts.append(block["text"])

                formatted.append({"role": role, "content": " ".join(text_parts)})

        return formatted

    @staticmethod
    def _extract_converse_response(result: dict) -> list:
        """Extract response from Converse API result"""
        output = result.get("output", {})
        message = output.get("message", {})

        role = message.get("role", "assistant")
        content_blocks = message.get("content", [])

        # Extract text from content blocks
        text_parts = []
        tool_calls = []

        for block in content_blocks:
            if isinstance(block, dict):
                if "text" in block:
                    text_parts.append(block["text"])
                elif "toolUse" in block:
                    tool_use = block["toolUse"]
                    tool_calls.append({
                        "id": tool_use.get("toolUseId"),
                        "name": tool_use.get("name"),
                        "arguments": str(tool_use.get("input", {})),
                    })

        msg_dict = {"role": role, "content": " ".join(text_parts)}
        if tool_calls:
            msg_dict["tool_calls"] = tool_calls

        return [msg_dict]

    @staticmethod
    def _extract_converse_tokens(result: dict) -> tuple:
        """Extract token counts from Converse API result"""
        usage = result.get("usage", {})
        tokens_in = usage.get("inputTokens")
        tokens_out = usage.get("outputTokens")
        return tokens_in, tokens_out

    @staticmethod
    def _extract_invoke_model_input(body: dict, model_type: str) -> list:
        """Extract input messages from InvokeModel body based on model type"""
        messages = []

        # Anthropic Claude format
        if model_type == "anthropic":
            if "messages" in body:
                for msg in body["messages"]:
                    role = msg.get("role", "user")
                    content = msg.get("content", "")
                    if isinstance(content, list):
                        text_parts = [c.get("text", "") for c in content if c.get("type") == "text"]
                        content = " ".join(text_parts)
                    messages.append({"role": role, "content": content})
            elif "prompt" in body:
                messages.append({"role": "user", "content": body["prompt"]})

        # Meta Llama format
        elif model_type == "meta":
            if "prompt" in body:
                messages.append({"role": "user", "content": body["prompt"]})

        # Amazon Titan format
        elif model_type == "amazon":
            if "inputText" in body:
                messages.append({"role": "user", "content": body["inputText"]})

        # Cohere format
        elif model_type == "cohere":
            if "message" in body:
                messages.append({"role": "user", "content": body["message"]})
            elif "prompt" in body:
                messages.append({"role": "user", "content": body["prompt"]})

        # Mistral format
        elif model_type == "mistral":
            if "prompt" in body:
                messages.append({"role": "user", "content": body["prompt"]})

        # Fallback
        if not messages:
            messages.append({"role": "user", "content": str(body)[:1000]})

        return messages

    @staticmethod
    def _extract_invoke_model_output(response: dict, model_type: str) -> tuple:
        """Extract output from InvokeModel response based on model type"""
        content = ""
        tokens_in = None
        tokens_out = None

        # Anthropic Claude format
        if model_type == "anthropic":
            if "content" in response:
                text_parts = []
                for block in response["content"]:
                    if isinstance(block, dict) and block.get("type") == "text":
                        text_parts.append(block.get("text", ""))
                content = " ".join(text_parts)
            if "usage" in response:
                tokens_in = response["usage"].get("input_tokens")
                tokens_out = response["usage"].get("output_tokens")

        # Meta Llama format
        elif model_type == "meta":
            content = response.get("generation", "")
            tokens_in = response.get("prompt_token_count")
            tokens_out = response.get("generation_token_count")

        # Amazon Titan format
        elif model_type == "amazon":
            results = response.get("results", [])
            if results:
                content = results[0].get("outputText", "")
            tokens_in = response.get("inputTextTokenCount")
            tokens_out = response.get("results", [{}])[0].get("tokenCount")

        # Cohere format
        elif model_type == "cohere":
            if "text" in response:
                content = response["text"]
            elif "generations" in response:
                content = response["generations"][0].get("text", "")

        # Mistral format
        elif model_type == "mistral":
            outputs = response.get("outputs", [])
            if outputs:
                content = outputs[0].get("text", "")

        # Fallback
        if not content:
            content = str(response)[:500]

        return [{"role": "assistant", "content": content}], tokens_in, tokens_out
