# 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 unittest
from http.client import HTTPResponse
from unittest.mock import Mock, patch

from gravitino import GravitinoClient, NameIdentifier
from gravitino.api.model.model import Model
from gravitino.api.model.model_version import ModelVersion
from gravitino.api.model_change import ModelChange
from gravitino.api.model_version_change import ModelVersionChange
from gravitino.dto.audit_dto import AuditDTO
from gravitino.dto.model_dto import ModelDTO
from gravitino.dto.model_version_dto import ModelVersionDTO
from gravitino.dto.responses.drop_response import DropResponse
from gravitino.dto.responses.entity_list_response import EntityListResponse
from gravitino.dto.responses.model_response import ModelResponse
from gravitino.dto.responses.model_version_list_response import (
    ModelVersionInfoListResponse,
    ModelVersionListResponse,
)
from gravitino.dto.responses.model_version_response import ModelVersionResponse
from gravitino.dto.responses.model_version_uri_response import ModelVersionUriResponse
from gravitino.namespace import Namespace
from gravitino.utils import Response
from tests.unittests import mock_base


@mock_base.mock_data
class TestModelCatalogApi(unittest.TestCase):
    _metalake_name: str = "metalake_demo"
    _catalog_name: str = "model_catalog"

    def test_list_models(self, *mock_method):
        gravitino_client = GravitinoClient(
            uri="http://localhost:8090", metalake_name=self._metalake_name
        )
        catalog = gravitino_client.load_catalog(self._catalog_name)

        ## test with response
        idents = [
            NameIdentifier.of(
                self._metalake_name, self._catalog_name, "schema", "model1"
            ),
            NameIdentifier.of(
                self._metalake_name, self._catalog_name, "schema", "model2"
            ),
        ]
        expected_idents = [
            NameIdentifier.of(ident.namespace().level(2), ident.name())
            for ident in idents
        ]
        entity_list_resp = EntityListResponse(_idents=idents, _code=0)
        json_str = entity_list_resp.to_json()
        mock_resp = self._mock_http_response(json_str)

        with patch(
            "gravitino.utils.http_client.HTTPClient.get",
            return_value=mock_resp,
        ):
            model_idents = catalog.as_model_catalog().list_models(
                Namespace.of("schema")
            )
            self.assertEqual(expected_idents, model_idents)

        ## test with empty response
        entity_list_resp_1 = EntityListResponse(_idents=[], _code=0)
        json_str_1 = entity_list_resp_1.to_json()
        mock_resp_1 = self._mock_http_response(json_str_1)

        with patch(
            "gravitino.utils.http_client.HTTPClient.get",
            return_value=mock_resp_1,
        ):
            model_idents = catalog.as_model_catalog().list_models(
                Namespace.of("schema")
            )
            self.assertEqual([], model_idents)

    def test_get_model(self, *mock_method):
        gravitino_client = GravitinoClient(
            uri="http://localhost:8090", metalake_name=self._metalake_name
        )
        catalog = gravitino_client.load_catalog(self._catalog_name)

        model_ident = NameIdentifier.of("schema", "model1")

        ## test with response
        model_dto = ModelDTO(
            _name="model1",
            _comment="this is test",
            _properties={"k": "v"},
            _latest_version=0,
            _audit=AuditDTO(_creator="test", _create_time="2022-01-01T00:00:00Z"),
        )
        model_resp = ModelResponse(_model=model_dto, _code=0)
        json_str = model_resp.to_json()
        mock_resp = self._mock_http_response(json_str)

        with patch(
            "gravitino.utils.http_client.HTTPClient.get",
            return_value=mock_resp,
        ):
            model = catalog.as_model_catalog().get_model(model_ident)
            self._compare_models(model_dto, model)

        ## test with empty response
        model_dto_1 = ModelDTO(
            _name="model1",
            _comment=None,
            _properties=None,
            _latest_version=0,
            _audit=AuditDTO(_creator="test", _create_time="2022-01-01T00:00:00Z"),
        )
        model_resp_1 = ModelResponse(_model=model_dto_1, _code=0)
        json_str_1 = model_resp_1.to_json()
        mock_resp_1 = self._mock_http_response(json_str_1)

        with patch(
            "gravitino.utils.http_client.HTTPClient.get",
            return_value=mock_resp_1,
        ):
            model = catalog.as_model_catalog().get_model(model_ident)
            self._compare_models(model_dto_1, model)

    def test_register_model(self, *mock_method):
        gravitino_client = GravitinoClient(
            uri="http://localhost:8090", metalake_name=self._metalake_name
        )
        catalog = gravitino_client.load_catalog(self._catalog_name)

        model_ident = NameIdentifier.of("schema", "model1")

        model_dto = ModelDTO(
            _name="model1",
            _comment="this is test",
            _properties={"k": "v"},
            _latest_version=0,
            _audit=AuditDTO(_creator="test", _create_time="2022-01-01T00:00:00Z"),
        )

        ## test with response
        model_resp = ModelResponse(_model=model_dto, _code=0)
        json_str = model_resp.to_json()
        mock_resp = self._mock_http_response(json_str)

        with patch(
            "gravitino.utils.http_client.HTTPClient.post",
            return_value=mock_resp,
        ):
            model = catalog.as_model_catalog().register_model(
                model_ident, "this is test", {"k": "v"}
            )
            self._compare_models(model_dto, model)

    def test_alter_model(self, *mock_method):
        gravitino_client = GravitinoClient(
            uri="http://localhost:8090", metalake_name=self._metalake_name
        )
        catalog = gravitino_client.load_catalog(self._catalog_name)

        model_ident = NameIdentifier.of("schema", "model1")

        model_dto = ModelDTO(
            _name="model2",
            _comment="this is test",
            _properties={"k": "v"},
            _latest_version=0,
            _audit=AuditDTO(_creator="test", _create_time="2022-01-01T00:00:00Z"),
        )

        ## test with response
        model_resp = ModelResponse(_model=model_dto, _code=0)
        json_str = model_resp.to_json()
        mock_resp = self._mock_http_response(json_str)

        with patch(
            "gravitino.utils.http_client.HTTPClient.put",
            return_value=mock_resp,
        ):
            model = catalog.as_model_catalog().alter_model(
                model_ident, ModelChange.rename("model2")
            )

            self._compare_models(model_dto, model)

    def test_delete_model(self, *mock_method):
        gravitino_client = GravitinoClient(
            uri="http://localhost:8090", metalake_name=self._metalake_name
        )
        catalog = gravitino_client.load_catalog(self._catalog_name)

        model_ident = NameIdentifier.of("schema", "model1")

        ## test with True response
        drop_resp = DropResponse(_dropped=True, _code=0)
        json_str = drop_resp.to_json()
        mock_resp = self._mock_http_response(json_str)

        with patch(
            "gravitino.utils.http_client.HTTPClient.delete",
            return_value=mock_resp,
        ):
            succ = catalog.as_model_catalog().delete_model(model_ident)
            self.assertTrue(succ)

        ## test with False response
        drop_resp_1 = DropResponse(_dropped=False, _code=0)
        json_str_1 = drop_resp_1.to_json()
        mock_resp_1 = self._mock_http_response(json_str_1)

        with patch(
            "gravitino.utils.http_client.HTTPClient.delete",
            return_value=mock_resp_1,
        ):
            succ = catalog.as_model_catalog().delete_model(model_ident)
            self.assertFalse(succ)

    def test_list_model_versions(self, *mock_method):
        gravitino_client = GravitinoClient(
            uri="http://localhost:8090", metalake_name=self._metalake_name
        )
        catalog = gravitino_client.load_catalog(self._catalog_name)

        model_ident = NameIdentifier.of("schema", "model1")

        ## test with response
        versions = [1, 2, 3]
        model_version_list_resp = ModelVersionListResponse(_versions=versions, _code=0)
        json_str = model_version_list_resp.to_json()
        mock_resp = self._mock_http_response(json_str)

        with patch(
            "gravitino.utils.http_client.HTTPClient.get",
            return_value=mock_resp,
        ):
            model_versions = catalog.as_model_catalog().list_model_versions(model_ident)
            self.assertEqual(versions, model_versions)

        ## test with empty response
        model_version_list_resp_1 = ModelVersionListResponse(_versions=[], _code=0)
        json_str_1 = model_version_list_resp_1.to_json()
        mock_resp_1 = self._mock_http_response(json_str_1)

        with patch(
            "gravitino.utils.http_client.HTTPClient.get",
            return_value=mock_resp_1,
        ):
            model_versions = catalog.as_model_catalog().list_model_versions(model_ident)
            self.assertEqual([], model_versions)

    def test_list_model_version_infos(self, *mock_method):
        gravitino_client = GravitinoClient(
            uri="http://localhost:8090", metalake_name=self._metalake_name
        )
        catalog = gravitino_client.load_catalog(self._catalog_name)

        model_ident = NameIdentifier.of("schema", "model1")

        model_versions_dto = [
            ModelVersionDTO(
                _version=0,
                _uris={"unknown": "http://localhost:8090"},
                _aliases=["alias1", "alias2"],
                _comment="this is test",
                _properties={"k": "v"},
                _audit=AuditDTO(_creator="test", _create_time="2022-01-01T00:00:00Z"),
            )
        ]
        model_version_info_list_resp = ModelVersionInfoListResponse(
            _versions=model_versions_dto, _code=0
        )
        json_str = model_version_info_list_resp.to_json()
        mock_resp = self._mock_http_response(json_str)

        with patch(
            "gravitino.utils.http_client.HTTPClient.get",
            return_value=mock_resp,
        ):
            model_versions = catalog.as_model_catalog().list_model_version_infos(
                model_ident
            )
            self.assertEqual(len(model_versions), 1)
            self._compare_model_versions(model_versions_dto[0], model_versions[0])

        model_version_info_list_resp_1 = ModelVersionInfoListResponse(
            _versions=[], _code=0
        )
        json_str_1 = model_version_info_list_resp_1.to_json()
        mock_resp_1 = self._mock_http_response(json_str_1)

        with patch(
            "gravitino.utils.http_client.HTTPClient.get",
            return_value=mock_resp_1,
        ):
            model_versions = catalog.as_model_catalog().list_model_version_infos(
                model_ident
            )
            self.assertEqual([], model_versions)

    def test_alter_model_version_comment(self, *mock_method):
        gravitino_client = GravitinoClient(
            uri="http://localhost:8090", metalake_name=self._metalake_name
        )
        catalog = gravitino_client.load_catalog(self._catalog_name)

        model_ident = NameIdentifier.of("schema", "model1")
        version = 1
        alias = "alias1"
        new_comment = "new comment"

        model_version_dto = ModelVersionDTO(
            _version=1,
            _uris={"unknown": "http://localhost:8090"},
            _aliases=["alias1", "alias2"],
            _comment="new comment",
            _properties={"k": "v"},
            _audit=AuditDTO(_creator="test", _create_time="2022-01-01T00:00:00Z"),
        )

        model_resp = ModelVersionResponse(_model_version=model_version_dto, _code=0)
        json_str = model_resp.to_json()
        mock_resp = self._mock_http_response(json_str)

        with patch(
            "gravitino.utils.http_client.HTTPClient.put",
            return_value=mock_resp,
        ):
            model_version = catalog.as_model_catalog().alter_model_version(
                model_ident, version, ModelVersionChange.update_comment(new_comment)
            )
            self._compare_model_versions(model_version_dto, model_version)

            model_version = catalog.as_model_catalog().alter_model_version(
                model_ident, alias, ModelVersionChange.update_comment(new_comment)
            )
            self._compare_model_versions(model_version_dto, model_version)

    def test_get_model_version(self, *mock_method):
        gravitino_client = GravitinoClient(
            uri="http://localhost:8090", metalake_name=self._metalake_name
        )
        catalog = gravitino_client.load_catalog(self._catalog_name)

        model_ident = NameIdentifier.of("schema", "model1")
        version = 1
        alias = "alias1"

        ## test with response
        model_version_dto = ModelVersionDTO(
            _version=1,
            _uris={"unknown": "http://localhost:8090"},
            _aliases=["alias1", "alias2"],
            _comment="this is test",
            _properties={"k": "v"},
            _audit=AuditDTO(_creator="test", _create_time="2022-01-01T00:00:00Z"),
        )
        model_resp = ModelVersionResponse(_model_version=model_version_dto, _code=0)
        json_str = model_resp.to_json()
        mock_resp = self._mock_http_response(json_str)

        with patch(
            "gravitino.utils.http_client.HTTPClient.get",
            return_value=mock_resp,
        ):
            model_version = catalog.as_model_catalog().get_model_version(
                model_ident, version
            )
            self._compare_model_versions(model_version_dto, model_version)

            model_version = catalog.as_model_catalog().get_model_version_by_alias(
                model_ident, alias
            )
            self._compare_model_versions(model_version_dto, model_version)

        ## test with empty response
        model_version_dto = ModelVersionDTO(
            _version=1,
            _uris={"unknown": "http://localhost:8090"},
            _aliases=None,
            _comment=None,
            _properties=None,
            _audit=AuditDTO(_creator="test", _create_time="2022-01-01T00:00:00Z"),
        )
        model_resp = ModelVersionResponse(_model_version=model_version_dto, _code=0)
        json_str = model_resp.to_json()
        mock_resp = self._mock_http_response(json_str)

        with patch(
            "gravitino.utils.http_client.HTTPClient.get",
            return_value=mock_resp,
        ):
            model_version = catalog.as_model_catalog().get_model_version(
                model_ident, version
            )
            self._compare_model_versions(model_version_dto, model_version)

            model_version = catalog.as_model_catalog().get_model_version_by_alias(
                model_ident, alias
            )
            self._compare_model_versions(model_version_dto, model_version)

    def test_link_model_version(self, *mock_method):
        gravitino_client = GravitinoClient(
            uri="http://localhost:8090", metalake_name=self._metalake_name
        )
        catalog = gravitino_client.load_catalog(self._catalog_name)

        model_ident = NameIdentifier.of("schema", "model1")

        ## test with response
        model_version_dto = ModelVersionDTO(
            _version=1,
            _uris={"unknown": "http://localhost:8090"},
            _aliases=["alias1", "alias2"],
            _comment="this is test",
            _properties={"k": "v"},
            _audit=AuditDTO(_creator="test", _create_time="2022-01-01T00:00:00Z"),
        )
        model_resp = ModelVersionResponse(_model_version=model_version_dto, _code=0)
        json_str = model_resp.to_json()
        mock_resp = self._mock_http_response(json_str)

        with patch(
            "gravitino.utils.http_client.HTTPClient.post",
            return_value=mock_resp,
        ):
            self.assertIsNone(
                catalog.as_model_catalog().link_model_version(
                    model_ident,
                    "http://localhost:8090",
                    ["alias1", "alias2"],
                    "this is test",
                    {"k": "v"},
                )
            )

    def test_link_model_version_with_multiple_uris(self, *mock_method):
        gravitino_client = GravitinoClient(
            uri="http://localhost:8090", metalake_name=self._metalake_name
        )
        catalog = gravitino_client.load_catalog(self._catalog_name)

        model_ident = NameIdentifier.of("schema", "model1")

        ## test with response
        model_version_dto = ModelVersionDTO(
            _version=1,
            _uris={"default-uri-name": "http://localhost:8090"},
            _aliases=["alias1", "alias2"],
            _comment="this is test",
            _properties={"k": "v"},
            _audit=AuditDTO(_creator="test", _create_time="2022-01-01T00:00:00Z"),
        )
        model_resp = ModelVersionResponse(_model_version=model_version_dto, _code=0)
        json_str = model_resp.to_json()
        mock_resp = self._mock_http_response(json_str)

        with patch(
            "gravitino.utils.http_client.HTTPClient.post",
            return_value=mock_resp,
        ):
            self.assertIsNone(
                catalog.as_model_catalog().link_model_version_with_multiple_uris(
                    model_ident,
                    {"default-uri-name": "http://localhost:8090"},
                    ["alias1", "alias2"],
                    "this is test",
                    {"k": "v"},
                )
            )

    def test_get_model_version_uri(self, *mock_method):
        gravitino_client = GravitinoClient(
            uri="http://localhost:8090", metalake_name=self._metalake_name
        )
        catalog = gravitino_client.load_catalog(self._catalog_name)

        model_ident = NameIdentifier.of("schema", "model1")
        version = 1
        alias = "alias1"

        ## test with response
        model_version_uri_resp = ModelVersionUriResponse(
            _code=0, _uri="s3://path/to/model"
        )
        json_str = model_version_uri_resp.to_json()
        mock_resp = self._mock_http_response(json_str)

        with patch(
            "gravitino.utils.http_client.HTTPClient.get",
            return_value=mock_resp,
        ):
            model_version_uri = catalog.as_model_catalog().get_model_version_uri(
                model_ident, version, "uri_name"
            )
            self.assertEqual("s3://path/to/model", model_version_uri)

            model_version_uri = (
                catalog.as_model_catalog().get_model_version_uri_by_alias(
                    model_ident, alias, "uri_name"
                )
            )
            self.assertEqual("s3://path/to/model", model_version_uri)

    def test_delete_model_version(self, *mock_method):
        gravitino_client = GravitinoClient(
            uri="http://localhost:8090", metalake_name=self._metalake_name
        )
        catalog = gravitino_client.load_catalog(self._catalog_name)

        model_ident = NameIdentifier.of("schema", "model1")
        version = 1
        alias = "alias1"

        ## test with True response
        drop_resp = DropResponse(_dropped=True, _code=0)
        json_str = drop_resp.to_json()
        mock_resp = self._mock_http_response(json_str)

        with patch(
            "gravitino.utils.http_client.HTTPClient.delete",
            return_value=mock_resp,
        ):
            succ = catalog.as_model_catalog().delete_model_version(model_ident, version)
            self.assertTrue(succ)

            succ = catalog.as_model_catalog().delete_model_version_by_alias(
                model_ident, alias
            )
            self.assertTrue(succ)

        ## test with False response
        drop_resp_1 = DropResponse(_dropped=False, _code=0)
        json_str_1 = drop_resp_1.to_json()
        mock_resp_1 = self._mock_http_response(json_str_1)

        with patch(
            "gravitino.utils.http_client.HTTPClient.delete",
            return_value=mock_resp_1,
        ):
            succ = catalog.as_model_catalog().delete_model_version(model_ident, version)
            self.assertFalse(succ)

            succ = catalog.as_model_catalog().delete_model_version_by_alias(
                model_ident, alias
            )
            self.assertFalse(succ)

    def _mock_http_response(self, json_str: str):
        mock_http_resp = Mock(HTTPResponse)
        mock_http_resp.getcode.return_value = 200
        mock_http_resp.read.return_value = json_str
        mock_http_resp.info.return_value = None
        mock_http_resp.url = None
        mock_resp = Response(mock_http_resp)
        return mock_resp

    def _compare_models(self, left: Model, right: Model):
        self.assertEqual(left.name(), right.name())
        self.assertEqual(left.comment(), right.comment())
        self.assertEqual(left.properties(), right.properties())
        self.assertEqual(left.latest_version(), right.latest_version())

    def _compare_model_versions(self, left: ModelVersion, right: ModelVersion):
        self.assertEqual(left.version(), right.version())
        self.assertEqual(left.uri(), right.uri())
        self.assertEqual(left.aliases(), right.aliases())
        self.assertEqual(left.comment(), right.comment())
        self.assertEqual(left.properties(), right.properties())
