"""
Azure AI Inference Instrumentation - Event-Driven Capture

Instruments Azure AI Inference SDK (azure-ai-inference) to automatically capture
model calls as events for detection and analysis.

This SDK is used for Azure AI Foundry models including:
- Azure OpenAI models (when accessed via Foundry)
- Mistral models
- Meta Llama models
- Cohere models
- And other Foundry models

The SDK provides a unified interface via ChatCompletionsClient and EmbeddingsClient.
"""
from typing import Collection, Optional
import logging

from wrapt import wrap_function_wrapper

from evoke.wrappers.base import BaseInstrumentor

logger = logging.getLogger(__name__)


class EvokeAzureInferenceInstrumentor(BaseInstrumentor):
    """Instrumentor for Azure AI Inference SDK - emits events for Foundry model calls"""

    def instrumentation_dependencies(self) -> Collection[str]:
        return ["azure-ai-inference >= 1.0.0b5"]

    def _instrument(self, **kwargs):
        """Instrument Azure AI Inference SDK"""
        # Wrap sync ChatCompletionsClient.complete
        try:
            wrap_function_wrapper(
                "azure.ai.inference",
                "ChatCompletionsClient.complete",
                self._complete_wrapper(),
            )
        except Exception as e:
            logger.debug(f"Could not instrument ChatCompletionsClient.complete: {e}")

        # Wrap async ChatCompletionsClient.complete
        try:
            wrap_function_wrapper(
                "azure.ai.inference.aio",
                "ChatCompletionsClient.complete",
                self._async_complete_wrapper(),
            )
        except Exception as e:
            logger.debug(f"Could not instrument async ChatCompletionsClient.complete: {e}")

        # Wrap sync EmbeddingsClient.embed
        try:
            wrap_function_wrapper(
                "azure.ai.inference",
                "EmbeddingsClient.embed",
                self._embed_wrapper(),
            )
        except Exception as e:
            logger.debug(f"Could not instrument EmbeddingsClient.embed: {e}")

        # Wrap async EmbeddingsClient.embed
        try:
            wrap_function_wrapper(
                "azure.ai.inference.aio",
                "EmbeddingsClient.embed",
                self._async_embed_wrapper(),
            )
        except Exception as e:
            logger.debug(f"Could not instrument async EmbeddingsClient.embed: {e}")

        logger.debug("Azure AI Inference SDK instrumented")

    def _uninstrument(self, **kwargs):
        """Uninstrument Azure AI Inference SDK"""
        try:
            import azure.ai.inference
            import azure.ai.inference.aio

            # Restore original methods
            if hasattr(azure.ai.inference.ChatCompletionsClient.complete, "__wrapped__"):
                azure.ai.inference.ChatCompletionsClient.complete = (
                    azure.ai.inference.ChatCompletionsClient.complete.__wrapped__
                )
            if hasattr(azure.ai.inference.aio.ChatCompletionsClient.complete, "__wrapped__"):
                azure.ai.inference.aio.ChatCompletionsClient.complete = (
                    azure.ai.inference.aio.ChatCompletionsClient.complete.__wrapped__
                )
            if hasattr(azure.ai.inference.EmbeddingsClient.embed, "__wrapped__"):
                azure.ai.inference.EmbeddingsClient.embed = (
                    azure.ai.inference.EmbeddingsClient.embed.__wrapped__
                )
            if hasattr(azure.ai.inference.aio.EmbeddingsClient.embed, "__wrapped__"):
                azure.ai.inference.aio.EmbeddingsClient.embed = (
                    azure.ai.inference.aio.EmbeddingsClient.embed.__wrapped__
                )
        except Exception as e:
            logger.warning(f"Error uninstrumenting Azure AI Inference: {e}")

    def _complete_wrapper(self):
        """Wrapper for ChatCompletionsClient.complete() - emits model call event"""

        def wrapper(wrapped, instance, args, kwargs):
            # Extract request info
            messages = kwargs.get("messages", args[0] if args else [])

            # Get model from client or kwargs
            model = kwargs.get("model") or getattr(instance, "_model", None) or "unknown"

            # Format input as structured messages
            input_messages = self._format_messages(messages)

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

                # Extract response
                output_messages = self._extract_completion(result)
                tokens_in, tokens_out = self._extract_tokens(result)

                # Get actual model from response
                actual_model = getattr(result, "model", model) or model

                # Build metadata
                metadata = {
                    "response_id": getattr(result, "id", None),
                    "response_model": actual_model,
                }

                # Extract finish reason
                if hasattr(result, "choices") and result.choices:
                    choice = result.choices[0]
                    if hasattr(choice, "finish_reason"):
                        metadata["finish_reason"] = str(choice.finish_reason)

                # Deduplicate
                response_id = metadata.get("response_id")
                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=str(actual_model),
                    provider="azure-foundry",
                    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": str(model), "provider": "azure-foundry"},
                )
                raise

        return wrapper

    def _async_complete_wrapper(self):
        """Wrapper for async ChatCompletionsClient.complete()"""

        async def wrapper(wrapped, instance, args, kwargs):
            messages = kwargs.get("messages", args[0] if args else [])
            model = kwargs.get("model") or getattr(instance, "_model", None) or "unknown"
            input_messages = self._format_messages(messages)

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

                output_messages = self._extract_completion(result)
                tokens_in, tokens_out = self._extract_tokens(result)

                actual_model = getattr(result, "model", model) or model

                metadata = {
                    "response_id": getattr(result, "id", None),
                    "response_model": actual_model,
                }

                if hasattr(result, "choices") and result.choices:
                    choice = result.choices[0]
                    if hasattr(choice, "finish_reason"):
                        metadata["finish_reason"] = str(choice.finish_reason)

                response_id = metadata.get("response_id")
                if not self._should_capture_global(response_id):
                    return result

                self.capture_model_event(
                    input_messages=input_messages,
                    output_messages=output_messages,
                    model_name=str(actual_model),
                    provider="azure-foundry",
                    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": str(model), "provider": "azure-foundry"},
                )
                raise

        return wrapper

    def _embed_wrapper(self):
        """Wrapper for EmbeddingsClient.embed() - emits model call event"""

        def wrapper(wrapped, instance, args, kwargs):
            input_data = kwargs.get("input", args[0] if args else "")
            model = kwargs.get("model") or getattr(instance, "_model", None) or "unknown"

            # Format input
            if isinstance(input_data, list):
                input_content = f"[{len(input_data)} inputs]"
            else:
                input_content = str(input_data)[:1000]

            input_messages = [{"role": "user", "content": input_content}]

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

                # Extract usage
                tokens_in = None
                if hasattr(result, "usage"):
                    tokens_in = getattr(result.usage, "total_tokens", None)
                    if tokens_in is None:
                        tokens_in = getattr(result.usage, "prompt_tokens", None)

                # Format output
                embedding_count = len(result.data) if hasattr(result, "data") else 0
                output_messages = [{"role": "assistant", "content": f"[{embedding_count} embeddings]"}]

                actual_model = getattr(result, "model", model) or model

                self.capture_model_event(
                    input_messages=input_messages,
                    output_messages=output_messages,
                    model_name=str(actual_model),
                    provider="azure-foundry",
                    tokens_in=tokens_in,
                    tokens_out=0,
                    metadata={"operation": "embeddings"},
                )

                return result

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

        return wrapper

    def _async_embed_wrapper(self):
        """Wrapper for async EmbeddingsClient.embed()"""

        async def wrapper(wrapped, instance, args, kwargs):
            input_data = kwargs.get("input", args[0] if args else "")
            model = kwargs.get("model") or getattr(instance, "_model", None) or "unknown"

            if isinstance(input_data, list):
                input_content = f"[{len(input_data)} inputs]"
            else:
                input_content = str(input_data)[:1000]

            input_messages = [{"role": "user", "content": input_content}]

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

                tokens_in = None
                if hasattr(result, "usage"):
                    tokens_in = getattr(result.usage, "total_tokens", None)
                    if tokens_in is None:
                        tokens_in = getattr(result.usage, "prompt_tokens", None)

                embedding_count = len(result.data) if hasattr(result, "data") else 0
                output_messages = [{"role": "assistant", "content": f"[{embedding_count} embeddings]"}]

                actual_model = getattr(result, "model", model) or model

                self.capture_model_event(
                    input_messages=input_messages,
                    output_messages=output_messages,
                    model_name=str(actual_model),
                    provider="azure-foundry",
                    tokens_in=tokens_in,
                    tokens_out=0,
                    metadata={"operation": "embeddings"},
                )

                return result

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

        return wrapper

    # === Helper Methods ===

    @staticmethod
    def _format_messages(messages) -> list:
        """Format Azure AI Inference messages as structured list of dicts"""
        if not messages:
            return []

        formatted = []
        for msg in messages:
            if isinstance(msg, dict):
                role = msg.get("role", "unknown")
                content = msg.get("content", "")

                # Handle content blocks
                if isinstance(content, list):
                    text_parts = []
                    for part in content:
                        if isinstance(part, dict) and part.get("type") == "text":
                            text_parts.append(part.get("text", ""))
                        elif isinstance(part, str):
                            text_parts.append(part)
                    content = " ".join(text_parts)

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

            elif hasattr(msg, "role") and hasattr(msg, "content"):
                # Object-style message
                content = msg.content
                if isinstance(content, list):
                    text_parts = []
                    for part in content:
                        if hasattr(part, "text"):
                            text_parts.append(part.text)
                        elif isinstance(part, str):
                            text_parts.append(part)
                    content = " ".join(text_parts)

                formatted.append({"role": msg.role, "content": content or ""})

        return formatted

    @staticmethod
    def _extract_completion(result) -> list:
        """Extract completion as structured message list from Azure AI Inference response"""
        if result is None:
            return []

        if not hasattr(result, "choices") or not result.choices:
            return []

        choice = result.choices[0]

        if not hasattr(choice, "message"):
            return []

        message = choice.message
        content = getattr(message, "content", "") or ""

        msg_dict = {"role": "assistant", "content": content}

        # Check for tool calls
        if hasattr(message, "tool_calls") and message.tool_calls:
            msg_dict["tool_calls"] = [
                {
                    "id": getattr(tc, "id", None),
                    "name": getattr(tc.function, "name", None) if hasattr(tc, "function") else None,
                    "arguments": getattr(tc.function, "arguments", None) if hasattr(tc, "function") else None,
                }
                for tc in message.tool_calls
            ]

        return [msg_dict]

    @staticmethod
    def _extract_tokens(result) -> tuple:
        """Extract token counts from Azure AI Inference response"""
        tokens_in = None
        tokens_out = None

        if hasattr(result, "usage") and result.usage:
            usage = result.usage
            tokens_in = getattr(usage, "prompt_tokens", None)
            tokens_out = getattr(usage, "completion_tokens", None)

        return tokens_in, tokens_out
