#
# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
#

import copy
from unittest import mock
from uuid import UUID

import pytest
import yaml
from google.cloud import storage
from metadata_service.models.generated.ConnectorRegistryDestinationDefinition import ConnectorRegistryDestinationDefinition
from metadata_service.models.generated.ConnectorRegistrySourceDefinition import ConnectorRegistrySourceDefinition
from metadata_service.models.generated.ConnectorRegistryV0 import ConnectorRegistryV0
from orchestrator.assets import registry
from orchestrator.assets.registry_entry import (
    get_connector_type_from_registry_entry,
    get_registry_entry_write_path,
    get_registry_status_lists,
    metadata_to_registry_entry,
    safe_parse_metadata_definition,
)
from orchestrator.assets.registry_report import (
    all_destinations_dataframe,
    all_sources_dataframe,
    cloud_destinations_dataframe,
    cloud_sources_dataframe,
    oss_destinations_dataframe,
    oss_sources_dataframe,
)
from orchestrator.models.metadata import LatestMetadataEntry, MetadataDefinition
from orchestrator.utils.blob_helpers import yaml_blob_to_dict
from pydantic import ValidationError

VALID_METADATA_DICT = {
    "metadataSpecVersion": "1.0",
    "data": {
        "name": "Test",
        "definitionId": str(UUID(int=1)),
        "connectorType": "source",
        "dockerRepository": "test_repo",
        "dockerImageTag": "test_tag",
        "license": "test_license",
        "documentationUrl": "https://test_documentation_url.com",
        "githubIssueLabel": "test_label",
        "connectorSubtype": "api",
        "supportLevel": "community",
        "releaseStage": "alpha",
        "registryOverrides": {"oss": {"enabled": True}, "cloud": {"enabled": True}},
    },
}
INVALID_METADATA_DICT = {"invalid": "metadata"}


@pytest.mark.parametrize(
    "blob_name, blob_content, expected_result, expected_exception",
    [
        ("1.2.3/metadata.yaml", yaml.dump(VALID_METADATA_DICT), MetadataDefinition.parse_obj(VALID_METADATA_DICT), None),
        ("latest/metadata.yaml", yaml.dump(VALID_METADATA_DICT), MetadataDefinition.parse_obj(VALID_METADATA_DICT), None),
        ("1.2.3/metadata.yaml", yaml.dump(INVALID_METADATA_DICT), None, None),
        ("latest/metadata.yaml", yaml.dump(INVALID_METADATA_DICT), None, ValidationError),
    ],
)
def test_safe_parse_metadata_definition(blob_name, blob_content, expected_result, expected_exception):
    # Mock the Blob object
    mock_blob = mock.create_autospec(storage.Blob, instance=True)
    mock_blob.name = blob_name
    mock_blob.download_as_string.return_value = blob_content.encode("utf-8")

    metadata_dict = yaml_blob_to_dict(mock_blob)

    if expected_exception:
        with pytest.raises(expected_exception):
            safe_parse_metadata_definition(mock_blob.name, metadata_dict)
    else:
        result = safe_parse_metadata_definition(mock_blob.name, metadata_dict)
        # assert the name is set correctly
        assert result == expected_result


@pytest.mark.parametrize(
    "registries_data, expected_enabled, expected_disabled",
    [
        (
            {"oss": {"enabled": True}, "cloud": {"enabled": True}},
            ["oss", "cloud"],
            [],
        ),
        (
            {"oss": {"enabled": False}, "cloud": {"enabled": True}},
            ["cloud"],
            ["oss"],
        ),
        (
            {"oss": {"enabled": False}, "cloud": {"enabled": False}},
            [],
            ["oss", "cloud"],
        ),
        (
            {"oss": {"enabled": False}, "cloud": None},
            [],
            ["oss", "cloud"],
        ),
        (
            {"oss": {"enabled": False}},
            [],
            ["oss", "cloud"],
        ),
        (
            None,
            [],
            ["oss", "cloud"],
        ),
        (
            {"oss": {"enabled": True}},
            ["oss"],
            ["cloud"],
        ),
    ],
)
def test_get_registry_status_lists(registries_data, expected_enabled, expected_disabled):
    metadata_dict = {
        "metadataSpecVersion": "1.0",
        "data": {
            "name": "Test",
            "definitionId": str(UUID(int=1)),
            "connectorType": "source",
            "dockerRepository": "test_repo",
            "dockerImageTag": "test_tag",
            "license": "test_license",
            "documentationUrl": "https://test_documentation_url.com",
            "githubIssueLabel": "test_label",
            "connectorSubtype": "api",
            "releaseStage": "alpha",
            "supportLevel": "community",
            "registryOverrides": registries_data,
        },
    }
    metadata_definition = MetadataDefinition.parse_obj(metadata_dict)
    registry_entry = LatestMetadataEntry(metadata_definition=metadata_definition)
    enabled, disabled = get_registry_status_lists(registry_entry)
    assert set(enabled) == set(expected_enabled)
    assert set(disabled) == set(expected_disabled)


@pytest.mark.parametrize(
    "registry_entry, expected_type, expected_class",
    [
        ({"sourceDefinitionId": "abc"}, "source", ConnectorRegistrySourceDefinition),
        ({"destinationDefinitionId": "abc"}, "destination", ConnectorRegistryDestinationDefinition),
        ({}, None, None),
    ],
)
def test_get_connector_type_from_registry_entry(registry_entry, expected_type, expected_class):
    if expected_type and expected_class:
        connector_type, connector_class = get_connector_type_from_registry_entry(registry_entry)
        assert connector_type == expected_type
        assert connector_class == expected_class
    else:
        with pytest.raises(Exception) as e_info:
            get_connector_type_from_registry_entry(registry_entry)
        assert str(e_info.value) == "Could not determine connector type from registry entry"


def test_merged_registry_dataframes(oss_registry_dict, cloud_registry_dict):
    oss_registry_model = ConnectorRegistryV0.parse_obj(oss_registry_dict)
    cloud_registry_model = ConnectorRegistryV0.parse_obj(cloud_registry_dict)

    github_connector_folders = []
    num_oss_destinations = len(oss_registry_dict["destinations"])
    num_cloud_destinations = len(cloud_registry_dict["destinations"])

    num_cloud_sources = len(cloud_registry_dict["sources"])
    num_oss_sources = len(oss_registry_dict["sources"])

    cloud_destinations_df = cloud_destinations_dataframe(cloud_registry_model).value
    assert len(cloud_destinations_df) == num_cloud_destinations

    cloud_sources_df = cloud_sources_dataframe(cloud_registry_model).value
    assert len(cloud_sources_df) == num_cloud_sources

    oss_destinations_df = oss_destinations_dataframe(oss_registry_model).value
    assert len(oss_destinations_df) == num_oss_destinations

    oss_sources_df = oss_sources_dataframe(oss_registry_model).value
    assert len(oss_sources_df) == num_oss_sources

    all_sources_df = all_sources_dataframe(cloud_sources_df, oss_sources_df, github_connector_folders)
    all_destinations_df = all_destinations_dataframe(cloud_destinations_df, oss_destinations_df, github_connector_folders)

    # assert that all_sources_df has a entry for each sourceDefinitionId in the cloud registry and oss registry
    oss_source_definition_ids = set([source["sourceDefinitionId"] for source in oss_registry_dict["sources"]])
    cloud_source_definition_ids = set([source["sourceDefinitionId"] for source in cloud_registry_dict["sources"]])
    all_source_definition_ids = set(all_sources_df["definitionId"])

    assert all_source_definition_ids == oss_source_definition_ids.union(cloud_source_definition_ids)

    # assert that all_destinations_df has a entry for each sourceDefinitionId in the cloud registry and oss registry
    oss_destination_definition_ids = set([destination["destinationDefinitionId"] for destination in oss_registry_dict["destinations"]])
    cloud_destination_definition_ids = set([destination["destinationDefinitionId"] for destination in cloud_registry_dict["destinations"]])
    all_destination_definition_ids = set(all_destinations_df["definitionId"])

    assert all_destination_definition_ids == oss_destination_definition_ids.union(cloud_destination_definition_ids)


@pytest.mark.parametrize(
    "registry_type, connector_type, expected_id_field",
    [
        ("cloud", "source", "sourceDefinitionId"),
        ("cloud", "destination", "destinationDefinitionId"),
        ("oss", "source", "sourceDefinitionId"),
        ("oss", "destination", "destinationDefinitionId"),
    ],
)
def test_definition_id_conversion(registry_type, connector_type, expected_id_field):
    """
    Test if the definitionId in the metadata is successfully converted to
    destinationDefinitionId or sourceDefinitionId in the registry entry.
    """
    metadata = {
        "data": {"connectorType": connector_type, "definitionId": "test-id", "registryOverrides": {registry_type: {"enabled": True}}}
    }

    mock_metadata_entry = mock.Mock()
    mock_metadata_entry.metadata_definition.dict.return_value = metadata
    mock_metadata_entry.icon_url = "test-icon-url"
    mock_metadata_entry.dependency_file_url = "test-dependency-file-url"
    result = metadata_to_registry_entry(mock_metadata_entry, registry_type)
    assert "definitionId" not in result
    assert "test-id" == result[expected_id_field]


def test_tombstone_custom_public_set():
    """
    Test if tombstone, custom and public are set correctly in the registry entry.
    """
    metadata = {"data": {"connectorType": "source", "definitionId": "test-id", "registryOverrides": {"oss": {"enabled": True}}}}

    mock_metadata_entry = mock.Mock()
    mock_metadata_entry.metadata_definition.dict.return_value = metadata
    mock_metadata_entry.icon_url = "test-icon-url"
    mock_metadata_entry.dependency_file_url = "test-dependency-file-url"

    result = metadata_to_registry_entry(mock_metadata_entry, "oss")
    assert result["tombstone"] is False
    assert result["custom"] is False
    assert result["public"] is True


def test_fields_deletion():
    """
    Test if registries, connectorType, and definitionId fields were deleted from the registry entry.
    """
    metadata = {"data": {"connectorType": "source", "definitionId": "test-id", "registryOverrides": {"oss": {"enabled": True}}}}

    mock_metadata_entry = mock.Mock()
    mock_metadata_entry.metadata_definition.dict.return_value = metadata
    mock_metadata_entry.icon_url = "test-icon-url"
    mock_metadata_entry.dependency_file_url = "test-dependency-file-url"

    result = metadata_to_registry_entry(mock_metadata_entry, "oss")
    assert "registryOverrides" not in result
    assert "connectorType" not in result
    assert "definitionId" not in result


@pytest.mark.parametrize(
    "registry_type, expected_docker_image_tag, expected_additional_field",
    [("cloud", "cloud_tag", "cloud_value"), ("oss", "oss_tag", "oss_value")],
)
def test_overrides_application(registry_type, expected_docker_image_tag, expected_additional_field):
    """
    Test if the overrides for cloud or oss are properly applied to the registry entry.
    """
    # Assuming 'overriddenField' is a field to be overridden in metadata for the sake of this test.
    metadata = {
        "data": {
            "connectorType": "source",
            "definitionId": "test-id",
            "dockerImageTag": "base_tag",
            "registryOverrides": {
                "oss": {"enabled": True, "dockerImageTag": "oss_tag", "additionalField": "oss_value"},
                "cloud": {"enabled": True, "dockerImageTag": "cloud_tag", "additionalField": "cloud_value"},
            },
        }
    }

    mock_metadata_entry = mock.Mock()
    mock_metadata_entry.metadata_definition.dict.return_value = metadata
    mock_metadata_entry.file_path = f"metadata/{expected_docker_image_tag}/metadata.yaml"
    mock_metadata_entry.icon_url = "test-icon-url"
    mock_metadata_entry.dependency_file_url = "test-dependency-file-url"

    registry_entry = metadata_to_registry_entry(mock_metadata_entry, registry_type)
    assert registry_entry["dockerImageTag"] == expected_docker_image_tag
    assert registry_entry["additionalField"] == expected_additional_field

    expected_write_path = f"metadata/{expected_docker_image_tag}/{registry_type}"
    assert get_registry_entry_write_path(registry_entry, mock_metadata_entry, registry_type) == expected_write_path


def test_source_type_extraction():
    """
    Test if sourceType is successfully extracted from connectorSubtype in the registry entry.
    """
    metadata = {
        "data": {
            "connectorType": "source",
            "connectorSubtype": "database",
            "definitionId": "test-id",
            "registryOverrides": {"oss": {"enabled": True}},
        }
    }

    mock_metadata_entry = mock.Mock()
    mock_metadata_entry.metadata_definition.dict.return_value = metadata
    mock_metadata_entry.icon_url = "test-icon-url"
    mock_metadata_entry.dependency_file_url = "test-dependency-file-url"

    result = metadata_to_registry_entry(mock_metadata_entry, "oss")
    assert result["sourceType"] == "database"


def test_erd_url():
    """
    Test that if when defined in the metadata, the erd_url will be populated in the registry
    """
    mock_metadata_entry = mock.Mock()
    mock_metadata_entry.metadata_definition.dict.return_value = {
        "data": {
            "connectorType": "source",
            "definitionId": "test-id",
            "registryOverrides": {"oss": {"enabled": True}},
            "erdUrl": "https://an-erd.com",
        }
    }
    mock_metadata_entry.icon_url = "test-icon-url"
    mock_metadata_entry.dependency_file_url = "test-dependency-file-url"

    result = metadata_to_registry_entry(mock_metadata_entry, "oss")
    assert result["erdUrl"] == "https://an-erd.com"


def test_sbom_url():
    """
    Test that if when defined in the metadata, the sbom url will be populated in the registry
    """
    mock_metadata_entry = mock.Mock()
    mock_metadata_entry.metadata_definition.dict.return_value = {
        "data": {
            "connectorType": "source",
            "definitionId": "test-id",
            "registryOverrides": {"oss": {"enabled": True}},
            "dockerRepository": "test-connector",
            "dockerImageTag": "test-tag",
            "generated": {"sbomUrl": "test-sbom-url"},
        }
    }
    mock_metadata_entry.icon_url = "test-icon-url"
    mock_metadata_entry.dependency_file_url = "test-dependency-file-url"

    result = metadata_to_registry_entry(mock_metadata_entry, "oss")
    assert result["generated"]["sbomUrl"] == "test-sbom-url"


def test_support_level_default():
    """
    Test if supportLevel is defaulted to alpha in the registry entry.
    """
    metadata = {"data": {"connectorType": "source", "definitionId": "test-id", "registryOverrides": {"oss": {"enabled": True}}}}

    mock_metadata_entry = mock.Mock()
    mock_metadata_entry.metadata_definition.dict.return_value = metadata
    mock_metadata_entry.icon_url = "test-icon-url"
    mock_metadata_entry.dependency_file_url = "test-dependency-file-url"

    result = metadata_to_registry_entry(mock_metadata_entry, "oss")
    assert result["supportLevel"] == "community"


def test_migration_documentation_url_default():
    """
    Test if migrationDocumentationUrl is successfully defaulted in releases.migrationDocumentationUrl in the registry entry.
    """
    metadata = {
        "data": {
            "connectorType": "source",
            "definitionId": "test-id",
            "documentationUrl": "test-doc-url",
            "registryOverrides": {"oss": {"enabled": True}},
            "releases": {"migrationDocumentationUrl": None, "breakingChanges": {"1.0.0": {"migrationDocumentationUrl": None}}},
        }
    }

    mock_metadata_entry = mock.Mock()
    mock_metadata_entry.metadata_definition.dict.return_value = metadata
    mock_metadata_entry.icon_url = "test-icon-url"
    mock_metadata_entry.dependency_file_url = "test-dependency-file-url"

    expected_top_migration_documentation_url = "test-doc-url-migrations"
    expected_version_migration_documentation_url = "test-doc-url-migrations#1.0.0"

    result = metadata_to_registry_entry(mock_metadata_entry, "oss")
    assert result["releases"]["migrationDocumentationUrl"] == expected_top_migration_documentation_url
    assert result["releases"]["breakingChanges"]["1.0.0"]["migrationDocumentationUrl"] == expected_version_migration_documentation_url


def test_breaking_changes_migration_documentation_url():
    """
    Test if migrationDocumentationUrl is successfully defaulted for all entries in releases.breakingChanges including the key as the version.
    """
    metadata = {
        "data": {
            "connectorType": "source",
            "definitionId": "test-id",
            "documentationUrl": "test-doc-url",
            "registryOverrides": {"oss": {"enabled": True}},
            "releases": {
                "migrationDocumentationUrl": "test-migration-doc-url",
                "breakingChanges": {"1.0.0": {"migrationDocumentationUrl": "test-migration-doc-url-version"}},
            },
        }
    }

    mock_metadata_entry = mock.Mock()
    mock_metadata_entry.metadata_definition.dict.return_value = metadata
    mock_metadata_entry.icon_url = "test-icon-url"
    mock_metadata_entry.dependency_file_url = "test-dependency-file-url"

    result = metadata_to_registry_entry(mock_metadata_entry, "oss")
    assert result["releases"]["migrationDocumentationUrl"] == "test-migration-doc-url"
    assert result["releases"]["breakingChanges"]["1.0.0"]["migrationDocumentationUrl"] == "test-migration-doc-url-version"


def test_icon_url():
    """
    Test if the iconUrl in the metadata entry is correctly set in the registry entry.
    """
    metadata = {"data": {"connectorType": "source", "definitionId": "test-id", "registryOverrides": {"oss": {"enabled": True}}}}

    mock_metadata_entry = mock.Mock()
    mock_metadata_entry.metadata_definition.dict.return_value = metadata
    mock_metadata_entry.icon_url = "test-icon-url"
    mock_metadata_entry.dependency_file_url = "test-dependency-file-url"

    result = metadata_to_registry_entry(mock_metadata_entry, "oss")
    assert result["iconUrl"] == "test-icon-url"


def test_set_language_from_tags():
    """
    Test if the language is set from the tags in the registry entry.
    """
    metadata = {
        "data": {
            "connectorType": "source",
            "definitionId": "test-id",
            "registryOverrides": {"oss": {"enabled": True}},
            "tags": ["language:manifest-only", "key:value"],
        }
    }

    mock_metadata_entry = mock.Mock()
    mock_metadata_entry.metadata_definition.dict.return_value = metadata
    mock_metadata_entry.dependency_file_url = "test-dependency-file-url"

    result = metadata_to_registry_entry(mock_metadata_entry, "oss")
    assert result["language"] == "manifest-only"


def test_language_from_tags_does_not_override_top_level_language():
    """
    Test if the language is not set from the tags if it is already set in the registry entry.
    """
    metadata = {
        "data": {
            "connectorType": "source",
            "definitionId": "test-id",
            "registryOverrides": {"oss": {"enabled": True}},
            "language": "python",
            "tags": ["language:manifest-only", "key:value"],
        }
    }

    mock_metadata_entry = mock.Mock()
    mock_metadata_entry.metadata_definition.dict.return_value = metadata
    mock_metadata_entry.dependency_file_url = "test-dependency-file-url"
    result = metadata_to_registry_entry(mock_metadata_entry, "oss")
    assert result["language"] == "python"


def test_apply_release_candidates_with_older_rc(mocker):
    uuid = UUID(int=1)
    latest_registry_entry = {
        "name": "source-test",
        "sourceDefinitionId": str(uuid),
        "dockerRepository": "test-repo",
        "documentationUrl": "https://test_documentation_url.com",
        "spec": {},
        "dockerImageTag": "1.1.0",
        "releases": {},
    }
    rc_registry_entry = ConnectorRegistrySourceDefinition.parse_obj(
        {
            "name": "source-test",
            "sourceDefinitionId": str(uuid),
            "dockerRepository": "test-repo",
            "documentationUrl": "https://test_documentation_url.com",
            "spec": {},
            "dockerImageTag": "1.1.0-rc.1",
            "releases": {"rolloutConfiguration": {"enableProgressiveRollout": True}},
        }
    )

    result = registry.apply_release_candidates(latest_registry_entry, rc_registry_entry)
    assert not result["releases"]


def test_apply_release_candidates_newer_version(mocker):
    uuid = UUID(int=1)
    latest_registry_entry = {
        "name": "source-test",
        "sourceDefinitionId": str(uuid),
        "dockerRepository": "test-repo",
        "documentationUrl": "https://test_documentation_url.com",
        "spec": {},
        "dockerImageTag": "1.0.0",
        "releases": {},
    }
    rc_registry_entry = ConnectorRegistrySourceDefinition.parse_obj(
        {
            "name": "source-test",
            "sourceDefinitionId": str(uuid),
            "dockerRepository": "test-repo",
            "documentationUrl": "https://test_documentation_url.com",
            "spec": {},
            "dockerImageTag": "1.1.0-rc.1",
            "releases": {"rolloutConfiguration": {"enableProgressiveRollout": True}},
        }
    )
    result = registry.apply_release_candidates(latest_registry_entry, rc_registry_entry)
    assert "1.1.0-rc.1" in result["releases"]["releaseCandidates"]
