# Copyright Contributors to the Amundsen project.
# SPDX-License-Identifier: Apache-2.0

import unittest
from unittest.mock import ANY

from databuilder.models.dashboard.dashboard_table import DashboardTable
from databuilder.models.graph_serializable import (
    RELATION_END_KEY, RELATION_END_LABEL, RELATION_REVERSE_TYPE, RELATION_START_KEY, RELATION_START_LABEL,
    RELATION_TYPE,
)
from databuilder.serializers import (
    mysql_serializer, neo4_serializer, neptune_serializer,
)
from databuilder.serializers.neptune_serializer import (
    METADATA_KEY_PROPERTY_NAME, NEPTUNE_CREATION_TYPE_JOB,
    NEPTUNE_CREATION_TYPE_RELATIONSHIP_PROPERTY_NAME_BULK_LOADER_FORMAT, NEPTUNE_HEADER_ID, NEPTUNE_HEADER_LABEL,
    NEPTUNE_LAST_EXTRACTED_AT_RELATIONSHIP_PROPERTY_NAME_BULK_LOADER_FORMAT, NEPTUNE_RELATIONSHIP_HEADER_FROM,
    NEPTUNE_RELATIONSHIP_HEADER_TO,
)


class TestDashboardTable(unittest.TestCase):
    def test_dashboard_table_nodes(self) -> None:
        dashboard_table = DashboardTable(table_ids=['hive://gold.schema/table1', 'hive://gold.schema/table2'],
                                         cluster='cluster_id', product='product_id',
                                         dashboard_id='dashboard_id', dashboard_group_id='dashboard_group_id')

        actual = dashboard_table.create_next_node()
        self.assertIsNone(actual)

    def test_dashboard_table_relations(self) -> None:
        dashboard_table = DashboardTable(table_ids=['hive://gold.schema/table1'],
                                         cluster='cluster_id', product='product_id',
                                         dashboard_id='dashboard_id', dashboard_group_id='dashboard_group_id')

        actual = dashboard_table.create_next_relation()
        actual_serialized = neo4_serializer.serialize_relationship(actual)
        actual_neptune_serialized = neptune_serializer.convert_relationship(actual)
        expected = {RELATION_END_KEY: 'hive://gold.schema/table1', RELATION_START_LABEL: 'Dashboard',
                    RELATION_END_LABEL: 'Table',
                    RELATION_START_KEY: 'product_id_dashboard://cluster_id.dashboard_group_id/dashboard_id',
                    RELATION_TYPE: 'DASHBOARD_WITH_TABLE',
                    RELATION_REVERSE_TYPE: 'TABLE_OF_DASHBOARD'}

        neptune_forward_expected = {
            NEPTUNE_HEADER_ID: "{label}:{from_vertex_id}_{to_vertex_id}".format(
                from_vertex_id='Dashboard:product_id_dashboard://cluster_id.dashboard_group_id/dashboard_id',
                to_vertex_id='Table:hive://gold.schema/table1',
                label='DASHBOARD_WITH_TABLE'
            ),
            METADATA_KEY_PROPERTY_NAME: "{label}:{from_vertex_id}_{to_vertex_id}".format(
                from_vertex_id='Dashboard:product_id_dashboard://cluster_id.dashboard_group_id/dashboard_id',
                to_vertex_id='Table:hive://gold.schema/table1',
                label='DASHBOARD_WITH_TABLE'
            ),
            NEPTUNE_RELATIONSHIP_HEADER_FROM:
                'Dashboard:product_id_dashboard://cluster_id.dashboard_group_id/dashboard_id',
            NEPTUNE_RELATIONSHIP_HEADER_TO: 'Table:hive://gold.schema/table1',
            NEPTUNE_HEADER_LABEL: 'DASHBOARD_WITH_TABLE',
            NEPTUNE_LAST_EXTRACTED_AT_RELATIONSHIP_PROPERTY_NAME_BULK_LOADER_FORMAT: ANY,
            NEPTUNE_CREATION_TYPE_RELATIONSHIP_PROPERTY_NAME_BULK_LOADER_FORMAT: NEPTUNE_CREATION_TYPE_JOB
        }

        neptune_reversed_expected = {
            NEPTUNE_HEADER_ID: "{label}:{from_vertex_id}_{to_vertex_id}".format(
                from_vertex_id='Table:hive://gold.schema/table1',
                to_vertex_id='Dashboard:product_id_dashboard://cluster_id.dashboard_group_id/dashboard_id',
                label='TABLE_OF_DASHBOARD'
            ),
            METADATA_KEY_PROPERTY_NAME: "{label}:{from_vertex_id}_{to_vertex_id}".format(
                from_vertex_id='Table:hive://gold.schema/table1',
                to_vertex_id='Dashboard:product_id_dashboard://cluster_id.dashboard_group_id/dashboard_id',
                label='TABLE_OF_DASHBOARD'
            ),
            NEPTUNE_RELATIONSHIP_HEADER_FROM: 'Table:hive://gold.schema/table1',
            NEPTUNE_RELATIONSHIP_HEADER_TO:
                'Dashboard:product_id_dashboard://cluster_id.dashboard_group_id/dashboard_id',
            NEPTUNE_HEADER_LABEL: 'TABLE_OF_DASHBOARD',
            NEPTUNE_LAST_EXTRACTED_AT_RELATIONSHIP_PROPERTY_NAME_BULK_LOADER_FORMAT: ANY,
            NEPTUNE_CREATION_TYPE_RELATIONSHIP_PROPERTY_NAME_BULK_LOADER_FORMAT: NEPTUNE_CREATION_TYPE_JOB
        }
        assert actual is not None
        self.assertDictEqual(actual_serialized, expected)
        self.assertDictEqual(actual_neptune_serialized[0], neptune_forward_expected)
        self.assertDictEqual(actual_neptune_serialized[1], neptune_reversed_expected)

    def test_dashboard_table_without_dot_as_name(self) -> None:
        dashboard_table = DashboardTable(table_ids=['bq-name://project-id.schema-name/table-name'],
                                         cluster='cluster_id', product='product_id',
                                         dashboard_id='dashboard_id', dashboard_group_id='dashboard_group_id')
        actual = dashboard_table.create_next_relation()
        actual_serialized = neo4_serializer.serialize_relationship(actual)
        expected = {RELATION_END_KEY: 'bq-name://project-id.schema-name/table-name', RELATION_START_LABEL: 'Dashboard',
                    RELATION_END_LABEL: 'Table',
                    RELATION_START_KEY: 'product_id_dashboard://cluster_id.dashboard_group_id/dashboard_id',
                    RELATION_TYPE: 'DASHBOARD_WITH_TABLE',
                    RELATION_REVERSE_TYPE: 'TABLE_OF_DASHBOARD'}
        assert actual is not None
        self.assertDictEqual(actual_serialized, expected)

    def test_dashboard_table_with_dot_as_name(self) -> None:
        dashboard_table = DashboardTable(table_ids=['bq-name://project.id.schema-name/table-name'],
                                         cluster='cluster_id', product='product_id',
                                         dashboard_id='dashboard_id', dashboard_group_id='dashboard_group_id')
        actual = dashboard_table.create_next_relation()
        self.assertIsNone(actual)

    def test_dashboard_table_with_slash_as_name(self) -> None:
        dashboard_table = DashboardTable(table_ids=['bq/name://project/id.schema/name/table/name'],
                                         cluster='cluster_id', product='product_id',
                                         dashboard_id='dashboard_id', dashboard_group_id='dashboard_group_id')
        actual = dashboard_table.create_next_relation()
        self.assertIsNone(actual)

    def test_dashboard_table_records(self) -> None:
        dashboard_table = DashboardTable(table_ids=['hive://gold.schema/table1', 'hive://gold.schema/table2'],
                                         cluster='cluster_id', product='product_id',
                                         dashboard_id='dashboard_id', dashboard_group_id='dashboard_group_id')
        actual1 = dashboard_table.create_next_record()
        actual1_serialized = mysql_serializer.serialize_record(actual1)
        expected1 = {
            'dashboard_rk': 'product_id_dashboard://cluster_id.dashboard_group_id/dashboard_id',
            'table_rk': 'hive://gold.schema/table1'
        }

        actual2 = dashboard_table.create_next_record()
        actual2_serialized = mysql_serializer.serialize_record(actual2)
        expected2 = {
            'dashboard_rk': 'product_id_dashboard://cluster_id.dashboard_group_id/dashboard_id',
            'table_rk': 'hive://gold.schema/table2'
        }

        assert actual1 is not None
        self.assertDictEqual(expected1, actual1_serialized)

        assert actual2 is not None
        self.assertDictEqual(expected2, actual2_serialized)
        self.assertIsNone(dashboard_table.create_next_record())
