"""
Google GenAI Instrumentation - Event-Driven Capture

Instruments Google's unified GenAI SDK (google-genai) to automatically capture
Gemini model calls as events for detection and analysis.

Supports both:
- Google AI Studio (Gemini Developer API)
- Vertex AI

Both use the same client.models.generate_content() API.
"""
from typing import Collection, Optional
import logging

from wrapt import wrap_function_wrapper

from evoke.wrappers.base import BaseInstrumentor

logger = logging.getLogger(__name__)


class EvokeGoogleInstrumentor(BaseInstrumentor):
    """Instrumentor for Google GenAI SDK - emits events for Gemini model calls"""

    def instrumentation_dependencies(self) -> Collection[str]:
        return ["google-genai >= 1.0.0"]

    def _instrument(self, **kwargs):
        """Instrument Google GenAI SDK"""
        # Wrap sync generate_content on the Models API client
        try:
            wrap_function_wrapper(
                "google.genai.models",
                "Models.generate_content",
                self._generate_content_wrapper(),
            )
        except Exception as e:
            logger.debug(f"Could not instrument google.genai.models.Models.generate_content: {e}")

        # Wrap async generate_content
        try:
            wrap_function_wrapper(
                "google.genai.models",
                "AsyncModels.generate_content",
                self._async_generate_content_wrapper(),
            )
        except Exception as e:
            logger.debug(f"Could not instrument google.genai.models.AsyncModels.generate_content: {e}")

        logger.debug("Google GenAI SDK instrumented")

    def _uninstrument(self, **kwargs):
        """Uninstrument Google GenAI SDK"""
        try:
            import google.genai.models

            # Restore original methods
            if hasattr(google.genai.models.Models.generate_content, "__wrapped__"):
                google.genai.models.Models.generate_content = (
                    google.genai.models.Models.generate_content.__wrapped__
                )
            if hasattr(google.genai.models.AsyncModels.generate_content, "__wrapped__"):
                google.genai.models.AsyncModels.generate_content = (
                    google.genai.models.AsyncModels.generate_content.__wrapped__
                )
        except Exception as e:
            logger.warning(f"Error uninstrumenting Google GenAI: {e}")

    def _generate_content_wrapper(self):
        """Wrapper for models.generate_content() - emits model call event"""

        def wrapper(wrapped, instance, args, kwargs):
            # Extract request info
            model = kwargs.get("model", args[0] if args else "unknown")
            contents = kwargs.get("contents", args[1] if len(args) > 1 else "")
            config = kwargs.get("config", {})

            # Format input as structured messages
            input_messages = self._format_contents(contents, config)

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

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

                # Build metadata
                metadata = {
                    "response_model": getattr(result, "model", model) if hasattr(result, "model") else model,
                }

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

                # Check for function calls
                function_calls = self._extract_function_calls(result)
                if function_calls:
                    metadata["function_calls"] = function_calls

                # Deduplicate - skip if already captured by another wrapper (e.g., LiteLLM)
                # Google responses don't have a standard ID, use a hash of the output
                response_text = output_messages[0].get("content", "") if output_messages else ""
                response_id = f"google-{hash(response_text)}" if response_text else None
                if not self._should_capture_global(response_id):
                    return result

                # Emit event with structured messages
                self.capture_model_event(
                    input_messages=input_messages,
                    output_messages=output_messages,
                    model_name=str(model),
                    provider="google",
                    tokens_in=tokens_in,
                    tokens_out=tokens_out,
                    metadata=metadata,
                )

                return result

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

        return wrapper

    def _async_generate_content_wrapper(self):
        """Wrapper for async models.generate_content()"""

        async def wrapper(wrapped, instance, args, kwargs):
            model = kwargs.get("model", args[0] if args else "unknown")
            contents = kwargs.get("contents", args[1] if len(args) > 1 else "")
            config = kwargs.get("config", {})
            input_messages = self._format_contents(contents, config)

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

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

                metadata = {
                    "response_model": getattr(result, "model", model) if hasattr(result, "model") else model,
                }

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

                function_calls = self._extract_function_calls(result)
                if function_calls:
                    metadata["function_calls"] = function_calls

                # Deduplicate
                response_text = output_messages[0].get("content", "") if output_messages else ""
                response_id = f"google-{hash(response_text)}" if response_text else None
                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(model),
                    provider="google",
                    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": "google"},
                )
                raise

        return wrapper

    # === Helper Methods ===

    @staticmethod
    def _format_contents(contents, config=None) -> list:
        """
        Format Google GenAI contents as structured list of message dicts.

        Contents can be:
        - A string (simple prompt)
        - A list of strings
        - A list of content dicts with role and parts
        """
        formatted = []

        # Handle system instruction from config
        if config:
            system_instruction = None
            if isinstance(config, dict):
                system_instruction = config.get("system_instruction")
            elif hasattr(config, "system_instruction"):
                system_instruction = config.system_instruction

            if system_instruction:
                if isinstance(system_instruction, str):
                    formatted.append({"role": "system", "content": system_instruction})
                elif hasattr(system_instruction, "text"):
                    formatted.append({"role": "system", "content": system_instruction.text})

        # Handle contents
        if isinstance(contents, str):
            # Simple string prompt
            formatted.append({"role": "user", "content": contents})

        elif isinstance(contents, list):
            for item in contents:
                if isinstance(item, str):
                    # List of strings
                    formatted.append({"role": "user", "content": item})

                elif isinstance(item, dict):
                    # Content dict with role and parts
                    role = item.get("role", "user")
                    # Map Google's "model" role to "assistant"
                    if role == "model":
                        role = "assistant"

                    parts = item.get("parts", [])
                    text_parts = []
                    for part in parts:
                        if isinstance(part, dict) and "text" in part:
                            text_parts.append(part["text"])
                        elif isinstance(part, str):
                            text_parts.append(part)
                        elif hasattr(part, "text"):
                            text_parts.append(part.text)

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

                elif hasattr(item, "role"):
                    # Object with role attribute
                    role = item.role
                    if role == "model":
                        role = "assistant"

                    text_parts = []
                    if hasattr(item, "parts"):
                        for part in item.parts:
                            if hasattr(part, "text"):
                                text_parts.append(part.text)
                            elif isinstance(part, str):
                                text_parts.append(part)

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

        elif hasattr(contents, "text"):
            # Single Part object
            formatted.append({"role": "user", "content": contents.text})

        return formatted

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

        # Try to get text directly
        text = ""
        if hasattr(result, "text"):
            text = result.text or ""
        elif hasattr(result, "candidates") and result.candidates:
            candidate = result.candidates[0]
            if hasattr(candidate, "content") and hasattr(candidate.content, "parts"):
                text_parts = []
                for part in candidate.content.parts:
                    if hasattr(part, "text"):
                        text_parts.append(part.text)
                text = " ".join(text_parts)

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

        # Check for function calls in the response
        if hasattr(result, "candidates") and result.candidates:
            candidate = result.candidates[0]
            if hasattr(candidate, "content") and hasattr(candidate.content, "parts"):
                function_calls = []
                for part in candidate.content.parts:
                    if hasattr(part, "function_call"):
                        fc = part.function_call
                        function_calls.append({
                            "name": getattr(fc, "name", None),
                            "arguments": str(getattr(fc, "args", {})),
                        })
                if function_calls:
                    msg_dict["tool_calls"] = function_calls

        return [msg_dict]

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

        if hasattr(result, "usage_metadata") and result.usage_metadata:
            usage = result.usage_metadata
            tokens_in = getattr(usage, "prompt_token_count", None)
            tokens_out = getattr(usage, "candidates_token_count", None)

            # Fallback attribute names
            if tokens_in is None:
                tokens_in = getattr(usage, "input_tokens", None)
            if tokens_out is None:
                tokens_out = getattr(usage, "output_tokens", None)

        return tokens_in, tokens_out

    @staticmethod
    def _extract_function_calls(result) -> list:
        """Extract function calls from Google response"""
        if not hasattr(result, "candidates") or not result.candidates:
            return []

        candidate = result.candidates[0]
        if not hasattr(candidate, "content") or not hasattr(candidate.content, "parts"):
            return []

        function_calls = []
        for part in candidate.content.parts:
            if hasattr(part, "function_call"):
                fc = part.function_call
                function_calls.append({
                    "name": getattr(fc, "name", None),
                    "args": str(getattr(fc, "args", {}))[:500],  # Limit size
                })

        return function_calls
