import os
from collections.abc import AsyncIterator
from typing import Any

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

from tests.fixtures.agents.prompt import SYSTEM_PROMPT
from tests.fixtures.agents.utils import try_execute_tool
from tests.fixtures.datasource import datasource


def _get_first_candidate_content(chunk: Any) -> Any:
    if (
        hasattr(chunk, "candidates")
        and chunk.candidates
        and hasattr(chunk.candidates[0], "content")
        and chunk.candidates[0].content
    ):
        return chunk.candidates[0].content
    return None


def _get_first_candidate_content_text(chunk: Any) -> str:
    content = _get_first_candidate_content(chunk)
    if content and hasattr(content, "parts") and content.parts:
        return content.parts[0].text

    return ""


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

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

    def __init__(self, api_key: str | None = None):
        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 ask(self, question: str) -> str:
        contents = [
            genai_types.Content(role="user", parts=[genai_types.Part(text=question)])
        ]

        for iteration in range(self.MAX_ITERATIONS):
            mode = "ANY" if iteration == 0 else "AUTO"
            config = self._create_config(mode)

            response = self.client.models.generate_content(
                model=self.model_name,
                contents=contents,
                config=config,
            )

            if not (first_candidate_content := _get_first_candidate_content(response)):
                break

            function_calls, text_parts = self._extract_parts(first_candidate_content)

            if not function_calls or iteration == self.MAX_ITERATIONS - 1:
                return _get_first_candidate_content_text(response)

            model_parts = self._build_model_parts(text_parts, function_calls)
            contents.append(genai_types.Content(role="model", parts=model_parts))

            function_response_parts = self._execute_function_calls(function_calls)
            contents.append(
                genai_types.Content(role="user", parts=function_response_parts)
            )

        return _get_first_candidate_content_text(response)

    async def stream(self, question: str) -> AsyncIterator[str]:
        contents = [
            genai_types.Content(role="user", parts=[genai_types.Part(text=question)])
        ]

        for iteration in range(self.MAX_ITERATIONS):
            mode = "ANY" if iteration == 0 else "AUTO"
            config = self._create_config(mode)

            stream = self.client.models.generate_content_stream(
                model=self.model_name,
                contents=contents,
                config=config,
            )

            full_response = None

            for chunk in stream:
                if content := _get_first_candidate_content(chunk):
                    if hasattr(content, "parts") and content.parts:
                        for part in content.parts:
                            if hasattr(part, "text") and part.text:
                                yield part.text
                    full_response = chunk

            if not full_response or not full_response.candidates:
                break

            content = _get_first_candidate_content(full_response)
            if not content:
                break

            function_calls, text_parts = self._extract_parts(content)
            if not function_calls or iteration == self.MAX_ITERATIONS - 1:
                break

            model_parts = self._build_model_parts(text_parts, function_calls)
            contents.append(genai_types.Content(role="model", parts=model_parts))

            function_response_parts = self._execute_function_calls(function_calls)
            contents.append(
                genai_types.Content(role="user", parts=function_response_parts)
            )

    def _build_model_parts(
        self, text_parts: list[str], function_calls: list
    ) -> list[genai_types.Part]:
        model_parts = []
        if text_parts:
            model_parts.extend([genai_types.Part(text=text) for text in text_parts])
        model_parts.extend(
            [genai_types.Part(function_call=fc) for fc in function_calls]
        )
        return model_parts

    def _create_config(self, mode: str) -> genai_types.GenerateContentConfig:
        return genai_types.GenerateContentConfig(
            max_output_tokens=500,
            temperature=0.7,
            tools=self.tools,
            system_instruction=SYSTEM_PROMPT,
            tool_config=genai_types.ToolConfig(
                function_calling_config=genai_types.FunctionCallingConfig(mode=mode)
            ),
        )

    def _create_function_map(self) -> dict[str, callable]:
        return {func.__name__: func for func in datasource}

    def _create_tools(self) -> list[genai_types.Tool]:
        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 _execute_function_call(
        self, func_call: genai_types.FunctionCall
    ) -> genai_types.Part:
        func = self.function_map.get(func_call.name)

        if not func:
            return genai_types.Part(
                function_response=genai_types.FunctionResponse(
                    name=func_call.name,
                    response={"error": f"Function '{func_call.name}' not found"},
                )
            )

        result, error_msg = try_execute_tool(func, func_call.args)

        if error_msg:
            response_data = {"error": error_msg}
        elif isinstance(result, dict):
            response_data = result
        else:
            response_data = {"result": result}

        return genai_types.Part(
            function_response=genai_types.FunctionResponse(
                name=func_call.name,
                response=response_data,
            )
        )

    def _execute_function_calls(self, function_calls: list) -> list[genai_types.Part]:
        return [self._execute_function_call(fc) for fc in function_calls]

    def _extract_parts(self, content: Any) -> tuple[list, list]:
        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
