import asyncio
import base64
from tempfile import NamedTemporaryFile
from threading import Thread
from unittest.mock import PropertyMock, create_autospec

import pytest
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa

from src.backend.config import SNOWFLAKE_CREDENTIALS
from src.backend.connectors import snowflake_db
from src.backend.connectors.snowflake_db import client

SnowflakeConnection = client.sf_connector.SnowflakeConnection

CONCURRENCY_COUNT: int = 10


class TestSnowflake:
    _class = snowflake_db.Client  # Test __init__.py

    @pytest.fixture
    def instance(self):
        """Create a Snowflake instance with mock credentials."""
        return self._class(
            user="mock_user",
            account="mock_account",
            warehouse="mock_warehouse",
            database="mock_database",
            schema="mock_schema",
            private_key="mock_private_key",
        )

    @pytest.fixture
    def instance_real(self):
        """Create a Snowflake instance with real credentials,
        using the FACTS database and PROD schema for testing.
        """
        return self._class(
            **SNOWFLAKE_CREDENTIALS,
            database="FACTS",
            schema="PROD",
        )

    @pytest.fixture
    def instance_mocked_connection(self, instance, mock_connection):
        """Return an instance with a mocked connection."""
        instance.connection = mock_connection
        return instance

    @pytest.fixture
    def mock_connection(self):
        """Return a mock connection."""
        return create_autospec(SnowflakeConnection)

    @pytest.fixture
    def patcher_connect(self, mocker, mock_connection):
        return mocker.patch.object(
            client.sf_connector, "connect", return_value=mock_connection
        )

    @pytest.fixture
    def patcher_private_key(self, mocker, instance):
        return mocker.patch.object(
            self._class, "private_key", new_callable=PropertyMock
        )

    def test_init(self, instance):
        assert instance.connection is None

    @pytest.mark.asyncio
    async def test_cursor(self, instance_mocked_connection):
        await instance_mocked_connection.cursor()
        assert instance_mocked_connection.connection.cursor.call_count == 1

    @pytest.mark.asyncio
    async def test_connect(self, instance, patcher_connect, patcher_private_key):
        connections = (instance.connect() for _ in range(CONCURRENCY_COUNT))
        await asyncio.gather(*connections)

        assert patcher_private_key.call_count == 1
        assert patcher_connect.call_count == 1
        assert not patcher_connect.call_args.kwargs.get(
            "insecure_mode"
        ), "Insecure mode should not be used in production."
        assert isinstance(
            instance.connection,
            SnowflakeConnection,
        ), "Connection was not correctly persisted."

    def test_disconnect(self, instance_mocked_connection):
        instance = instance_mocked_connection
        mock_connection = instance.connection

        # Create 10 threads to test concurrency.
        threads = [Thread(target=instance.disconnect) for _ in range(CONCURRENCY_COUNT)]

        for thread in threads:
            thread.start()

        for thread in threads:
            thread.join()

        assert mock_connection.close.call_count == 1  # noqa
        assert instance.connection is None, "Connection was not purged."

    class TestPrivateKey:
        def test_private_key(self, mocker, instance):
            pem_private, private_key = mock_private_key()
            mocker.patch.object(
                instance, "_get_private_key_content", return_value=pem_private
            )
            result = instance.private_key
            assert isinstance(result, bytes)
            assert result == private_key.private_bytes(
                encoding=client.serialization.Encoding.DER,
                format=client.serialization.PrivateFormat.PKCS8,
                encryption_algorithm=client.serialization.NoEncryption(),
            )
            assert instance.private_key == result  # Coverage of persistence.

        def test_private_key_file(self, instance):
            """Test that a mock private key is loaded correctly."""
            pem_private, private_key = mock_private_key()

            with NamedTemporaryFile() as mock_private_key_file:
                instance._private_key = mock_private_key_file.name
                mock_private_key_file.write(pem_private)
                mock_private_key_file.seek(0)

                result = instance.private_key

            assert result
            assert isinstance(result, bytes)

        def test_get_private_key_content_from_pk_as_string(self, instance):
            pem_private, _ = mock_private_key()
            pem_private_as_str = pem_private.decode("utf-8")
            result = instance._get_private_key_content(pem_private_as_str)
            assert isinstance(result, bytes)
            assert result == pem_private

        def test_get_private_key_content_from_pk_as_file(self, instance):
            pem_private, _ = mock_private_key()

            with NamedTemporaryFile() as mock_private_key_file:
                mock_private_key_file.write(pem_private)
                mock_private_key_file.seek(0)

                result = instance._get_private_key_content(mock_private_key_file.name)
            assert isinstance(result, bytes)
            assert result == pem_private

        def test_get_private_key_content_from_pk_as_base64(self, instance):
            pem_private, _ = mock_private_key()
            pem_private_as_base64 = base64.b64encode(pem_private).decode("utf-8")
            result = instance._get_private_key_content(pem_private_as_base64)
            assert isinstance(result, bytes)
            assert result == pem_private


def mock_private_key():
    """Return a mock private key for testing purposes."""
    private_key = rsa.generate_private_key(
        public_exponent=65537, key_size=2048, backend=default_backend()
    )

    pem_private = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.TraditionalOpenSSL,
        encryption_algorithm=serialization.NoEncryption(),
    )

    return pem_private, private_key