import json
import os
from typing import Any

from google import genai
from google.genai import types as genai_types

from tests.fixtures.datasource import datasource


class GoogleAgent:
    """Agent for interacting with Google's Gemini API with function calling support."""

    DEFAULT_MODEL = "gemini-2.0-flash-exp"
    MAX_ITERATIONS = 3

    def __init__(self, api_key: str | None = None):
        """
        Initialize the Google Agent.

        Args:
            api_key: Google API key. If not provided, uses GOOGLE_API_KEY env variable.
        """
        self.client = genai.Client(api_key=api_key or os.getenv("GOOGLE_API_KEY"))
        self.model_name = self.DEFAULT_MODEL
        self.tools = self._create_tools()
        self.function_map = self._create_function_map()

    def _create_tools(self) -> list[genai_types.Tool]:
        """Convert datasource functions to Google GenAI tools."""
        function_declarations = [
            genai_types.FunctionDeclaration.from_callable(
                callable=func, client=self.client
            )
            for func in datasource
        ]
        return [genai_types.Tool(function_declarations=function_declarations)]

    def _create_function_map(self) -> dict[str, callable]:
        """Create a mapping of function names to callable functions."""
        return {func.__name__: func for func in datasource}

    def _execute_function_call(
        self, func_call: genai_types.FunctionCall
    ) -> genai_types.Part:
        """
        Execute a single function call and return the response part.

        Args:
            func_call: The function call to execute

        Returns:
            Part containing the function response
        """
        func_name = func_call.name
        func_args = func_call.args

        if isinstance(func_args, str):
            try:
                func_args = json.loads(func_args)
            except json.JSONDecodeError:
                func_args = {}

        if func_name not in self.function_map:
            return genai_types.Part(
                function_response=genai_types.FunctionResponse(
                    name=func_name,
                    response={"error": f"Function '{func_name}' not found"},
                )
            )

        try:
            result = (
                self.function_map[func_name](**func_args)
                if isinstance(func_args, dict)
                else self.function_map[func_name](func_args)
            )
            return genai_types.Part(
                function_response=genai_types.FunctionResponse(
                    name=func_name,
                    response=result,
                )
            )
        except Exception as e:
            return genai_types.Part(
                function_response=genai_types.FunctionResponse(
                    name=func_name,
                    response={"error": str(e)},
                )
            )

    def _extract_parts(self, content: Any) -> tuple[list, list]:
        """
        Extract function calls and text from response content.

        Returns:
            Tuple of (function_calls, text_parts)
        """
        parts = content.parts if hasattr(content, "parts") else []
        function_calls = []
        text_parts = []

        for part in parts:
            if hasattr(part, "function_call") and part.function_call:
                function_calls.append(part.function_call)
            elif hasattr(part, "text") and part.text:
                text_parts.append(part.text)

        return function_calls, text_parts

    def ask(self, question: str) -> list[Any]:
        """
        Ask a question and get a response, handling function calls automatically.

        Args:
            question: The question to ask

        Returns:
            List of response candidates
        """
        contents = [
            {
                "role": "user",
                "parts": [{"text": question}],
            }
        ]

        config = {
            "max_output_tokens": 500,
            "temperature": 0.7,
            "tools": self.tools,
            "tool_config": genai_types.ToolConfig(
                function_calling_config=genai_types.FunctionCallingConfig(mode="ANY")
            ),
        }

        for _ in range(self.MAX_ITERATIONS):
            response = self.client.models.generate_content(
                model=self.model_name,
                contents=contents,
                config=config,
            )

            if not response.candidates or not response.candidates[0].content:
                break

            content = response.candidates[0].content
            function_calls = self._extract_parts(content)

            if not function_calls:
                return response.candidates

            contents.append(
                {
                    "role": "model",
                    "parts": [
                        genai_types.Part(function_call=fc) for fc in function_calls
                    ],
                }
            )

            for func_call in function_calls:
                function_response = self._execute_function_call(func_call)
                contents.append(
                    {
                        "role": "user",
                        "parts": [function_response],
                    }
                )

        return response.candidates

    async def stream(self, question: str):
        """
        Stream responses from Google GenAI with tool call support.

        Args:
            question: The question to ask

        Yields:
            Text chunks from the streaming response
        """
        contents = [
            {
                "role": "user",
                "parts": [{"text": question}],
            }
        ]

        config = {
            "max_output_tokens": 500,
            "temperature": 0.7,
            "tools": self.tools,
            "tool_config": genai_types.ToolConfig(
                function_calling_config=genai_types.FunctionCallingConfig(mode="ANY")
            ),
        }

        for _ in range(self.MAX_ITERATIONS):
            stream = self.client.models.generate_content_stream(
                model=self.model_name,
                contents=contents,
                config=config,
            )

            full_response = None
            accumulated_text = []
            function_calls = []

            for chunk in stream:
                if (
                    hasattr(chunk, "candidates")
                    and chunk.candidates
                    and hasattr(chunk.candidates[0], "content")
                    and chunk.candidates[0].content
                ):
                    candidate = chunk.candidates[0]
                    content = candidate.content
                    if hasattr(content, "parts") and content.parts:
                        for part in content.parts:
                            if hasattr(part, "text") and part.text:
                                accumulated_text.append(part.text)
                                yield part.text
                            elif hasattr(part, "function_call") and part.function_call:
                                function_calls.append(part.function_call)
                    full_response = chunk

            if (
                function_calls
                and full_response
                and hasattr(full_response, "candidates")
                and full_response.candidates
                and hasattr(full_response.candidates[0], "content")
                and full_response.candidates[0].content
            ):
                candidate = full_response.candidates[0]
                content = candidate.content
                extracted_calls, _ = self._extract_parts(content)
                if extracted_calls:
                    contents.append(
                        {
                            "role": "model",
                            "parts": [
                                genai_types.Part(function_call=fc)
                                for fc in extracted_calls
                            ],
                        }
                    )

                    for func_call in extracted_calls:
                        function_response = self._execute_function_call(func_call)
                        contents.append(
                            {
                                "role": "user",
                                "parts": [function_response],
                            }
                        )
                    continue

            break
