import base64
from types import SimpleNamespace
from unittest.mock import patch

import pytest

from swarm.tool import function_to_json
from swarm.util import image_to_base64


def test_basic_function():
    def basic_function(arg1, arg2):
        return arg1 + arg2

    result = function_to_json(basic_function)
    assert result == {
        "type": "function",
        "function": {
            "name": "basic_function",
            "description": "",
            "parameters": {
                "type": "object",
                "properties": {"arg1": {"type": "string"}, "arg2": {"type": "string"}},
                "required": ["arg1", "arg2"],
            },
        },
    }


def test_complex_function():
    def complex_function_with_types_and_descriptions(
        arg1: int, arg2: str, arg3: float = 3.14, arg4: bool = False
    ):
        """This is a complex function with a docstring."""
        pass

    result = function_to_json(complex_function_with_types_and_descriptions)
    assert result == {
        "type": "function",
        "function": {
            "name": "complex_function_with_types_and_descriptions",
            "description": "This is a complex function with a docstring.",
            "parameters": {
                "type": "object",
                "properties": {
                    "arg1": {"type": "integer"},
                    "arg2": {"type": "string"},
                    "arg3": {"type": "number"},
                    "arg4": {"type": "boolean"},
                },
                "required": ["arg1", "arg2"],
            },
        },
    }


def test_image_to_base64_raises_install_error_when_numpy_missing(tmp_path):
    image_path = tmp_path / "image.bin"
    image_path.write_bytes(b"\x00\x01\x02\x03\x04\x05")

    with patch("swarm.util.importlib.import_module", side_effect=ImportError):
        with pytest.raises(ImportError, match=r"swarm-agents\[image\]"):
            image_to_base64(str(image_path))


def test_image_to_base64_returns_base64_when_numpy_available(tmp_path):
    image_path = tmp_path / "image.bin"
    image_path.write_bytes(b"\x00\x01\x02\x03\x04\x05")

    class _FakeArray:
        def reshape(self, shape):
            _ = shape
            return self

        def __getitem__(self, key):
            _ = key
            return self

    class _FakeNumpy:
        uint8 = object()

        @staticmethod
        def frombuffer(data, dtype):
            _ = (data, dtype)
            return _FakeArray()

        @staticmethod
        def save(buffer, resized_img, allow_pickle=False):
            _ = (resized_img, allow_pickle)
            buffer.write(b"fake-npy")

    with patch(
        "swarm.util.importlib.import_module",
        return_value=SimpleNamespace(
            frombuffer=_FakeNumpy.frombuffer,
            save=_FakeNumpy.save,
            uint8=_FakeNumpy.uint8,
        ),
    ):
        result = image_to_base64(str(image_path), scale_factor=0.5)

    assert isinstance(result, str)
    assert base64.b64decode(result) == b"fake-npy"
