from base_google_genai_proxy_sdk_test import BaseGoogleGenAIProxySDKTest
from base_google_test import BaseGoogleGenAITest

import pytest
import litellm
import unittest.mock
import json


class TestGoogleGenAIStudio(BaseGoogleGenAITest, BaseGoogleGenAIProxySDKTest):
    """Test Google GenAI Studio"""

    @property
    def model_config(self):
        return {
            "model": "gemini/gemini-2.5-flash-lite",
        }

    @property
    def proxy_model_name(self) -> str:
        return "gemini-2.5-flash-lite"


@pytest.mark.asyncio
async def test_mock_stream_generate_content_with_tools():
    """Test streaming function call response parsing and validation"""
    from litellm.types.google_genai.main import ToolConfigDict

    litellm._turn_on_debug()
    contents = [
        {
            "role": "user",
            "parts": [
                {
                    "text": "Schedule a meeting with Bob and Alice for 03/27/2025 at 10:00 AM about the Q3 planning"
                }
            ],
        }
    ]

    # Mock streaming response chunks that represent a function call response
    mock_response_chunk = {
        "candidates": [
            {
                "content": {
                    "parts": [
                        {
                            "functionCall": {
                                "name": "schedule_meeting",
                                "args": {
                                    "attendees": ["Bob", "Alice"],
                                    "date": "2025-03-27",
                                    "time": "10:00",
                                    "topic": "Q3 planning",
                                },
                            }
                        }
                    ],
                    "role": "model",
                },
                "finishReason": "STOP",
                "index": 0,
            }
        ],
        "usageMetadata": {
            "promptTokenCount": 15,
            "candidatesTokenCount": 5,
            "totalTokenCount": 20,
        },
    }

    # Mock the HTTP handler
    with unittest.mock.patch(
        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
        new_callable=unittest.mock.AsyncMock,
    ) as mock_post:
        # Create mock response object
        mock_response = unittest.mock.MagicMock()
        mock_response.status_code = 200
        mock_response.headers = {"content-type": "application/json"}

        # Mock aiter_lines: yield one line at a time (no trailing newlines),
        # with a blank line between events, matching httpx aiter_lines behaviour.
        async def mock_aiter_lines():
            yield f"data: {json.dumps(mock_response_chunk)}"
            yield ""
            yield "data: [DONE]"
            yield ""

        mock_response.aiter_lines = mock_aiter_lines
        mock_post.return_value = mock_response

        print(
            "\n--- Testing async agenerate_content_stream with function call parsing ---"
        )
        response = await litellm.google_genai.agenerate_content_stream(
            model="gemini/gemini-2.5-flash-lite",
            contents=contents,
            tools=[
                {
                    "functionDeclarations": [
                        {
                            "name": "schedule_meeting",
                            "description": "Schedules a meeting with specified attendees at a given time and date.",
                            "parameters": {
                                "type": "object",
                                "properties": {
                                    "attendees": {
                                        "type": "array",
                                        "items": {"type": "string"},
                                        "description": "List of people attending the meeting.",
                                    },
                                    "date": {
                                        "type": "string",
                                        "description": "Date of the meeting (e.g., '2024-07-29')",
                                    },
                                    "time": {
                                        "type": "string",
                                        "description": "Time of the meeting (e.g., '15:00')",
                                    },
                                    "topic": {
                                        "type": "string",
                                        "description": "The subject or topic of the meeting.",
                                    },
                                },
                                "required": ["attendees", "date", "time", "topic"],
                            },
                        }
                    ]
                }
            ],
        )

        # Collect all chunks and parse function calls
        chunks = []
        function_calls = []

        chunk_count = 0
        async for chunk in response:
            chunk_count += 1
            print(f"Received chunk {chunk_count}: {chunk}")
            chunks.append(chunk)

            # Stop after a reasonable number of chunks to prevent infinite loop
            if chunk_count > 10:
                break

            # Parse function calls from byte chunks
            if isinstance(chunk, bytes):
                try:
                    # Decode bytes to string
                    chunk_str = chunk.decode("utf-8")
                    print(f"Decoded chunk: {chunk_str}")

                    # Extract JSON from Server-Sent Events format (data: {...})
                    if chunk_str.startswith("data: ") and not chunk_str.startswith(
                        "data: [DONE]"
                    ):
                        json_str = chunk_str[6:].strip()  # Remove 'data: ' prefix
                        try:
                            parsed_json = json.loads(json_str)
                            print(f"Parsed JSON: {parsed_json}")

                            # Parse function calls from the JSON
                            if "candidates" in parsed_json:
                                for candidate in parsed_json["candidates"]:
                                    if (
                                        "content" in candidate
                                        and "parts" in candidate["content"]
                                    ):
                                        for part in candidate["content"]["parts"]:
                                            if "functionCall" in part:
                                                function_calls.append(
                                                    {
                                                        "name": part["functionCall"][
                                                            "name"
                                                        ],
                                                        "args": part["functionCall"][
                                                            "args"
                                                        ],
                                                    }
                                                )
                                                print(
                                                    f"Found function call: {part['functionCall']}"
                                                )
                        except json.JSONDecodeError as e:
                            print(f"Failed to parse JSON: {e}")
                except UnicodeDecodeError as e:
                    print(f"Failed to decode bytes: {e}")

            # Handle dict responses (in case some chunks are already parsed)
            elif isinstance(chunk, dict):
                # Direct dict response
                if "candidates" in chunk:
                    for candidate in chunk["candidates"]:
                        if "content" in candidate and "parts" in candidate["content"]:
                            for part in candidate["content"]["parts"]:
                                if "functionCall" in part:
                                    function_calls.append(
                                        {
                                            "name": part["functionCall"]["name"],
                                            "args": part["functionCall"]["args"],
                                        }
                                    )

            # Handle object responses with attributes
            elif hasattr(chunk, "candidates") and chunk.candidates:
                for candidate in chunk.candidates:
                    if hasattr(candidate, "content") and candidate.content:
                        if (
                            hasattr(candidate.content, "parts")
                            and candidate.content.parts
                        ):
                            for part in candidate.content.parts:
                                if (
                                    hasattr(part, "function_call")
                                    and part.function_call
                                ):
                                    function_calls.append(
                                        {
                                            "name": part.function_call.name,
                                            "args": part.function_call.args,
                                        }
                                    )

        # Assertions
        print(f"\nFunction calls found: {function_calls}")
        print(f"Total chunks received: {chunk_count}")

        # Assert we found at least one function call
        assert (
            len(function_calls) > 0
        ), "Expected at least one function call in the streaming response"

        # Check the first function call
        function_call = function_calls[0]

        # Assert function name
        assert (
            function_call["name"] == "schedule_meeting"
        ), f"Expected function name 'schedule_meeting', got '{function_call['name']}'"

        # Assert function arguments
        args = function_call["args"]
        assert "attendees" in args, "Expected 'attendees' in function call arguments"
        assert "date" in args, "Expected 'date' in function call arguments"
        assert "time" in args, "Expected 'time' in function call arguments"
        assert "topic" in args, "Expected 'topic' in function call arguments"

        # Assert specific argument values
        assert args["attendees"] == [
            "Bob",
            "Alice",
        ], f"Expected attendees ['Bob', 'Alice'], got {args['attendees']}"
        assert (
            args["date"] == "2025-03-27"
        ), f"Expected date '2025-03-27', got {args['date']}"
        assert args["time"] == "10:00", f"Expected time '10:00', got {args['time']}"
        assert (
            args["topic"] == "Q3 planning"
        ), f"Expected topic 'Q3 planning', got {args['topic']}"

        print("✅ All function call assertions passed!")


@pytest.mark.asyncio
async def test_validate_post_request_parameters():
    """
    Test that the correct parameters are sent in the POST request to Google GenAI API

    Params validated
        1. model
        2. contents
        3. tools
    """
    from litellm.types.google_genai.main import ToolConfigDict

    contents = [
        {
            "role": "user",
            "parts": [
                {
                    "text": "Schedule a meeting with Bob and Alice for 03/27/2025 at 10:00 AM about the Q3 planning"
                }
            ],
        }
    ]

    tools = [
        {
            "functionDeclarations": [
                {
                    "name": "schedule_meeting",
                    "description": "Schedules a meeting with specified attendees at a given time and date.",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "attendees": {
                                "type": "array",
                                "items": {"type": "string"},
                                "description": "List of people attending the meeting.",
                            },
                            "date": {
                                "type": "string",
                                "description": "Date of the meeting (e.g., '2024-07-29')",
                            },
                            "time": {
                                "type": "string",
                                "description": "Time of the meeting (e.g., '15:00')",
                            },
                            "topic": {
                                "type": "string",
                                "description": "The subject or topic of the meeting.",
                            },
                        },
                        "required": ["attendees", "date", "time", "topic"],
                    },
                }
            ]
        }
    ]

    # Mock the HTTP handler to capture the request
    with unittest.mock.patch(
        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
        new_callable=unittest.mock.AsyncMock,
    ) as mock_post:
        # Create mock response object
        mock_response = unittest.mock.MagicMock()
        mock_response.status_code = 200
        mock_response.headers = {"content-type": "application/json"}

        # Mock aiter_lines: yield one line at a time (no trailing newlines),
        # with a blank line between events, matching httpx aiter_lines behaviour.
        async def mock_aiter_lines():
            yield "data: [DONE]"
            yield ""

        mock_response.aiter_lines = mock_aiter_lines
        mock_post.return_value = mock_response

        print("\n--- Testing POST request parameters validation ---")

        # Make the API call
        response = await litellm.google_genai.agenerate_content_stream(
            model="gemini/gemini-2.5-flash-lite", contents=contents, tools=tools
        )

        # Consume the response to ensure the request is made
        async for chunk in response:
            pass

        # Validate that the HTTP post was called
        assert mock_post.called, "Expected HTTP POST to be called"

        # Get the call arguments
        call_args, call_kwargs = mock_post.call_args

        print(f"POST call args: {call_args}")
        print(f"POST call kwargs: {call_kwargs}")

        # Validate URL contains the correct endpoint
        if call_args:
            url = call_args[0] if len(call_args) > 0 else call_kwargs.get("url")
            assert url is not None, "Expected URL to be provided"
            assert (
                "generativelanguage.googleapis.com" in url
            ), f"Expected Google API URL, got: {url}"
            assert (
                "streamGenerateContent" in url
            ), f"Expected streamGenerateContent endpoint, got: {url}"
            print(f"✅ URL validation passed: {url}")

        # Get the request data/json from the call
        request_data = None
        if "data" in call_kwargs:
            # If data is passed as bytes, decode it
            if isinstance(call_kwargs["data"], bytes):
                request_data = json.loads(call_kwargs["data"].decode("utf-8"))
            else:
                request_data = call_kwargs["data"]
        elif "json" in call_kwargs:
            request_data = call_kwargs["json"]

        assert request_data is not None, "Expected request data to be provided"
        print(f"Request data: {json.dumps(request_data, indent=2)}")

        # Validate model field
        assert "model" in request_data, "Expected 'model' field in request data"
        # Model might be transformed, but should contain gemini-2.5-flash-lite
        model_value = request_data["model"]
        assert (
            "gemini-2.5-flash-lite" in model_value
        ), f"Expected model to contain 'gemini-2.5-flash-lite', got: {model_value}"
        print(f"✅ Model validation passed: {model_value}")

        # Validate contents field
        assert "contents" in request_data, "Expected 'contents' field in request data"
        request_contents = request_data["contents"]
        assert isinstance(request_contents, list), "Expected contents to be a list"
        assert len(request_contents) > 0, "Expected at least one content item"

        # Check the first content item
        first_content = request_contents[0]
        assert "role" in first_content, "Expected 'role' in content item"
        assert (
            first_content["role"] == "user"
        ), f"Expected role 'user', got: {first_content['role']}"
        assert "parts" in first_content, "Expected 'parts' in content item"
        assert isinstance(first_content["parts"], list), "Expected parts to be a list"
        assert len(first_content["parts"]) > 0, "Expected at least one part"

        # Check the text content
        first_part = first_content["parts"][0]
        assert "text" in first_part, "Expected 'text' in part"
        expected_text = "Schedule a meeting with Bob and Alice for 03/27/2025 at 10:00 AM about the Q3 planning"
        assert (
            first_part["text"] == expected_text
        ), f"Expected text '{expected_text}', got: {first_part['text']}"
        print(f"✅ Contents validation passed")

        # Validate tools field
        assert "tools" in request_data, "Expected 'tools' field in request data"
        request_tools = request_data["tools"]
        assert isinstance(request_tools, list), "Expected tools to be a list"
        assert len(request_tools) > 0, "Expected at least one tool"

        # Check the first tool
        first_tool = request_tools[0]
        assert (
            "functionDeclarations" in first_tool
        ), "Expected 'functionDeclarations' in tool"
        function_declarations = first_tool["functionDeclarations"]
        assert isinstance(
            function_declarations, list
        ), "Expected functionDeclarations to be a list"
        assert (
            len(function_declarations) > 0
        ), "Expected at least one function declaration"

        # Check the function declaration
        func_decl = function_declarations[0]
        assert "name" in func_decl, "Expected 'name' in function declaration"
        assert (
            func_decl["name"] == "schedule_meeting"
        ), f"Expected function name 'schedule_meeting', got: {func_decl['name']}"
        assert (
            "description" in func_decl
        ), "Expected 'description' in function declaration"
        assert (
            "parameters" in func_decl
        ), "Expected 'parameters' in function declaration"

        # Check function parameters
        params = func_decl["parameters"]
        assert "type" in params, "Expected 'type' in parameters"
        assert (
            params["type"] == "object"
        ), f"Expected parameters type 'object', got: {params['type']}"
        assert "properties" in params, "Expected 'properties' in parameters"
        assert "required" in params, "Expected 'required' in parameters"

        # Check required fields
        required_fields = params["required"]
        expected_required = ["attendees", "date", "time", "topic"]
        assert set(required_fields) == set(
            expected_required
        ), f"Expected required fields {expected_required}, got: {required_fields}"
        print(f"✅ Tools validation passed")

        print("✅ All POST request parameter validations passed!")
