from unittest import mock
from urllib.parse import quote

import pytest

from snowflake.core import PollingOperation, Root
from snowflake.core.tag import TagValue
from snowflake.core.user_defined_function import (
    ColumnType,
    ReturnDataType,
    ReturnTable,
    SQLFunction,
    UserDefinedFunction,
    UserDefinedFunctionResource,
)
from snowflake.core.user_defined_function._generated import TagAssignment, TagReference

from ...utils import BASE_URL, extra_params, mock_http_response, random_string


API_CLIENT_REQUEST = "snowflake.core._generated.api_client.ApiClient.request"
UDF = UserDefinedFunction(
    name="my_udf", arguments=[], return_type=ReturnDataType(datatype="VARCHAR"), language_config=SQLFunction()
)


@pytest.fixture
def user_defined_functions(schema):
    return schema.user_defined_functions


@pytest.fixture
def user_defined_function(user_defined_functions):
    return user_defined_functions["my_udf()"]


def test_create_user_defined_function(fake_root, user_defined_functions):
    args = (fake_root, "POST", BASE_URL + "/databases/my_db/schemas/my_schema/user-defined-functions")
    kwargs = extra_params(
        query_params=[],
        body={
            "name": "my_udf",
            "arguments": [],
            "return_type": {"datatype": "VARCHAR", "type": "DATATYPE"},
            "language_config": {"language": "SQL"},
        },
    )

    with mock.patch(API_CLIENT_REQUEST) as mocked_request:
        udf_res = user_defined_functions.create(UDF)
        assert isinstance(udf_res, UserDefinedFunctionResource)
        assert udf_res.name_with_args == "my_udf()"
    mocked_request.assert_called_once_with(*args, **kwargs)

    with mock.patch(API_CLIENT_REQUEST) as mocked_request:
        op = user_defined_functions.create_async(UDF)
        assert isinstance(op, PollingOperation)
        udf_res = op.result()
        assert udf_res.name_with_args == "my_udf()"
    mocked_request.assert_called_once_with(*args, **kwargs)


def test_iter_user_defined_function(fake_root, user_defined_functions):
    args = (fake_root, "GET", BASE_URL + "/databases/my_db/schemas/my_schema/user-defined-functions")
    kwargs = extra_params()

    with mock.patch(API_CLIENT_REQUEST) as mocked_request:
        mocked_request.return_value = mock_http_response()
        user_defined_functions.iter()
    mocked_request.assert_called_once_with(*args, **kwargs)

    with mock.patch(API_CLIENT_REQUEST) as mocked_request:
        mocked_request.return_value = mock_http_response()
        op = user_defined_functions.iter_async()
        assert isinstance(op, PollingOperation)
        it = op.result()
        assert list(it) == []
    mocked_request.assert_called_once_with(*args, **kwargs)


def test_fetch_user_defined_function(fake_root, user_defined_function):
    from snowflake.core.user_defined_function._generated.models import ReturnDataType as ReturnDataTypeModel
    from snowflake.core.user_defined_function._generated.models import SQLFunction as SQLFunctionModel
    from snowflake.core.user_defined_function._generated.models import UserDefinedFunction as UserDefinedFunctionModel

    model = UserDefinedFunctionModel(
        name="my_udf",
        arguments=[],
        return_type=ReturnDataTypeModel(datatype="VARCHAR"),
        language_config=SQLFunctionModel(),
    )
    args = (
        fake_root,
        "GET",
        BASE_URL + f"/databases/my_db/schemas/my_schema/user-defined-functions/{quote('my_udf()')}",
    )
    kwargs = extra_params()

    with mock.patch(API_CLIENT_REQUEST) as mocked_request:
        mocked_request.return_value = mock_http_response(model.to_json())
        user_defined_function.fetch()
    mocked_request.assert_called_once_with(*args, **kwargs)

    with mock.patch(API_CLIENT_REQUEST) as mocked_request:
        mocked_request.return_value = mock_http_response(model.to_json())
        op = user_defined_function.fetch_async()
        assert isinstance(op, PollingOperation)
        tab = op.result()
        assert tab.to_dict() == UDF.to_dict()
    mocked_request.assert_called_once_with(*args, **kwargs)


def test_drop_user_defined_function(fake_root, user_defined_function):
    args = (
        fake_root,
        "DELETE",
        BASE_URL + f"/databases/my_db/schemas/my_schema/user-defined-functions/{quote('my_udf()')}",
    )
    kwargs = extra_params()

    with mock.patch(API_CLIENT_REQUEST) as mocked_request:
        user_defined_function.drop()
    mocked_request.assert_called_once_with(*args, **kwargs)

    with mock.patch(API_CLIENT_REQUEST) as mocked_request:
        op = user_defined_function.drop_async()
        assert isinstance(op, PollingOperation)
        op.result()
    mocked_request.assert_called_once_with(*args, **kwargs)


def test_rename_user_defined_function(fake_root, user_defined_function, user_defined_functions):
    def format_args(udf_name: str) -> tuple[Root, str, str]:
        return (
            fake_root,
            "POST",
            BASE_URL
            + f"/databases/my_db/schemas/my_schema/user-defined-functions/{quote(udf_name)}:rename?"
            + "targetDatabase=my_db&targetSchema=my_schema&targetName=new_udf",
        )

    kwargs = extra_params(
        query_params=[("targetDatabase", "my_db"), ("targetSchema", "my_schema"), ("targetName", "new_udf")]
    )

    with mock.patch(API_CLIENT_REQUEST) as mocked_request:
        user_defined_function.rename("new_udf")
        assert user_defined_function.name_with_args == "new_udf()"
    mocked_request.assert_called_once_with(*format_args("my_udf()"), **kwargs)

    with mock.patch(API_CLIENT_REQUEST) as mocked_request:
        udf_res2 = user_defined_functions["my_udf2()"]
        op = udf_res2.rename_async("new_udf")
        assert isinstance(op, PollingOperation)
        op.result()
        assert udf_res2.name_with_args == "new_udf()"
    mocked_request.assert_called_once_with(*format_args("my_udf2()"), **kwargs)


def test_execute_udf_table_raises_not_implemented_error(fake_root, user_defined_functions):
    udf_name = random_string(10, "test_create_user_defined_function_sql_")
    func_body = """
        SELECT 1, 2
        UNION ALL
        SELECT 3, 4
    """
    model = UserDefinedFunction(
        name=udf_name,
        arguments=[],
        return_type=ReturnTable(
            column_list=[ColumnType(name="x", datatype="INTEGER"), ColumnType(name="y", datatype="INTEGER")]
        ),
        language_config=SQLFunction(),
        body=func_body,
    )
    with mock.patch(API_CLIENT_REQUEST) as mocked_request:
        udf_handle = user_defined_functions.create(model)
        mocked_request.return_value = mock_http_response(model.to_json())
        with pytest.raises(NotImplementedError):
            udf_handle.execute([])


def test_set_tags(fake_root, user_defined_function, tag):
    args = (
        fake_root,
        "POST",
        BASE_URL + f"/databases/my_db/schemas/my_schema/user-defined-functions/{quote('my_udf()')}:set-tags",
    )
    tags = {tag: TagValue(value="value")}
    kwargs = extra_params(
        body=[
            TagAssignment(
                tag_value=v.value, tag_name=k.name, tag_schema=k.schema.name, tag_database=k.database.name
            ).to_dict()
            for k, v in tags.items()
        ]
    )

    with mock.patch(API_CLIENT_REQUEST) as mocked_request:
        user_defined_function.set_tags(tags)
    mocked_request.assert_called_once_with(*args, **kwargs)

    with mock.patch(API_CLIENT_REQUEST) as mocked_request:
        op = user_defined_function.set_tags_async(tags)
        assert isinstance(op, PollingOperation)
        op.result()
    mocked_request.assert_called_once_with(*args, **kwargs)


def test_unset_tags(fake_root, user_defined_function, tag):
    args = (
        fake_root,
        "POST",
        BASE_URL + f"/databases/my_db/schemas/my_schema/user-defined-functions/{quote('my_udf()')}:unset-tags",
    )
    tag_resources = {tag}
    kwargs = extra_params(
        body=[
            TagReference(
                tag_name=tag_res.name, tag_schema=tag_res.schema.name, tag_database=tag_res.database.name
            ).to_dict()
            for tag_res in tag_resources
        ]
    )

    with mock.patch(API_CLIENT_REQUEST) as mocked_request:
        user_defined_function.unset_tags(tag_resources)
    mocked_request.assert_called_once_with(*args, **kwargs)

    with mock.patch(API_CLIENT_REQUEST) as mocked_request:
        op = user_defined_function.unset_tags_async(tag_resources)
        assert isinstance(op, PollingOperation)
        op.result()
    mocked_request.assert_called_once_with(*args, **kwargs)


def test_get_tags(fake_root, user_defined_function):
    args = (
        fake_root,
        "GET",
        BASE_URL + f"/databases/my_db/schemas/my_schema/user-defined-functions/{quote('my_udf()')}:get-tags",
    )
    kwargs = extra_params()

    with mock.patch(API_CLIENT_REQUEST) as mocked_request:
        mocked_request.return_value = mock_http_response()
        assert user_defined_function.get_tags() == {}
    mocked_request.assert_called_once_with(*args, **kwargs)

    with mock.patch(API_CLIENT_REQUEST) as mocked_request:
        mocked_request.return_value = mock_http_response()
        op = user_defined_function.get_tags_async()
        assert isinstance(op, PollingOperation)
        assert op.result() == {}
    mocked_request.assert_called_once_with(*args, **kwargs)
