#!/usr/bin/env python3
"""
Unit tests for the Antigravity gRPC fallback module.

Tests cover:
1. Module import and availability detection
2. Protobuf conversion helpers (JSON <-> protobuf)
3. Request building from wrapped REST dict
4. Reverse alias map correctness
5. GrpcFallbackResult type
6. Integration: _try_grpc_fallback triggers correctly on REST 404
"""

import json
import os
import sys
import unittest
from unittest.mock import patch, MagicMock

# Add src to path so we can import the antigravity_grpc package
_src_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "src")
if _src_dir not in sys.path:
    sys.path.insert(0, _src_dir)


class TestGrpcModuleAvailability(unittest.TestCase):
    """Tests for is_grpc_available() and module loading."""

    def test_is_grpc_available_returns_bool(self):
        """is_grpc_available should return a boolean."""
        from antigravity_grpc import is_grpc_available
        result = is_grpc_available()
        self.assertIsInstance(result, bool)

    def test_is_grpc_available_true_when_installed(self):
        """If grpcio is installed and stubs are loadable, should return True."""
        from antigravity_grpc import is_grpc_available
        # grpcio was installed at test time, so this should be True
        self.assertTrue(is_grpc_available())

    def test_client_instantiation(self):
        """AntigravityGrpcClient should be instantiatable."""
        from antigravity_grpc import AntigravityGrpcClient
        client = AntigravityGrpcClient()
        self.assertIsNotNone(client)

    def test_get_client_singleton(self):
        """get_client should return the same singleton."""
        from antigravity_grpc import get_client
        c1 = get_client()
        c2 = get_client()
        self.assertIs(c1, c2)


class TestGrpcFallbackResult(unittest.TestCase):
    """Tests for GrpcFallbackResult type."""

    def test_default_values(self):
        from antigravity_grpc import GrpcFallbackResult
        r = GrpcFallbackResult()
        self.assertFalse(r.ok)
        self.assertIsNone(r.response_data)
        self.assertIsNone(r.stream_chunks)
        self.assertEqual(r.error_message, "")
        self.assertEqual(r.endpoint_used, "")
        self.assertEqual(r.model_used, "")
        self.assertEqual(r.elapsed_s, 0.0)

    def test_success_result(self):
        from antigravity_grpc import GrpcFallbackResult
        r = GrpcFallbackResult(ok=True, response_data={"response": {"candidates": []}},
                                endpoint_used="daily-cloudcode-pa.googleapis.com:443",
                                model_used="Gemini 3.5 Flash (High)",
                                elapsed_s=2.5)
        self.assertTrue(r.ok)
        self.assertIsNotNone(r.response_data)
        self.assertEqual(r.elapsed_s, 2.5)

    def test_failure_result(self):
        from antigravity_grpc import GrpcFallbackResult
        r = GrpcFallbackResult(ok=False, error_message="All gRPC endpoints failed")
        self.assertFalse(r.ok)
        self.assertIn("failed", r.error_message)

    def test_repr(self):
        from antigravity_grpc import GrpcFallbackResult
        r_ok = GrpcFallbackResult(ok=True, response_data={"response": {"candidates": []}})
        self.assertIn("OK", repr(r_ok))
        r_fail = GrpcFallbackResult(ok=False, error_message="timeout")
        self.assertIn("FAIL", repr(r_fail))


class TestReverseAliasMap(unittest.TestCase):
    """Tests for the _GRPC_REVERSE_ALIAS map in translate-proxy.py."""

    def test_import_reverse_alias(self):
        """The reverse alias map should be importable from the proxy module."""
        import importlib
        _spec = importlib.util.spec_from_file_location(
            "translate_proxy",
            os.path.join(_src_dir, "translate-proxy.py"),
        )
        tp = importlib.util.module_from_spec(_spec)
        _spec.loader.exec_module(tp)
        self.assertIsInstance(tp._GRPC_REVERSE_ALIAS, dict)

    def test_key_models_have_reverse_aliases(self):
        """All key REST model slugs should have gRPC display name mappings."""
        import importlib
        _spec = importlib.util.spec_from_file_location(
            "translate_proxy",
            os.path.join(_src_dir, "translate-proxy.py"),
        )
        tp = importlib.util.module_from_spec(_spec)
        _spec.loader.exec_module(tp)

        required_slugs = [
            "gemini-3-flash",
            "gemini-3.5-flash-low",
            "gemini-3.1-pro-low",
            "claude-sonnet-4-6",
            "claude-opus-4-6-thinking",
            "gemini-2.5-flash",
        ]
        for slug in required_slugs:
            self.assertIn(slug, tp._GRPC_REVERSE_ALIAS,
                         f"Missing reverse alias for REST slug '{slug}'")

    def test_reverse_alias_values_are_display_names(self):
        """gRPC display names should contain spaces and parentheses, not hyphens."""
        import importlib
        _spec = importlib.util.spec_from_file_location(
            "translate_proxy",
            os.path.join(_src_dir, "translate-proxy.py"),
        )
        tp = importlib.util.module_from_spec(_spec)
        _spec.loader.exec_module(tp)

        for slug, display_name in tp._GRPC_REVERSE_ALIAS.items():
            # Display names typically have spaces (e.g. "Gemini 3.5 Flash (High)")
            # while slugs use hyphens (e.g. "gemini-3-flash")
            self.assertNotEqual(slug, display_name,
                               f"Reverse alias for '{slug}' should differ from slug (gRPC uses display names)")


class TestProtobufConversion(unittest.TestCase):
    """Tests for JSON -> protobuf conversion helpers."""

    def test_struct_to_protobuf(self):
        """_struct_to_protobuf should convert a simple dict to Struct."""
        from antigravity_grpc.client import _struct_to_protobuf
        result = _struct_to_protobuf({"key": "value", "num": 42})
        self.assertIsNotNone(result)
        # Verify round-trip
        from antigravity_grpc.client import _protobuf_struct_to_dict
        d = _protobuf_struct_to_dict(result)
        self.assertEqual(d["key"], "value")
        self.assertEqual(d["num"], 42.0)

    def test_struct_round_trip_nested(self):
        """Nested dicts should survive a round-trip through protobuf."""
        from antigravity_grpc.client import _struct_to_protobuf, _protobuf_struct_to_dict
        original = {"outer": {"inner": "hello"}, "list_val": [1, 2, 3]}
        proto = _struct_to_protobuf(original)
        result = _protobuf_struct_to_dict(proto)
        self.assertEqual(result["outer"]["inner"], "hello")
        self.assertEqual(result["list_val"], [1.0, 2.0, 3.0])

    def test_json_parts_to_proto_text(self):
        """Text parts should convert to protobuf Part with text field."""
        from antigravity_grpc.client import _json_parts_to_proto
        parts = _json_parts_to_proto([{"text": "Hello world"}])
        self.assertEqual(len(parts), 1)
        self.assertEqual(parts[0].text, "Hello world")

    def test_json_parts_to_proto_function_call(self):
        """FunctionCall parts should convert correctly."""
        from antigravity_grpc.client import _json_parts_to_proto
        parts = _json_parts_to_proto([{
            "functionCall": {
                "name": "exec_command",
                "args": {"cmd": "ls -la"},
                "id": "call_123"
            }
        }])
        self.assertEqual(len(parts), 1)
        self.assertTrue(parts[0].HasField("function_call"))
        self.assertEqual(parts[0].function_call.name, "exec_command")
        self.assertEqual(parts[0].function_call.id, "call_123")

    def test_json_parts_to_proto_function_response(self):
        """FunctionResponse parts should convert correctly."""
        from antigravity_grpc.client import _json_parts_to_proto
        parts = _json_parts_to_proto([{
            "functionResponse": {
                "name": "exec_command",
                "response": {"result": "file1.txt"},
                "id": "call_123"
            }
        }])
        self.assertEqual(len(parts), 1)
        self.assertTrue(parts[0].HasField("function_response"))
        self.assertEqual(parts[0].function_response.name, "exec_command")

    def test_json_contents_to_proto(self):
        """Content objects should convert correctly."""
        from antigravity_grpc.client import _json_contents_to_proto
        contents = _json_contents_to_proto([
            {"role": "user", "parts": [{"text": "Hello"}]},
            {"role": "model", "parts": [{"text": "Hi there"}]},
        ])
        self.assertEqual(len(contents), 2)
        self.assertEqual(contents[0].role, "user")
        self.assertEqual(contents[1].role, "model")

    def test_proto_candidate_to_json(self):
        """Protobuf candidates should convert back to JSON-compatible dicts."""
        from antigravity_grpc.client import _json_contents_to_proto, _proto_candidate_to_json
        from antigravity_grpc import cloudcode_pb2 as pb2

        # Build a candidate manually
        candidate = pb2.Candidate()
        candidate.content.role = "model"
        candidate.content.parts.add().text = "Hello from gRPC"
        candidate.finish_reason = "STOP"
        candidate.index = 0

        result = _proto_candidate_to_json(candidate)
        self.assertEqual(result["finishReason"], "STOP")
        self.assertEqual(result["content"]["role"], "model")
        self.assertEqual(result["content"]["parts"][0]["text"], "Hello from gRPC")


class TestGrpcRequestBuilding(unittest.TestCase):
    """Tests for _build_request (wrapped REST dict → protobuf)."""

    def _get_client(self):
        from antigravity_grpc import AntigravityGrpcClient
        return AntigravityGrpcClient()

    def test_build_request_basic(self):
        """Basic request fields should be populated correctly."""
        client = self._get_client()
        wrapped = {
            "project": "test-project-123",
            "model": "Gemini 3.5 Flash (High)",
            "requestType": "agent",
            "userAgent": "antigravity/2.0.6",
            "requestId": "agent-test123",
            "request": {
                "contents": [
                    {"role": "user", "parts": [{"text": "Say hello"}]}
                ],
                "safetySettings": [
                    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"},
                ],
            }
        }
        req = client._build_request(wrapped)
        self.assertEqual(req.project, "test-project-123")
        self.assertEqual(req.model, "Gemini 3.5 Flash (High)")
        self.assertEqual(req.request_type, "agent")
        self.assertEqual(len(req.request.contents), 1)
        self.assertEqual(req.request.contents[0].role, "user")

    def test_build_request_with_tools(self):
        """Tools should be converted to function declarations."""
        client = self._get_client()
        wrapped = {
            "project": "test-project",
            "model": "gemini-3-flash",
            "request": {
                "contents": [],
                "tools": [{
                    "functionDeclarations": [{
                        "name": "exec_command",
                        "description": "Run a shell command",
                        "parameters": {"type": "object", "properties": {"cmd": {"type": "string"}}}
                    }]
                }],
            }
        }
        req = client._build_request(wrapped)
        self.assertEqual(len(req.request.tools), 1)
        self.assertEqual(req.request.tools[0].function_declarations[0].name, "exec_command")

    def test_build_request_with_generation_config(self):
        """Generation config should be populated correctly."""
        client = self._get_client()
        wrapped = {
            "project": "test-project",
            "model": "gemini-3-flash",
            "request": {
                "contents": [],
                "generationConfig": {
                    "maxOutputTokens": 64000,
                    "temperature": 0.7,
                    "stopSequences": ["\n\nHuman:"],
                    "thinkingConfig": {
                        "includeThoughts": True,
                        "thinkingBudget": 8192,
                    }
                }
            }
        }
        req = client._build_request(wrapped)
        self.assertEqual(req.request.generation_config.max_output_tokens, 64000)
        self.assertAlmostEqual(req.request.generation_config.temperature, 0.7, places=2)
        self.assertTrue(req.request.generation_config.thinking_config.include_thoughts)
        self.assertEqual(req.request.generation_config.thinking_config.thinking_budget, 8192)

    def test_build_request_with_function_call_history(self):
        """Function call/response pairs in contents should be preserved."""
        client = self._get_client()
        wrapped = {
            "project": "test-project",
            "model": "gemini-3-flash",
            "request": {
                "contents": [
                    {"role": "user", "parts": [{"text": "List files"}]},
                    {"role": "model", "parts": [{
                        "functionCall": {"name": "exec_command", "args": {"cmd": "ls"}, "id": "call_1"}
                    }]},
                    {"role": "user", "parts": [{
                        "functionResponse": {"name": "exec_command", "response": {"result": "file.txt"}, "id": "call_1"}
                    }]},
                ]
            }
        }
        req = client._build_request(wrapped)
        self.assertEqual(len(req.request.contents), 3)
        # Verify function call preserved
        self.assertTrue(req.request.contents[1].parts[0].HasField("function_call"))
        self.assertEqual(req.request.contents[1].parts[0].function_call.name, "exec_command")
        # Verify function response preserved
        self.assertTrue(req.request.contents[2].parts[0].HasField("function_response"))
        self.assertEqual(req.request.contents[2].parts[0].function_response.name, "exec_command")


class TestGrpcEndpointsConfig(unittest.TestCase):
    """Tests for gRPC endpoint configuration."""

    def test_default_endpoints(self):
        """Default endpoints should include production and daily."""
        from antigravity_grpc.client import _GRPC_ENDPOINTS
        self.assertGreaterEqual(len(_GRPC_ENDPOINTS), 2)
        hostnames = [ep.split(":")[0] for ep in _GRPC_ENDPOINTS]
        self.assertIn("daily-cloudcode-pa.googleapis.com", hostnames)
        self.assertIn("cloudcode-pa.googleapis.com", hostnames)

    def test_staging_env_var(self):
        """Staging endpoints should be controlled by env var."""
        from antigravity_grpc.client import _ALLOW_STAGING_ENV
        self.assertEqual(_ALLOW_STAGING_ENV, "ALLOW_ANTIGRAVITY_STAGING")


class TestProxyIntegration(unittest.TestCase):
    """Tests for the proxy's gRPC fallback integration."""

    def _load_proxy_module(self):
        import importlib
        _spec = importlib.util.spec_from_file_location(
            "translate_proxy",
            os.path.join(_src_dir, "translate-proxy.py"),
        )
        tp = importlib.util.module_from_spec(_spec)
        _spec.loader.exec_module(tp)
        return tp

    def test_get_grpc_client_function_exists(self):
        """_get_grpc_client should exist as a module-level function."""
        tp = self._load_proxy_module()
        self.assertTrue(callable(tp._get_grpc_client))

    def test_grpc_fallback_errors_set(self):
        """_GRPC_FALLBACK_REST_ERRORS should include 404."""
        tp = self._load_proxy_module()
        self.assertIn(404, tp._GRPC_FALLBACK_REST_ERRORS)

    def test_versions_bug_fixed(self):
        """The _versions[0] NameError should be fixed (should be _fetched_ver)."""
        # Read the source file and verify _versions is not used incorrectly
        with open(os.path.join(_src_dir, "translate-proxy.py")) as f:
            source = f.read()
        # The bug was: ver={_versions[0]}  -- should be ver={_fetched_ver}
        self.assertNotIn("_versions[0]", source,
                         "Bug: _versions[0] should have been replaced with _fetched_ver")


if __name__ == "__main__":
    print("=" * 70)
    print("Antigravity gRPC Fallback - Unit Tests")
    print("=" * 70)
    print()

    unittest.main(verbosity=2)
