from __future__ import annotations

import asyncio
import unittest
from unittest.mock import patch

from openclaw_lite.providers.base import (
    AsyncJsonHttpClient,
    ProviderError,
    build_gemini_contents,
    build_gemini_tools,
    build_openai_messages,
    build_openai_tools,
    extract_gemini_model_names,
    extract_openai_model_names,
    normalise_gemini_response,
    normalise_openai_response,
)
from openclaw_lite.providers.factory import ProviderRegistry, build_provider_registry
from openclaw_lite.providers.anthropic import AnthropicProvider
from openclaw_lite.providers.gemini import GeminiProvider
from openclaw_lite.providers.openai import OpenAIProvider
from openclaw_lite.providers.openai_compatible import OpenAICompatibleProvider


class OpenAIHelpersTest(unittest.TestCase):
    def test_build_openai_tools_accepts_short_form(self) -> None:
        tools = build_openai_tools(
            [
                {
                    "name": "run_shell",
                    "description": "Run a command",
                    "parameters": {"type": "object", "properties": {"command": {"type": "string"}}},
                }
            ]
        )
        self.assertEqual(tools[0]["type"], "function")
        self.assertEqual(tools[0]["function"]["name"], "run_shell")

    def test_build_openai_messages_preserves_tool_context(self) -> None:
        messages = build_openai_messages(
            [
                {"role": "user", "content": "hello"},
                {
                    "role": "assistant",
                    "content": "",
                    "tool_calls": [{"id": "call_1", "name": "ping", "arguments": {"value": 1}}],
                },
                {"role": "tool", "name": "ping", "tool_call_id": "call_1", "content": {"ok": True}},
            ],
            system_prompt="sys",
        )
        self.assertEqual(messages[0]["role"], "system")
        self.assertEqual(messages[2]["tool_calls"][0]["function"]["arguments"], '{"value":1}')
        self.assertEqual(messages[3]["tool_call_id"], "call_1")

    def test_normalise_openai_response_extracts_text_and_tool_calls(self) -> None:
        payload = {
            "model": "gpt-test",
            "usage": {"total_tokens": 42},
            "choices": [
                {
                    "message": {
                        "content": "Done",
                        "tool_calls": [
                            {
                                "id": "call_1",
                                "function": {
                                    "name": "send_message",
                                    "arguments": '{"chat_id":123,"text":"hi"}',
                                },
                            }
                        ],
                    }
                }
            ],
        }

        normalized = normalise_openai_response(payload, fallback_model="fallback", provider_name="openai")
        self.assertEqual(normalized["text"], "Done")
        self.assertEqual(normalized["tool_calls"][0]["name"], "send_message")
        self.assertEqual(normalized["tool_calls"][0]["arguments"]["chat_id"], 123)
        self.assertEqual(normalized["usage"]["total_tokens"], 42)

    def test_extract_openai_model_names(self) -> None:
        names = extract_openai_model_names({"data": [{"id": "gpt-4.1-mini"}, {"id": "gpt-4.1-mini"}, {"id": "gpt-4.1"}]})
        self.assertEqual(names, ["gpt-4.1", "gpt-4.1-mini"])

    def test_http_client_wraps_raw_timeout_as_provider_error(self) -> None:
        client = AsyncJsonHttpClient(timeout=12.5)

        with patch("openclaw_lite.providers.base.request.urlopen", side_effect=TimeoutError("timed out")):
            with self.assertRaisesRegex(ProviderError, r"Request timeout for https://example\.com after 12\.5s"):
                client._request_json("POST", "https://example.com", None, {"ping": True})


class GeminiHelpersTest(unittest.TestCase):
    def test_build_gemini_tools_accepts_openai_style_tools(self) -> None:
        tools = build_gemini_tools(
            [
                {
                    "type": "function",
                    "function": {
                        "name": "search_web",
                        "description": "Search docs",
                        "parameters": {"type": "object", "properties": {"query": {"type": "string"}}},
                    },
                }
            ]
        )
        self.assertEqual(tools[0]["functionDeclarations"][0]["name"], "search_web")

    def test_build_gemini_contents_handles_tool_response(self) -> None:
        contents = build_gemini_contents(
            [
                {"role": "user", "content": "hello"},
                {"role": "tool", "name": "lookup", "content": {"answer": "world"}},
            ]
        )
        self.assertEqual(contents[0]["parts"][0]["text"], "hello")
        function_response = contents[1]["parts"][0]["functionResponse"]
        self.assertEqual(function_response["name"], "lookup")
        self.assertEqual(function_response["response"]["answer"], "world")

    def test_normalise_gemini_response_extracts_text_and_tools(self) -> None:
        payload = {
            "modelVersion": "gemini-2.5-flash",
            "usageMetadata": {"totalTokenCount": 10},
            "candidates": [
                {
                    "content": {
                        "parts": [
                            {"text": "I can do that"},
                            {"functionCall": {"name": "run_shell", "args": {"command": "pwd"}}},
                        ]
                    }
                }
            ],
        }

        normalized = normalise_gemini_response(payload, fallback_model="fallback", provider_name="gemini")
        self.assertEqual(normalized["text"], "I can do that")
        self.assertEqual(normalized["tool_calls"][0]["name"], "run_shell")
        self.assertEqual(normalized["tool_calls"][0]["arguments"]["command"], "pwd")
        self.assertEqual(normalized["usage"]["totalTokenCount"], 10)

    def test_extract_gemini_model_names_filters_non_generation_models(self) -> None:
        names = extract_gemini_model_names(
            {
                "models": [
                    {"name": "models/gemini-2.5-flash", "supportedGenerationMethods": ["generateContent"]},
                    {"name": "models/text-embedding-004", "supportedGenerationMethods": ["embedContent"]},
                ]
            }
        )
        self.assertEqual(names, ["gemini-2.5-flash"])


class ProviderRegistryTest(unittest.TestCase):
    def test_build_provider_registry_supports_all_provider_types(self) -> None:
        registry = build_provider_registry(
            {
                "providers": {
                    "items": {
                        "openai": {"type": "openai", "api_key": "k1"},
                        "router": {
                            "type": "openai_compatible",
                            "api_key": "k2",
                            "base_url": "https://router.example/v1",
                            "headers": {"X-Test": "1"},
                        },
                        "gemini": {"type": "gemini", "api_key": "k3"},
                        "claude": {"type": "claude", "api_key": "k4"},
                    }
                }
            }
        )
        self.assertIsInstance(registry, ProviderRegistry)
        self.assertIsInstance(registry.get("openai"), OpenAIProvider)
        self.assertIsInstance(registry.get("router"), OpenAICompatibleProvider)
        self.assertIsInstance(registry.get("gemini"), GeminiProvider)
        self.assertIsInstance(registry.get("claude"), AnthropicProvider)
        self.assertEqual(registry.list_provider_names(), ["claude", "gemini", "openai", "router"])

    def test_registry_reload_rebuilds_instances(self) -> None:
        registry = build_provider_registry({"providers": {"items": {"a": {"type": "openai", "api_key": "k1"}}}})
        before = registry.get("a")
        registry.reload({"providers": {"items": {"b": {"type": "gemini", "api_key": "k2"}}}})
        self.assertNotIn("a", registry.list_provider_names())
        self.assertIn("b", registry.list_provider_names())
        self.assertNotEqual(before.provider_name, registry.get("b").provider_name)

    def test_disabled_providers_are_skipped(self) -> None:
        registry = build_provider_registry(
            {
                "providers": {
                    "items": {
                        "openai": {"type": "openai", "api_key": "k1", "enabled": True},
                        "gemini": {"type": "gemini", "api_key": "k2", "enabled": False},
                    }
                }
            }
        )
        self.assertEqual(registry.list_provider_names(), ["openai"])

    def test_list_all_models_aggregates_async_results(self) -> None:
        registry = ProviderRegistry(providers={}, config={})

        class DummyProvider:
            def __init__(self, models):
                self._models = models

            async def list_models(self):
                return self._models

        registry.providers = {"demo": DummyProvider(["demo", "demo-2"]), "aux": DummyProvider(["x"])}
        models = asyncio.run(registry.list_all_models())
        self.assertEqual(models, {"demo": ["demo", "demo-2"], "aux": ["x"]})


if __name__ == "__main__":
    unittest.main()


    def test_registry_can_build_without_keys_for_bot_first_setup(self) -> None:
        registry = build_provider_registry(
            {
                "providers": {
                    "items": {
                        "openai": {"type": "openai", "api_key": ""},
                        "gemini": {"type": "gemini", "api_key": ""},
                        "claude": {"type": "claude", "api_key": ""},
                    }
                }
            }
        )
        self.assertEqual(registry.list_provider_names(), ["claude", "gemini", "openai"])
