# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License
import json
import time
import unittest
import uuid
from unittest import mock
from uuid import uuid4

from azure.storage.queue import QueueMessage, QueueClient

from azure.kusto.ingest import QueuedIngestClient, IngestionResult, IngestionStatus
from azure.kusto.ingest._resource_manager import _ResourceUri
from azure.kusto.ingest.status import KustoIngestStatusQueues, SuccessMessage, FailureMessage

SAS = "sp=rl&st=2020-05-20T13:38:37Z&se=2020-05-21T13:38:37Z&sv=2019-10-10&sr=c&sig=xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
ENDPOINT_SUFFIX = "core.windows.net"


def mock_message(success):
    m = QueueMessage()
    m.id = uuid4()
    m.insertion_time = time.time()
    m.expiration_time = None
    m.dequeue_count = None

    if success:
        content = {
            "OperationId": str(m.id),
            "Database": "db1",
            "Table": "table1",
            "IngestionSourceId": str(m.id),
            "IngestionSourcePath": "blob/path",
            "RootActivityId": "1",
            "SucceededOn": time.time(),
        }
    else:
        content = {
            "OperationId": str(m.id),
            "Database": "db1",
            "Table": "table1",
            "IngestionSourceId": str(m.id),
            "IngestionSourcePath": "blob/path",
            "RootActivityId": "1",
            "FailedOn": time.time(),
            "Details": "",
            "ErrorCode": "1",
            "FailureStatus": "",
            "OriginatesFromUpdatePolicy": "",
            "ShouldRetry": False,
        }

    m.content = json.dumps(content)
    m.pop_receipt = None
    m.time_next_visible = None

    return m


def fake_peek_factory(f):
    def fake_peek(self, max_messages):
        return f(self.queue_name, max_messages)

    return fake_peek


def fake_receive_factory(f):
    def fake_receive(self, messages_per_page, *args, **kwargs):
        return f(self.queue_name, messages_per_page)

    return fake_receive


def fake_delete_factory(f):
    def fake_delete(self, n):
        return f(self.queue_name, n)

    return fake_delete


def get_resource_uri(account: str, container: str) -> _ResourceUri:
    return _ResourceUri("https://{0}.{1}.{2}/{3}?{4}".format(account, "queue", ENDPOINT_SUFFIX, container, SAS))


class StatusQTests(unittest.TestCase):
    def test_init(self):
        client = QueuedIngestClient("some-cluster")
        qs = KustoIngestStatusQueues(client)

        assert qs.success.message_cls == SuccessMessage
        assert qs.failure.message_cls == FailureMessage

    def test_isempty(self):
        client = QueuedIngestClient("some-cluster")

        fake_peek = fake_peek_factory(
            lambda queue_name, num_messages=1: [mock_message(success=True) for _ in range(0, num_messages)] if "qs" in queue_name else []
        )
        with mock.patch.object(client._resource_manager, "get_successful_ingestions_queues") as mocked_get_success_qs, mock.patch.object(
            client._resource_manager, "get_failed_ingestions_queues"
        ) as mocked_get_failed_qs, mock.patch.object(QueueClient, "peek_messages", autospec=True, side_effect=fake_peek) as q_mock:
            fake_failed_queue = get_resource_uri("mocked_storage_account_f1", "mocked_qf_name")
            fake_success_queue = get_resource_uri("mocked_storage_account2", "mocked_qs_name")

            mocked_get_success_qs.return_value = [fake_success_queue]
            mocked_get_failed_qs.return_value = [fake_failed_queue]

            qs = KustoIngestStatusQueues(client)

            assert qs.success.is_empty() is False
            assert qs.failure.is_empty() is True

            assert q_mock.call_count == 2
            assert q_mock.call_args_list[0][1]["max_messages"] == 2
            assert q_mock.call_args_list[1][1]["max_messages"] == 2

    def test_peek(self):
        client = QueuedIngestClient("some-cluster")

        fake_peek = fake_peek_factory(
            lambda queue_name, num_messages=1: [
                mock_message(success=True) if "qs" in queue_name else mock_message(success=False) for _ in range(0, num_messages)
            ]
        )

        with mock.patch.object(client._resource_manager, "get_successful_ingestions_queues") as mocked_get_success_qs, mock.patch.object(
            client._resource_manager, "get_failed_ingestions_queues"
        ) as mocked_get_failed_qs, mock.patch.object(QueueClient, "peek_messages", autospec=True, side_effect=fake_peek) as q_mock:
            fake_failed_queue1 = get_resource_uri("mocked_storage_account_f1", "mocked_qf_name")
            fake_failed_queue2 = get_resource_uri("mocked_storage_account_f2", "mocked_qf_2_name")
            fake_success_queue = get_resource_uri("mocked_storage_account2", "mocked_qs_name")

            mocked_get_success_qs.return_value = [fake_success_queue]
            mocked_get_failed_qs.return_value = [fake_failed_queue1, fake_failed_queue2]

            qs = KustoIngestStatusQueues(client)

            peek_success_actual = qs.success.peek()
            peek_failure_actual = qs.failure.peek(6)

            assert len(peek_success_actual) == 1

            for m in peek_failure_actual:
                assert isinstance(m, FailureMessage) is True

            for m in peek_success_actual:
                assert isinstance(m, SuccessMessage) is True

            assert len(peek_failure_actual) == 6

            actual = {}

            assert len(QueueClient.peek_messages.call_args_list) == 3

            for call_args in q_mock.call_args_list:
                actual[call_args[0][0].queue_name] = actual.get(call_args[0][0].queue_name, 0) + call_args[1]["max_messages"]

            assert actual[fake_failed_queue2.object_name] == 4
            assert actual[fake_failed_queue1.object_name] == 4
            assert actual[fake_success_queue.object_name] == 2

    def test_pop(self):
        client = QueuedIngestClient("some-cluster")

        fake_receive = fake_receive_factory(
            lambda queue_name, num_messages=1: [
                mock_message(success=True) if "qs" in queue_name else mock_message(success=False) for _ in range(0, num_messages)
            ]
        )

        with mock.patch.object(client._resource_manager, "get_successful_ingestions_queues") as mocked_get_success_qs, mock.patch.object(
            client._resource_manager, "get_failed_ingestions_queues"
        ) as mocked_get_failed_qs, mock.patch.object(
            QueueClient,
            "receive_messages",
            autospec=True,
            side_effect=fake_receive,
        ) as q_receive_mock, mock.patch.object(
            QueueClient, "delete_message", return_value=None
        ) as q_del_mock:
            fake_failed_queue1 = get_resource_uri("mocked_storage_account_f1", "mocked_qf_name")
            fake_failed_queue2 = get_resource_uri("mocked_storage_account_f2", "mocked_qf_2_name")
            fake_success_queue = get_resource_uri("mocked_storage_account2", "mocked_qs_name")

            mocked_get_success_qs.return_value = [fake_success_queue]
            mocked_get_failed_qs.return_value = [fake_failed_queue1, fake_failed_queue2]

            qs = KustoIngestStatusQueues(client)

            get_success_actual = qs.success.pop()
            get_failure_actual = qs.failure.pop(6)

            assert len(get_success_actual) == 1
            assert len(get_failure_actual) == 6

            for m in get_failure_actual:
                assert isinstance(m, FailureMessage)

            for m in get_success_actual:
                assert isinstance(m, SuccessMessage)

            assert q_receive_mock.call_count == 3
            assert q_del_mock.call_count == len(get_success_actual) + len(get_failure_actual)

            assert q_receive_mock.call_args_list[0][1]["messages_per_page"] == 2

            actual = {
                q_receive_mock.call_args_list[1][0][0].queue_name: q_receive_mock.call_args_list[1][1]["messages_per_page"],
                q_receive_mock.call_args_list[2][0][0].queue_name: q_receive_mock.call_args_list[2][1]["messages_per_page"],
            }

            assert actual[fake_failed_queue2.object_name] == 4
            assert actual[fake_failed_queue1.object_name] == 4

        client.close()

    def test_pop_unbalanced_queues(self):
        client = QueuedIngestClient("some-cluster")

        fake_receive = fake_receive_factory(
            lambda queue_name, messages_per_page=1: [mock_message(success=False) for _ in range(0, messages_per_page)] if "1" in queue_name else []
        )
        with mock.patch.object(client._resource_manager, "get_successful_ingestions_queues"), mock.patch.object(
            client._resource_manager, "get_failed_ingestions_queues"
        ) as mocked_get_failed_qs, mock.patch.object(
            QueueClient,
            "receive_messages",
            autospec=True,
            side_effect=fake_receive,
        ) as q_receive_mock, mock.patch.object(
            QueueClient, "delete_message", return_value=None
        ):
            fake_failed_queue1 = get_resource_uri("mocked_storage_account_f1", "mocked_qf_1_name")
            fake_failed_queue2 = get_resource_uri("mocked_storage_account_f2", "mocked_qf_2_name")

            mocked_get_failed_qs.return_value = [fake_failed_queue1, fake_failed_queue2]

            qs = KustoIngestStatusQueues(client)

            get_failure_actual = qs.failure.pop(6)

            assert len(get_failure_actual) == 6

            for m in get_failure_actual:
                assert isinstance(m, FailureMessage)

            assert q_receive_mock.call_count == 3

            actual = {}

            for call_args in q_receive_mock.call_args_list:
                actual[call_args[0][0].queue_name] = actual.get(call_args[0][0].queue_name, 0) + call_args[1]["messages_per_page"]

            assert actual[fake_failed_queue2.object_name] + actual[fake_failed_queue1.object_name] == (4 + 4 + 6)

    def test_status_sanity(self):
        i = IngestionResult(IngestionStatus.QUEUED, "db", "table", uuid.UUID("11111111-1111-1111-1111-111111111111"), "path")
        assert (
            repr(i)
            == "IngestionResult(status=IngestionStatus.QUEUED, database=db, table=table, source_id=11111111-1111-1111-1111-111111111111, obfuscated_blob_uri=path)"
        )

        no_path = IngestionResult(IngestionStatus.QUEUED, "db", "table", uuid.UUID("11111111-1111-1111-1111-111111111111"))
        assert repr(no_path) == "IngestionResult(status=IngestionStatus.QUEUED, database=db, table=table, source_id=11111111-1111-1111-1111-111111111111)"
