# 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.

from __future__ import annotations

import json as _json
import unittest

from gravitino.dto.rel.partitions.json_serdes.partition_dto_serdes import (
    PartitionDTOSerdes,
)
from gravitino.dto.responses.credential_response import CredentialResponse
from gravitino.dto.responses.file_location_response import FileLocationResponse
from gravitino.dto.responses.function_list_response import FunctionListResponse
from gravitino.dto.responses.function_response import FunctionResponse
from gravitino.dto.responses.model_response import ModelResponse
from gravitino.dto.responses.model_version_list_response import ModelVersionListResponse
from gravitino.dto.responses.model_version_response import ModelVersionResponse
from gravitino.dto.responses.model_version_uri_response import ModelVersionUriResponse
from gravitino.dto.responses.name_list_response import NameListResponse
from gravitino.dto.responses.partition_list_response import PartitionListResponse
from gravitino.dto.responses.partition_name_list_response import (
    PartitionNameListResponse,
)
from gravitino.dto.responses.partition_response import PartitionResponse
from gravitino.dto.responses.table_response import TableResponse
from gravitino.exceptions.base import IllegalArgumentException


class TestResponses(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.PARTITION_JSON_STRING = """
            {
                "type": "identity",
                "name": "test_identity_partition",
                "fieldNames": [
                    [
                        "upper"
                    ],
                    [
                        "lower"
                    ]
                ],
                "values": [
                    {
                        "type": "literal",
                        "dataType": "integer",
                        "value": "0"
                    },
                    {
                        "type": "literal",
                        "dataType": "integer",
                        "value": "100"
                    }
                ],
                "properties": {
                    "key1": "value1",
                    "key2": "value2"
                }
            }
        """
        cls.TABLE_JSON_STRING = """
        {
            "name": "example_table",
            "comment": "This is an example table",
            "audit": {
                "creator": "anonymous",
                "createTime":"2025-10-10T00:00:00"
            },
            "columns": [
                {
                    "name": "id",
                    "type": "integer",
                    "comment": "id column comment",
                    "nullable": false,
                    "autoIncrement": true,
                    "defaultValue": {
                        "type": "literal",
                        "dataType": "integer",
                        "value": "-1"
                    }
                },
                {
                    "name": "name",
                    "type": "varchar(500)",
                    "comment": "name column comment",
                    "nullable": true,
                    "autoIncrement": false,
                    "defaultValue": {
                        "type": "literal",
                        "dataType": "null",
                        "value": "null"
                    }
                },
                {
                    "name": "StartingDate",
                    "type": "timestamp",
                    "comment": "StartingDate column comment",
                    "nullable": false,
                    "autoIncrement": false,
                    "defaultValue": {
                        "type": "function",
                        "funcName": "current_timestamp",
                        "funcArgs": []
                    }
                },
                {
                    "name": "info",
                    "type": {
                        "type": "struct",
                        "fields": [
                            {
                                "name": "position",
                                "type": "string",
                                "nullable": true,
                                "comment": "position field comment"
                            },
                            {
                                "name": "contact",
                                "type": {
                                "type": "list",
                                "elementType": "integer",
                                "containsNull": false
                                },
                                "nullable": true,
                                "comment": "contact field comment"
                            },
                            {
                                "name": "rating",
                                "type": {
                                "type": "map",
                                "keyType": "string",
                                "valueType": "integer",
                                "valueContainsNull": false
                                },
                                "nullable": true,
                                "comment": "rating field comment"
                            }
                        ]
                    },
                    "comment": "info column comment",
                    "nullable": true
                },
                {
                    "name": "dt",
                    "type": "date",
                    "comment": "dt column comment",
                    "nullable": true
                }
            ],
            "partitioning": [
                {
                    "strategy": "identity",
                    "fieldName": [ "dt" ]
                }
            ],
            "distribution": {
                "strategy": "hash",
                "number": 32,
                "funcArgs": [
                    {
                        "type": "field",
                        "fieldName": [ "id" ]
                    }
                ]
            },
            "sortOrders": [
                {
                    "sortTerm": {
                        "type": "field",
                        "fieldName": [ "id" ]
                    },
                    "direction": "asc",
                    "nullOrdering": "nulls_first"
                }
            ],
            "indexes": [
                {
                    "indexType": "primary_key",
                    "name": "PRIMARY",
                    "fieldNames": [["id"]]
                }
            ],
            "properties": {
                "format": "ORC"
            }
        }
        """

    def test_file_location_response(self):
        json_data = {"code": 0, "fileLocation": "file:/test/1"}
        json_str = _json.dumps(json_data)
        file_location_resp: FileLocationResponse = FileLocationResponse.from_json(
            json_str
        )
        self.assertEqual(file_location_resp.file_location(), "file:/test/1")
        file_location_resp.validate()

    def test_file_location_response_exception(self):
        json_data = {"code": 0, "fileLocation": ""}
        json_str = _json.dumps(json_data)
        file_location_resp: FileLocationResponse = FileLocationResponse.from_json(
            json_str
        )
        with self.assertRaises(IllegalArgumentException):
            file_location_resp.validate()

    def test_credential_response(self):
        json_data = {"code": 0, "credentials": []}
        json_str = _json.dumps(json_data)
        credential_resp: CredentialResponse = CredentialResponse.from_json(json_str)
        self.assertEqual(0, len(credential_resp.credentials()))
        credential_resp.validate()

        json_data = {
            "code": 0,
            "credentials": [
                {
                    "credentialType": "s3-token",
                    "expireTimeInMs": 1000,
                    "credentialInfo": {
                        "s3-access-key-id": "access-id",
                        "s3-secret-access-key": "secret-key",
                        "s3-session-token": "token",
                    },
                }
            ],
        }
        json_str = _json.dumps(json_data)
        credential_resp: CredentialResponse = CredentialResponse.from_json(json_str)
        credential_resp.validate()
        self.assertEqual(1, len(credential_resp.credentials()))
        credential = credential_resp.credentials()[0]
        self.assertEqual("s3-token", credential.credential_type())
        self.assertEqual(1000, credential.expire_time_in_ms())
        self.assertEqual("access-id", credential.credential_info()["s3-access-key-id"])
        self.assertEqual(
            "secret-key", credential.credential_info()["s3-secret-access-key"]
        )
        self.assertEqual("token", credential.credential_info()["s3-session-token"])

    def test_model_response(self):
        json_data = {
            "code": 0,
            "model": {
                "name": "test_model",
                "comment": "test comment",
                "properties": {"key1": "value1"},
                "latestVersion": 0,
                "audit": {
                    "creator": "anonymous",
                    "createTime": "2024-04-05T10:10:35.218Z",
                },
            },
        }
        json_str = _json.dumps(json_data)
        model_resp: ModelResponse = ModelResponse.from_json(
            json_str, infer_missing=True
        )
        model_resp.validate()
        self.assertEqual("test_model", model_resp.model().name())
        self.assertEqual(0, model_resp.model().latest_version())
        self.assertEqual("test comment", model_resp.model().comment())
        self.assertEqual({"key1": "value1"}, model_resp.model().properties())
        self.assertEqual("anonymous", model_resp.model().audit_info().creator())
        self.assertEqual(
            "2024-04-05T10:10:35.218Z", model_resp.model().audit_info().create_time()
        )

        json_data_missing = {
            "code": 0,
            "model": {
                "name": "test_model",
                "latestVersion": 0,
                "audit": {
                    "creator": "anonymous",
                    "createTime": "2024-04-05T10:10:35.218Z",
                },
            },
        }
        json_str_missing = _json.dumps(json_data_missing)
        model_resp_missing: ModelResponse = ModelResponse.from_json(
            json_str_missing, infer_missing=True
        )
        model_resp_missing.validate()
        self.assertEqual("test_model", model_resp_missing.model().name())
        self.assertEqual(0, model_resp_missing.model().latest_version())
        self.assertIsNone(model_resp_missing.model().comment())
        self.assertIsNone(model_resp_missing.model().properties())

    def test_model_version_list_response(self):
        json_data = {"code": 0, "versions": [0, 1, 2]}
        json_str = _json.dumps(json_data)
        resp: ModelVersionListResponse = ModelVersionListResponse.from_json(
            json_str, infer_missing=True
        )
        resp.validate()
        self.assertEqual(3, len(resp.versions()))
        self.assertEqual([0, 1, 2], resp.versions())

        json_data_missing = {"code": 0, "versions": []}
        json_str_missing = _json.dumps(json_data_missing)
        resp_missing: ModelVersionListResponse = ModelVersionListResponse.from_json(
            json_str_missing, infer_missing=True
        )
        resp_missing.validate()
        self.assertEqual(0, len(resp_missing.versions()))
        self.assertEqual([], resp_missing.versions())

        json_data_missing_1 = {
            "code": 0,
        }
        json_str_missing_1 = _json.dumps(json_data_missing_1)
        resp_missing_1: ModelVersionListResponse = ModelVersionListResponse.from_json(
            json_str_missing_1, infer_missing=True
        )
        self.assertRaises(IllegalArgumentException, resp_missing_1.validate)

    def test_model_version_response(self):
        json_data = {
            "code": 0,
            "modelVersion": {
                "version": 0,
                "aliases": ["alias1", "alias2"],
                "uris": {"unknown": "http://localhost:8080"},
                "comment": "test comment",
                "properties": {"key1": "value1"},
                "audit": {
                    "creator": "anonymous",
                    "createTime": "2024-04-05T10:10:35.218Z",
                },
            },
        }
        json_str = _json.dumps(json_data)
        resp: ModelVersionResponse = ModelVersionResponse.from_json(
            json_str, infer_missing=True
        )
        resp.validate()
        self.assertEqual(0, resp.model_version().version())
        self.assertEqual(["alias1", "alias2"], resp.model_version().aliases())
        self.assertEqual("test comment", resp.model_version().comment())
        self.assertEqual({"key1": "value1"}, resp.model_version().properties())
        self.assertEqual("anonymous", resp.model_version().audit_info().creator())
        self.assertEqual(
            "2024-04-05T10:10:35.218Z", resp.model_version().audit_info().create_time()
        )

        json_data = {
            "code": 0,
            "modelVersion": {
                "version": 0,
                "uris": {"unknown": "http://localhost:8080"},
                "audit": {
                    "creator": "anonymous",
                    "createTime": "2024-04-05T10:10:35.218Z",
                },
            },
        }
        json_str = _json.dumps(json_data)
        resp: ModelVersionResponse = ModelVersionResponse.from_json(
            json_str, infer_missing=True
        )
        resp.validate()
        self.assertEqual(0, resp.model_version().version())
        self.assertIsNone(resp.model_version().aliases())
        self.assertIsNone(resp.model_version().comment())
        self.assertIsNone(resp.model_version().properties())

        json_data = {
            "code": 0,
            "modelVersion": {
                "uris": {"unknown": "http://localhost:8080"},
                "audit": {
                    "creator": "anonymous",
                    "createTime": "2024-04-05T10:10:35.218Z",
                },
            },
        }
        json_str = _json.dumps(json_data)
        resp: ModelVersionResponse = ModelVersionResponse.from_json(
            json_str, infer_missing=True
        )
        self.assertRaises(IllegalArgumentException, resp.validate)

        json_data = {
            "code": 0,
            "modelVersion": {
                "version": 0,
                "audit": {
                    "creator": "anonymous",
                    "createTime": "2024-04-05T10:10:35.218Z",
                },
            },
        }
        json_str = _json.dumps(json_data)
        resp: ModelVersionResponse = ModelVersionResponse.from_json(
            json_str, infer_missing=True
        )
        self.assertRaises(IllegalArgumentException, resp.validate)

        json_data = {
            "code": 0,
            "modelVersion": {
                "version": 0,
                "uris": {"unknown": "http://localhost:8080"},
            },
        }
        json_str = _json.dumps(json_data)
        resp: ModelVersionResponse = ModelVersionResponse.from_json(
            json_str, infer_missing=True
        )
        self.assertRaises(IllegalArgumentException, resp.validate)

    def test_model_version_uri_response(self):
        json_data = {"code": 0, "uri": "s3://path/to/model"}
        json_str = _json.dumps(json_data)
        resp: ModelVersionUriResponse = ModelVersionUriResponse.from_json(
            json_str, infer_missing=True
        )
        resp.validate()
        self.assertEqual("s3://path/to/model", resp.uri())

        json_data_missing = {"code": 0, "uri": ""}
        json_str_missing = _json.dumps(json_data_missing)
        resp_missing: ModelVersionUriResponse = ModelVersionUriResponse.from_json(
            json_str_missing, infer_missing=True
        )
        self.assertRaises(IllegalArgumentException, resp_missing.validate)

        json_data_missing_1 = {
            "code": 0,
        }
        json_str_missing_1 = _json.dumps(json_data_missing_1)
        resp_missing_1: ModelVersionUriResponse = ModelVersionUriResponse.from_json(
            json_str_missing_1, infer_missing=True
        )
        self.assertRaises(IllegalArgumentException, resp_missing_1.validate)

    def test_partition_name_list_response(self):
        partition_names = [f"partition_{i}" for i in range(3)]
        json_data = {"code": 0, "names": partition_names}
        json_str = _json.dumps(json_data)
        resp: PartitionNameListResponse = PartitionNameListResponse.from_json(json_str)
        self.assertListEqual(resp.partition_names(), partition_names)
        resp.validate()

    def test_partition_name_list_response_exception(self):
        json_data = {"code": 0, "names": None}
        json_str = _json.dumps(json_data)
        resp: PartitionNameListResponse = PartitionNameListResponse.from_json(json_str)
        with self.assertRaises(IllegalArgumentException):
            resp.validate()

    def test_partition_response(self):
        json_string = f"""
        {{
            "code": 0,
            "partition": {TestResponses.PARTITION_JSON_STRING}
        }}
        """
        partition = PartitionDTOSerdes.deserialize(
            _json.loads(TestResponses.PARTITION_JSON_STRING)
        )
        resp: PartitionResponse = PartitionResponse.from_json(json_string)
        resp.validate()
        self.assertEqual(resp.get_partition(), partition)

    def test_partition_list_response(self):
        json_string = f"""
        {{
            "code": 0,
            "partitions": [{TestResponses.PARTITION_JSON_STRING}]
        }}
        """
        partitions = [
            PartitionDTOSerdes.deserialize(
                _json.loads(TestResponses.PARTITION_JSON_STRING)
            )
        ]
        resp: PartitionListResponse = PartitionListResponse.from_json(json_string)
        resp.validate()
        self.assertListEqual(resp.get_partitions(), partitions)

    def test_table_response(self):
        json_string = f"""
        {{
            "code": 0,
            "table": {TestResponses.TABLE_JSON_STRING}
        }}
        """
        resp: TableResponse = TableResponse.from_json(json_string)
        resp.validate()

    def test_function_response(self):
        """Test FunctionResponse."""
        json_str = """
            {
                "code": 0,
                "function": {
                    "name": "func1",
                    "functionType": "SCALAR",
                    "deterministic": true,
                    "definitions": [
                        {
                            "parameters": [],
                            "returnType": "integer",
                            "impls": [
                                {
                                    "language": "SQL",
                                    "runtime": "SPARK",
                                    "sql": "SELECT 1"
                                }
                            ]
                        }
                    ],
                    "comment": "comment",
                    "audit": {
                        "creator": "anonymous",
                        "createTime": "2024-04-05T10:10:35.218Z"
                    }
                }
            }
        """
        resp = FunctionResponse.from_json(json_str)
        resp.validate()
        self.assertEqual("func1", resp.function().name())
        self.assertEqual("SCALAR", resp.function().function_type().name)

        with self.assertRaises(IllegalArgumentException):
            FunctionResponse.from_json('{"code": 0, "function": null}').validate()

    def test_function_list_response(self):
        """Test FunctionListResponse."""
        json_str = """
            {
                "code": 0,
                "functions": [
                    {
                        "name": "func1",
                        "functionType": "SCALAR",
                        "deterministic": true,
                        "definitions": [
                            {
                                "parameters": [],
                                "returnType": "integer",
                                "impls": [
                                    {
                                        "language": "SQL",
                                        "runtime": "SPARK",
                                        "sql": "SELECT 1"
                                    }
                                ]
                            }
                        ],
                        "comment": "comment",
                        "audit": {
                            "creator": "anonymous",
                            "createTime": "2024-04-05T10:10:35.218Z"
                        }
                    }
                ]
            }
        """
        resp = FunctionListResponse.from_json(json_str)
        resp.validate()
        self.assertEqual(1, len(resp.functions()))
        self.assertEqual("func1", resp.functions()[0].name())

    def test_role_names_list_response(self) -> None:
        role_response = NameListResponse(0, ["role1", "role2", "role3"])
        role_response.validate()
        json_str = _json.dumps(
            {
                "code": 0,
                "names": ["role1", "role2", "role3"],
            }
        )

        ser_json = _json.dumps(role_response.to_dict())
        self.assertEqual(json_str, ser_json)
        deser_dict = _json.loads(ser_json)

        self.assertEqual(0, deser_dict.get("code"))
        self.assertIsNotNone(deser_dict.get("names"))
        self.assertEqual(3, len(deser_dict.get("names")))
        self.assertEqual("role1", deser_dict["names"][0])
        self.assertEqual("role2", deser_dict["names"][1])
        self.assertEqual("role3", deser_dict["names"][2])

    def test_role_names_list_response_validation(self) -> None:
        with self.assertRaises(IllegalArgumentException):
            role_response = NameListResponse(0, None)  # type: ignore
            role_response.validate()

        with self.assertRaises(IllegalArgumentException):
            role_response = NameListResponse(0, ["role1", ""])
            role_response.validate()
