# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import shutil
import tempfile
import unittest
import uuid

from pypaimon.api.api_response import ConfigResponse
from pypaimon.api.auth import BearTokenAuthProvider
from pypaimon.api.rest_api import RESTApi, IllegalArgumentError
from pypaimon.catalog.catalog_exception import (
    FunctionNotExistException,
    FunctionAlreadyExistException,
    DefinitionAlreadyExistException,
    DefinitionNotExistException,
)
from pypaimon.catalog.catalog_context import CatalogContext
from pypaimon.catalog.rest.rest_catalog import RESTCatalog
from pypaimon.common.identifier import Identifier
from pypaimon.common.options import Options
from pypaimon.function.function import FunctionImpl
from pypaimon.function.function_change import FunctionChange
from pypaimon.function.function_definition import (
    FunctionDefinition, FunctionFileResource,
)
from pypaimon.schema.data_types import AtomicType, DataField
from pypaimon.tests.rest.rest_server import RESTCatalogServer


def _mock_function(identifier: Identifier) -> FunctionImpl:
    input_params = [
        DataField(0, "length", AtomicType("DOUBLE")),
        DataField(1, "width", AtomicType("DOUBLE")),
    ]
    return_params = [
        DataField(0, "area", AtomicType("DOUBLE")),
    ]
    flink_function = FunctionDefinition.file(
        file_resources=[FunctionFileResource("jar", "/a/b/c.jar")],
        language="java",
        class_name="className",
        function_name="eval",
    )
    spark_function = FunctionDefinition.lambda_def(
        "(Double length, Double width) -> length * width", "java"
    )
    trino_function = FunctionDefinition.sql("length * width")
    definitions = {
        "flink": flink_function,
        "spark": spark_function,
        "trino": trino_function,
    }
    return FunctionImpl(
        identifier=identifier,
        input_params=input_params,
        return_params=return_params,
        deterministic=False,
        definitions=definitions,
        comment="comment",
        options={},
    )


class RESTFunctionTest(unittest.TestCase):

    def setUp(self):
        self.temp_dir = tempfile.mkdtemp(prefix="function_test_")
        self.config = ConfigResponse(defaults={"prefix": "mock-test"})
        self.token = str(uuid.uuid4())
        self.server = RESTCatalogServer(
            data_path=self.temp_dir,
            auth_provider=BearTokenAuthProvider(self.token),
            config=self.config,
            warehouse="warehouse",
        )
        self.server.start()

        options = Options({
            "metastore": "rest",
            "uri": f"http://localhost:{self.server.port}",
            "warehouse": "warehouse",
            "token.provider": "bear",
            "token": self.token,
        })
        self.catalog = RESTCatalog(CatalogContext.create_from_options(options))

    def tearDown(self):
        self.server.shutdown()
        import gc
        gc.collect()
        shutil.rmtree(self.temp_dir, ignore_errors=True)

    def test_function(self):
        self.catalog.create_database("rest_catalog_db", True)

        identifier_with_slash = Identifier.create("rest_catalog_db", "function/")
        with self.assertRaises(IllegalArgumentError):
            self.catalog.create_function(
                identifier_with_slash,
                _mock_function(identifier_with_slash),
                False,
            )
        with self.assertRaises(FunctionNotExistException):
            self.catalog.get_function(identifier_with_slash)
        with self.assertRaises(IllegalArgumentError):
            self.catalog.drop_function(identifier_with_slash, True)

        identifier_without_alphabet = Identifier.create("rest_catalog_db", "-")
        with self.assertRaises(IllegalArgumentError):
            self.catalog.create_function(
                identifier_without_alphabet,
                _mock_function(identifier_without_alphabet),
                False,
            )
        with self.assertRaises(FunctionNotExistException):
            self.catalog.get_function(identifier_without_alphabet)
        with self.assertRaises(IllegalArgumentError):
            self.catalog.drop_function(identifier_without_alphabet, True)

        identifier = Identifier.from_string("rest_catalog_db.function.na_me-01")
        function = _mock_function(identifier)

        # Drop first to ensure fresh data (not stale from a previous failed run)
        self.catalog.drop_function(identifier, True)
        self.catalog.create_function(identifier, function, True)
        with self.assertRaises(FunctionAlreadyExistException):
            self.catalog.create_function(identifier, function, False)

        self.assertIn(function.name(), self.catalog.list_functions(identifier.get_database_name()))

        get_function = self.catalog.get_function(identifier)
        self.assertEqual(get_function.name(), function.name())
        for dialect in function.definitions().keys():
            self.assertEqual(get_function.definition(dialect), function.definition(dialect))

        self.catalog.drop_function(identifier, True)
        self.assertNotIn(function.name(), self.catalog.list_functions(identifier.get_database_name()))

        with self.assertRaises(FunctionNotExistException):
            self.catalog.drop_function(identifier, False)
        with self.assertRaises(FunctionNotExistException):
            self.catalog.get_function(identifier)

    def test_list_functions(self):
        db1 = "db_rest_catalog_db"
        db2 = "db2_rest_catalog"
        identifier = Identifier.create(db1, "list_function")
        identifier1 = Identifier.create(db1, "function")
        identifier2 = Identifier.create(db2, "list_function")
        identifier3 = Identifier.create(db2, "function")

        self.catalog.create_database(db1, True)
        self.catalog.create_database(db2, True)
        self.catalog.create_function(identifier, _mock_function(identifier), True)
        self.catalog.create_function(identifier1, _mock_function(identifier1), True)
        self.catalog.create_function(identifier2, _mock_function(identifier2), True)
        self.catalog.create_function(identifier3, _mock_function(identifier3), True)

        result = self.catalog.list_functions_paged(db1, None, None, None)
        self.assertEqual(
            set(result.elements),
            {identifier.get_object_name(), identifier1.get_object_name()},
        )

        result = self.catalog.list_functions_paged(db1, 1, None, None)
        self.assertEqual(len(result.elements), 1)
        self.assertIn(
            result.elements[0],
            [identifier.get_object_name(), identifier1.get_object_name()],
        )

        result = self.catalog.list_functions_paged(
            db1, 1, identifier1.get_object_name(), None)
        self.assertEqual(
            result.elements,
            [identifier.get_object_name()],
        )

        result = self.catalog.list_functions_paged(db1, None, None, "func%")
        self.assertEqual(result.elements, [identifier1.get_object_name()])

        result = self.catalog.list_functions_paged_globally("db2_rest%", "func%", None, None)
        self.assertEqual(len(result.elements), 1)
        self.assertEqual(result.elements[0].get_full_name(), identifier3.get_full_name())

        result = self.catalog.list_functions_paged_globally(
            "db2_rest%", None, 1, None)
        self.assertEqual(len(result.elements), 1)
        self.assertIn(
            result.elements[0].get_full_name(),
            [identifier2.get_full_name(), identifier3.get_full_name()],
        )

        result = self.catalog.list_functions_paged_globally(
            "db2_rest%", None, 1, identifier3.get_full_name())
        self.assertEqual(len(result.elements), 1)
        self.assertEqual(
            result.elements[0].get_full_name(),
            identifier2.get_full_name(),
        )

        result = self.catalog.list_function_details_paged(db1, 1, None, None)
        self.assertEqual(len(result.elements), 1)
        self.assertIn(
            result.elements[0].full_name(),
            [identifier.get_full_name(), identifier1.get_full_name()],
        )

        result = self.catalog.list_function_details_paged(db2, 4, None, "func%")
        self.assertEqual(len(result.elements), 1)
        self.assertEqual(
            result.elements[0].full_name(), identifier3.get_full_name())

        result = self.catalog.list_function_details_paged(
            db2, 1, identifier3.get_object_name(), None)
        full_names = [f.full_name() for f in result.elements]
        self.assertIn(identifier2.get_full_name(), full_names)

    def test_alter_function(self):
        identifier = Identifier.create("rest_catalog_db", "alter_function_name")
        self.catalog.create_database(identifier.get_database_name(), True)
        self.catalog.drop_function(identifier, True)
        function = _mock_function(identifier)
        definition = FunctionDefinition.sql("x * y + 1")
        add_definition = FunctionChange.add_definition("flink_1", definition)

        self.catalog.alter_function(identifier, [add_definition], True)

        with self.assertRaises(FunctionNotExistException):
            self.catalog.alter_function(identifier, [add_definition], False)

        self.catalog.create_function(identifier, function, True)

        key = str(uuid.uuid4())
        value = str(uuid.uuid4())
        set_option = FunctionChange.set_option(key, value)
        self.catalog.alter_function(identifier, [set_option], False)
        catalog_function = self.catalog.get_function(identifier)
        self.assertEqual(catalog_function.options().get(key), value)

        self.catalog.alter_function(identifier, [FunctionChange.remove_option(key)], False)
        catalog_function = self.catalog.get_function(identifier)
        self.assertNotIn(key, catalog_function.options())

        new_comment = "new comment"
        self.catalog.alter_function(
            identifier, [FunctionChange.update_comment(new_comment)], False
        )
        catalog_function = self.catalog.get_function(identifier)
        self.assertEqual(catalog_function.comment(), new_comment)

        self.catalog.alter_function(identifier, [add_definition], False)
        catalog_function = self.catalog.get_function(identifier)
        self.assertEqual(
            catalog_function.definition(add_definition.name),
            add_definition.definition,
        )

        with self.assertRaises(DefinitionAlreadyExistException):
            self.catalog.alter_function(identifier, [add_definition], False)

        update_definition = FunctionChange.update_definition("flink_1", definition)
        self.catalog.alter_function(identifier, [update_definition], False)
        catalog_function = self.catalog.get_function(identifier)
        self.assertEqual(
            catalog_function.definition(update_definition.name),
            update_definition.definition,
        )

        with self.assertRaises(DefinitionNotExistException):
            self.catalog.alter_function(
                identifier,
                [FunctionChange.update_definition("no_exist", definition)],
                False,
            )

        drop_definition = FunctionChange.drop_definition(update_definition.name)
        self.catalog.alter_function(identifier, [drop_definition], False)
        catalog_function = self.catalog.get_function(identifier)
        self.assertIsNone(catalog_function.definition(update_definition.name))

        with self.assertRaises(DefinitionNotExistException):
            self.catalog.alter_function(identifier, [drop_definition], False)

    def test_validate_function_name(self):
        self.assertTrue(RESTApi.is_valid_function_name("a"))
        self.assertTrue(RESTApi.is_valid_function_name("a1_"))
        self.assertTrue(RESTApi.is_valid_function_name("a-b_c"))
        self.assertTrue(RESTApi.is_valid_function_name("a-b_c.1"))

        self.assertFalse(RESTApi.is_valid_function_name("a\\/b"))
        self.assertFalse(RESTApi.is_valid_function_name("a$?b"))
        self.assertFalse(RESTApi.is_valid_function_name("a@b"))
        self.assertFalse(RESTApi.is_valid_function_name("a*b"))
        self.assertFalse(RESTApi.is_valid_function_name("123"))
        self.assertFalse(RESTApi.is_valid_function_name("_-"))
        self.assertFalse(RESTApi.is_valid_function_name(""))
        self.assertFalse(RESTApi.is_valid_function_name(None))

        with self.assertRaises(IllegalArgumentError):
            RESTApi.check_function_name("a\\/b")
        with self.assertRaises(IllegalArgumentError):
            RESTApi.check_function_name("123")
        with self.assertRaises(IllegalArgumentError):
            RESTApi.check_function_name("")
        with self.assertRaises(IllegalArgumentError):
            RESTApi.check_function_name(None)


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