import asyncio
import os
import tempfile
from asyncio import AbstractEventLoop
from collections import defaultdict
from pathlib import Path
from shutil import rmtree
from threading import Thread
from unittest import IsolatedAsyncioTestCase
from unittest.mock import AsyncMock, MagicMock, mock_open, patch

from aio_pika import IncomingMessage
from aio_pika.abc import AbstractRobustConnection

from amqp.adapter import file_handler
from amqp.adapter.file_handler import StreamingFileHandler


class TestFileHandler(IsolatedAsyncioTestCase):
    def test_ensure_directory_exists(self):
        """
        Test function for ensure_directory_exists.  Creates a temporary directory
        and files to test the versioning logic.
        """

        # Create a temporary directory
        file_handler = StreamingFileHandler()
        with tempfile.TemporaryDirectory() as temp_dir:
            # Test case 1: File does not exist
            filepath1 = os.path.join(temp_dir, "test1.txt")
            result1 = file_handler.ensure_directory_exists(filepath1)
            assert result1 == filepath1, f"Test Case 1 Failed: {result1} != {filepath1}"
            # Create the file
            with open(filepath1, "w") as f:
                f.write("test")

            # Test case 2: File exists, should create versioned file
            filepath2 = os.path.join(temp_dir, "test2.txt")
            result2 = file_handler.ensure_directory_exists(filepath2)
            assert result2 == filepath2, f"Test Case 2 Failed: {result2} != {filepath2}"
            with open(filepath2, "w") as f:
                f.write("test")
            result2_v1 = file_handler.ensure_directory_exists(filepath2)
            assert result2_v1 == os.path.join(
                temp_dir, "test2.1.txt"
            ), f"Test Case 2 Version 1 Failed: {result2_v1} != {os.path.join(temp_dir, 'test2.1.txt')}"

            # Create a few versioned files
            with open(os.path.join(temp_dir, "test3.txt"), "w") as f:
                f.write("test")
            with open(os.path.join(temp_dir, "test3.1.txt"), "w") as f:
                f.write("test")
            with open(os.path.join(temp_dir, "test3.2.txt"), "w") as f:
                f.write("test")
            filepath3 = os.path.join(temp_dir, "test3.txt")
            result3_v3 = file_handler.ensure_directory_exists(filepath3)
            assert result3_v3 == os.path.join(
                temp_dir, "test3.3.txt"
            ), f"Test Case 3 Version 3 Failed: {result3_v3} != {os.path.join(temp_dir, 'test3.3.txt')}"

            print("All test cases passed!")

    async def test_download_from_queue_to_file(self):
        loop = asyncio.new_event_loop()
        tempFile = tempfile.mktemp()
        file_handler = StreamingFileHandler()
        with patch.object(
            StreamingFileHandler, "_get_connection", new_callable=AsyncMock
        ) as mock_connect:
            mock_connection: AbstractRobustConnection = MagicMock(spec=AbstractRobustConnection)
            mock_connect.return_value = mock_connection
            mock_connection.channel = AsyncMock()
            mock_channel = mock_connection.channel.return_value
            mock_channel.get_queue = AsyncMock()
            mock_queue = mock_channel.get_queue.return_value
            mock_queue.iterator = MagicMock()
            mock_queue_iter = mock_queue.iterator().__aenter__.return_value
            mock_message: IncomingMessage = MagicMock(spec=IncomingMessage)
            mock_message.body = b"test_data"
            mock_message.body_size = len(b"test_data")
            mock_message.headers = {"eof": 1}
            mock_queue_iter.__aiter__.return_value = [mock_message]

            # Act
            result = await file_handler._download_from_queue_to_file(
                "test_rabbitmq",
                "test_queue",
                filepath=tempFile,
                loop=loop,
            )
            # clean up
            os.remove(tempFile)

            # Assert
            assert result == (len(b"test_data"), 1)
            mock_connect.assert_called_once_with("test_rabbitmq", loop=loop)
            mock_channel.get_queue.assert_called_once_with("test_queue", ensure=False)

    async def test_download_file_success(self):
        # Arrange
        file_handler = StreamingFileHandler()
        loop = asyncio.new_event_loop()
        loop_thread = Thread(target=loop.run_forever, daemon=True)
        loop_thread.start()
        rabbitmq_url = "amqp://test.rabbitmq"
        file_info = {
            "filename": "test_file.txt",
            "size": 1024,
            "streamName": "test_stream",
            "queue_name": "test_queue",
            "key": "test_key",
        }
        message_data = {}
        output_dir = tempfile.mkdtemp()

        file_handler.ensure_directory_exists = MagicMock(return_value=f"{output_dir}/test_file.txt")
        file_handler._download_from_queue_to_file = AsyncMock(return_value=(123, 1))

        # Act
        result = await file_handler.download_file(
            loop, rabbitmq_url, file_info, message_data, output_dir
        )

        # clean up
        loop.stop()
        rmtree(output_dir, ignore_errors=True)

        # Assert
        assert result == (123, 1)
        file_handler.ensure_directory_exists.assert_called_once_with(f"{output_dir}/test_file.txt")
        file_handler._download_from_queue_to_file.assert_called_once_with(
            "amqp://test.rabbitmq", "test_queue", f"{output_dir}/test_file.txt", loop
        )

    async def test_download_file_queue_not_found(self):
        # Arrange
        file_handler = StreamingFileHandler()
        loop = MagicMock(spec=AbstractEventLoop)
        rabbitmq_url = "amqp://test.rabbitmq"
        file_info = {
            "filename": "test_file.txt",
            "size": 1024,
            "streamName": "test_stream",
            "queue_name": None,
            "key": "test_key",
        }
        message_data = {}
        output_dir = "test_output"

        # Act
        result = await file_handler.download_file(
            loop, rabbitmq_url, file_info, message_data, output_dir
        )

        # Assert
        assert result == (0, 2)

    async def test_download_from_queue_to_buffer(self):
        loop = asyncio.new_event_loop()
        file_handler = StreamingFileHandler()
        with patch.object(
            StreamingFileHandler, "_get_connection", new_callable=AsyncMock
        ) as mock_connect:
            mock_connection: AbstractRobustConnection = MagicMock(spec=AbstractRobustConnection)
            mock_connect.return_value = mock_connection
            mock_connection.channel = AsyncMock()
            mock_channel = mock_connection.channel.return_value
            mock_channel.get_queue = AsyncMock()
            mock_queue = mock_channel.get_queue.return_value
            mock_queue.iterator = MagicMock()
            mock_queue_iter = mock_queue.iterator().__aenter__.return_value
            mock_message: IncomingMessage = MagicMock(spec=IncomingMessage)
            mock_message.body = b"test_data"
            mock_message.body_size = len(b"test_data")
            mock_message.headers = {"eof": 1}
            mock_queue_iter.__aiter__.return_value = [mock_message]

            # Act
            result = await file_handler._download_from_queue_to_buffer(
                "test_rabbitmq",
                "test_queue",
                loop=loop,
            )

            # Assert
            assert result == (b"test_data", 1)
            mock_connect.assert_called_once_with("test_rabbitmq", loop=loop)
            mock_channel.get_queue.assert_called_once_with("test_queue", ensure=False)

    async def test_download_buffer_success(self):
        # Arrange
        file_handler = StreamingFileHandler()
        loop = asyncio.new_event_loop()
        loop_thread = Thread(target=loop.run_forever, daemon=True)
        loop_thread.start()
        rabbitmq_url = "amqp://test.rabbitmq"

        file_handler._download_from_queue_to_buffer = AsyncMock(return_value=(b"hello", 1))

        # Act
        result = await file_handler.download_buffer(loop, rabbitmq_url, "test_queue")

        # clean up
        loop.stop()

        # Assert
        assert result == (b"hello", 1)
        file_handler._download_from_queue_to_buffer.assert_called_once_with(
            "amqp://test.rabbitmq", "test_queue", loop
        )

    async def test_download_buffer_queue_not_found(self):
        # Arrange
        file_handler = StreamingFileHandler()
        loop = MagicMock(spec=AbstractEventLoop)
        rabbitmq_url = "amqp://test.rabbitmq"

        # Act
        result = await file_handler.download_buffer(loop, rabbitmq_url, None)

        # Assert
        assert result == (0, 2)

    async def test_on_message_no_files_to_download(self):
        handler = StreamingFileHandler()
        message = MagicMock()
        message.body = b"{}"
        json_data = b'{"files": {}}'
        rabbitmq_url = "amqp://localhost"
        loop = AsyncMock()
        output_dir = "downloaded_files"

        handler.download_file = AsyncMock()
        handler.download_file.return_value = (123, 1)
        with patch("amqp.adapter.file_handler.logging_info"):
            result = await handler.on_message(message, json_data, rabbitmq_url, loop, output_dir)

        assert isinstance(result, defaultdict)
        self.assertEqual(result, defaultdict())

    async def test_on_message_with_files_to_download(self):
        handler = StreamingFileHandler()
        message = MagicMock()
        message.body = b'{"fileStream1": "queue1"}'
        json_data = b'{"files": {"key1": [{"streamName": "fileStream1"}]}}'
        rabbitmq_url = "amqp://localhost"
        loop = AsyncMock()
        output_dir = "downloaded_files"

        handler.download_file = AsyncMock()

        with patch("amqp.adapter.file_handler.logging_info") as mock_logging_info:
            result = await handler.on_message(message, json_data, rabbitmq_url, loop, output_dir)

        self.assertEqual(result, file_handler.download_locations)
        handler.download_file.assert_awaited_once()
        mock_logging_info.assert_any_call(
            f" [{message.delivery_tag}] about to create task for {{'streamName': 'fileStream1', 'queue_name': 'queue1', 'key': 'key1'}}"
        )
        mock_logging_info.assert_called_with(f"All files downloaded. d-tag: {message.delivery_tag}")

    async def test_on_message_handles_exception(self):
        handler = StreamingFileHandler()
        message = MagicMock()
        message.body = b"invalid_json"
        message.ack = AsyncMock()
        json_data = b"invalid_json"
        rabbitmq_url = "amqp://localhost"
        loop = AsyncMock()
        output_dir = "downloaded_files"

        with patch("amqp.adapter.file_handler.logging_error") as mock_logging_error:
            await handler.on_message(message, json_data, rabbitmq_url, loop, output_dir)

        mock_logging_error.assert_called()

    async def test_publish_file_file_not_found(self):
        file_handler = StreamingFileHandler()
        loop = asyncio.get_event_loop()
        exchange_mock = AsyncMock()
        file_path = Path("/path/to/nonexistent_file.txt")

        with self.assertRaises(FileNotFoundError):
            await file_handler.publish_file(
                exchange=exchange_mock,
                file_path=file_path,
                routing_key="test_routing_key",
                loop=loop,
            )

    async def test_publish_file_success(self):
        file_handler = StreamingFileHandler()
        loop = asyncio.get_event_loop()
        exchange_mock = AsyncMock()
        file_path = Path("/path/to/test_file.txt")
        data = b"test-data"
        with (
            patch("amqp.adapter.file_handler.Path.is_file", return_value=True),
            patch("amqp.adapter.file_handler.os.path.getsize", return_value=len(data)),
            patch("builtins.open", new=mock_open(read_data=data)) as mock_file,
        ):
            result = await file_handler.publish_file(
                exchange=exchange_mock,
                file_path=file_path,
                routing_key="test_routing_key",
                loop=loop,
                max_chunk_size=512,
            )

            mock_file.assert_called_with(file_path, "rb")
            self.assertEqual(result, ("test_routing_key", len(data)))
