"""
implements the file downloader for the amqp adapter, where file is sent in chunks over RabbitMQ
"""

import asyncio
import json
import logging
import os
from asyncio import AbstractEventLoop
from collections import defaultdict
from pathlib import Path

from aio_pika import Message, connect_robust
from aio_pika.abc import (
    AbstractExchange,
    AbstractIncomingMessage,
    AbstractQueue,
    AbstractRobustChannel,
    DeliveryMode,
)

from amqp.adapter.logging_utils import logging_debug, logging_error, logging_info
from amqp.router.utils import await_result

# Global dictionary to track completed downloads
download_locations = defaultdict(str)
IN_PROGRESS = 0
END_OF_FILE = 1
CANCELLED = 2


class FileHandler:
    """
    The abstract base class for file handling. This is mainly to make the code testable, as there are only 2
    options for pushing (uploading) a file(s) to a service: as HTTP multipart message delivered via RabbitMQ
    or via stream in RabbitMQ. Downloading a file is always done via RabbitMQ stream. Thus the method signiture
    is tighly coupled to the RabbitMQ API in v1. Only alternative implementation is a mock object in unittests.
    """

    def __init__(self):
        pass

    def ensure_directory_exists(self, filepath) -> str:
        """
        Ensures that the directory for the given filepath exists, and handles file versioning.

        Args:
            filepath: The path to the file.

        Returns:
            The original filepath if the file does not exist, or a modified filepath
            pointing to the next available version of the file if it does.
        """
        # 1. Ensure directory exists
        directory = os.path.dirname(filepath)
        if directory and not os.path.exists(directory):
            os.makedirs(directory)

        # 2. Handle file versioning
        if not os.path.exists(filepath):
            return filepath  # File doesn't exist, return original path

        # File exists, find next available version
        base, ext = os.path.splitext(filepath)
        version = 1
        while True:
            new_filepath = f"{base}.{version}{ext}"
            if not os.path.exists(new_filepath):
                return new_filepath  # Found a free version, return this path
            version += 1

    async def download_buffer(self, loop, rabbitmq_url, queue_name) -> (bytearray, int):
        """
        Asynchronously downloads a bytearray from RabbitMQ based on the provided file information.
        It assumes entire array is in the first (and only) chunk
        Args:
            loop: The asyncio event loop.
            rabbitmq_url: The RabbitMQ connection URL.
            queue_name (str): The directory where the file should be saved.
        """
        raise NotImplementedError("Subclasses must implement this method.")

    async def download_file(
        self, loop, rabbitmq_url, file_info, message_data, output_dir
    ) -> (int, int):
        """
        Asynchronously downloads a file from RabbitMQ based on the provided file information.
        takes the file_info and message_data and output_dir, creates queue consumer and waits for the file
        to be downloaded. File may consist of multiple chunks, so the function will wait until all chunks are received.
        Args:
            loop: The asyncio event loop.
            rabbitmq_url: The RabbitMQ connection URL.
            file_info (dict): A dictionary containing file metadata
            message_data (dict): The entire message
            output_dir (str): The directory where the file should be saved.
        """
        raise NotImplementedError("Subclasses must implement this method.")

    async def on_message(
        self,
        message: AbstractIncomingMessage,
        json_data: bytes,
        rabbitmq_url: str,
        loop,
        output_dir="downloaded_files",
    ) -> defaultdict:
        """
        Asynchronous callback function for handling messages from the 'amq-cleverswarm' queue.

        Args:
            message: The incoming RabbitMQ message.
            rabbitmq_url: The RabbitMQ connection URL.
        """
        raise NotImplementedError("Subclasses must implement this method.")

    async def publish_file(
        self,
        exchange: AbstractExchange,
        file_path: Path,
        routing_key: str,
        loop: AbstractEventLoop,
        max_chunk_size: int = 2**20,  # 1MB default chunk size
        content_type: str = "application/octet-stream",  # Default content type
        delivery_mode: DeliveryMode = DeliveryMode.PERSISTENT,
    ) -> tuple[str, int]:
        """
        Reads a file from the file system and publishes its content to a RabbitMQ exchange
        in chunks.  It declares a queue with a random name, binds it to the specified
        exchange, and uses the queue name as the routing key.

        Args:
            exchange: The aio_pika RobustChannel to use for communication with RabbitMQ.
            file_path: The path to the file to publish.
            max_chunk_size: The maximum size of each chunk (in bytes).
            routing_key: The name of the RabbitMQ exchange routing key to publish to correct queue.
            loop: correct event loop
            max_chunk_size: chunk size to send
            content_type: The content type of the file data.
            delivery_mode: The delivery mode (e.g., PERSISTENT or TRANSIENT).

        Returns:
            The name of the randomly generated queue where the file content was published.
        """
        raise NotImplementedError("Subclasses must implement this method.")


# The StreamingFileHandler class is a subclass of FileHandler that implements the download_file method
class StreamingFileHandler(FileHandler):
    """
    Handles file downloads from RabbitMQ streams.
    """

    def __init__(self):
        super().__init__()

    # This method MUST NOT be static.
    # For testing purpose, we will mock this method and inject a mock connection.
    async def _get_connection(self, rabbitmq_url, loop):
        return await connect_robust(rabbitmq_url, loop=loop)

    async def _download_from_queue_to_file(
        self, rabbitmq_url, queue_name, filepath, loop
    ) -> tuple[int, int] | None:
        """
        Internal method for downloading a file from a queue to disk.
        Return (file size, eof status).
        """
        # Here the normal mock doesn't work since this method
        # is async and will be run in another loop,
        # out of the scope of method mock.
        connection = await self._get_connection(rabbitmq_url, loop=loop)
        _file_size = 0
        _eof = IN_PROGRESS
        try:
            channel: AbstractRobustChannel = await connection.channel()
            queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False)
            logging_debug(f"Awaiting all file chunks being downloaded for {filepath}")
            async with queue.iterator() as queue_iter:
                async for message in queue_iter:
                    async with message.process():
                        with open(filepath, "ab") as f:
                            f.write(message.body)
                        _file_size = _file_size + message.body_size
                        _eof = int(message.headers.get("eof", END_OF_FILE))
                        if _eof != IN_PROGRESS:
                            break
        except Exception as e:
            logging_error(f"File download failed: {e}")
        finally:
            await connection.close()

        return _file_size, _eof

    async def download_file(
        self, loop, rabbitmq_url, file_info, message_data, output_dir
    ) -> (int, int):
        """
        Asynchronously downloads a file from RabbitMQ based on the provided file information.
        takes the file_info and message_data and output_dir, creates queue consumer and waits for the file to be downloaded
        Args:
            loop: The asyncio event loop.
            rabbitmq_url: The RabbitMQ connection URL.
            file_info (dict): A dictionary containing file metadata
            message_data (dict): The entire message
            output_dir (str): The directory where the file should be saved.
        """
        filename = file_info["filename"]
        size = file_info["size"]
        stream_name = file_info["streamName"]
        queue_name = file_info["queue_name"]
        filepath = os.path.join(output_dir, filename)
        filepath = self.ensure_directory_exists(
            filepath
        )  # if filename exists, create a new (next) version
        download_locations[file_info["key"]] = filepath
        filename = os.path.basename(
            filepath
        )  # if a new version was created, ensure to use that new version
        file_info["filename"] = filename

        logging_debug(
            f"Downloading file: {filename} (size: {size} bytes) from stream: {stream_name}:{queue_name} to {filepath}"
        )

        if queue_name:
            _local_file_size, _local_eof = await await_result(
                asyncio.run_coroutine_threadsafe(
                    self._download_from_queue_to_file(rabbitmq_url, queue_name, filepath, loop),
                    loop=loop,
                )
            )
            return _local_file_size, _local_eof

        return 0, CANCELLED

    async def _download_from_queue_to_buffer(
        self, rabbitmq_url, queue_name, loop
    ) -> tuple[bytearray, int]:
        """
        Internal method for downloading a file from a queue to RAM.
        Return (bytearray, eof status).
        """
        _connection = await self._get_connection(rabbitmq_url, loop=loop)
        _eof = IN_PROGRESS
        _res = bytearray()
        try:
            channel: AbstractRobustChannel = await _connection.channel()
            queue: AbstractQueue = await channel.get_queue(queue_name, ensure=False)
            logging_debug(f"Awaiting all buffer chunks being downloaded for {queue_name}")
            async with queue.iterator() as queue_iter:
                async for message in queue_iter:
                    async with message.process():
                        _res.extend(message.body)
                        _eof = message.headers.get("eof", END_OF_FILE)
                        if _eof != IN_PROGRESS:
                            break
            logging_info(f"Finished downloading buffer {queue_name}")
        except Exception as e:
            logging_error(f"Downloading buffer: {e}")

        await _connection.close()
        return _res, _eof

    async def download_buffer(self, loop, rabbitmq_url, queue_name: str | None) -> (bytearray, int):
        """
        Asynchronously downloads a bytearay from RabbitMQ based on the provided file information.
        It assumes entire array is in the first (and only) chunk
        Args:
            loop: The asyncio event loop.
            rabbitmq_url: The RabbitMQ connection URL.
            queue_name (str): The directory where the file should be saved.
        """
        logging_debug(f"Downloading buffer from stream: {queue_name}")
        if queue_name:
            _local_res, _local_eof = await await_result(
                asyncio.run_coroutine_threadsafe(
                    self._download_from_queue_to_buffer(rabbitmq_url, queue_name, loop), loop=loop
                )
            )
            return _local_res, _local_eof
        return 0, CANCELLED

    async def on_message(
        self,
        message: AbstractIncomingMessage,
        json_data: bytes,
        rabbitmq_url: str,
        loop,
        output_dir="downloaded_files",
    ) -> defaultdict:
        """
        Asynchronous callback function for handling messages from the 'amq-cleverswarm' queue.

        Args:
            message: The incoming RabbitMQ message.
            json_data: The JSON data from the message.
            rabbitmq_url: The RabbitMQ connection URL.
            loop: The asyncio event loop of the RabbitMQ API
            output_dir
        """
        try:
            _message_data = json.loads(message.body.decode())
            _amq_message_data = json.loads(json_data.decode())

            files_to_download = []
            for _key, _file_type in _amq_message_data["files"].items():
                if len(_file_type) > 0:
                    files_info = _file_type[0]
                    stream_name = files_info["streamName"]
                    files_info["queue_name"] = (
                        _message_data.get(stream_name, "") if stream_name else ""
                    )
                    files_info["key"] = _key
                    files_to_download.append(files_info)

            if not files_to_download:
                logging_info(" [%s] No files to download in this message.", message.delivery_tag)
                # await message.ack()
                return defaultdict()

            # Create a task for each file download
            _tasks = []
            for file_info in files_to_download:
                logging_info(f" [{message.delivery_tag}] about to create task for {file_info}")
                _task = asyncio.create_task(
                    self.download_file(loop, rabbitmq_url, file_info, _message_data, output_dir)
                )
                _tasks.append(_task)
            logging_info(
                f" [{message.delivery_tag}] Waiting for all files to download...gather(*tasks)"
            )
            await asyncio.gather(*_tasks)  # Wait for all downloads to complete

            logging_info(f"All files downloaded. d-tag: {message.delivery_tag}")
        except Exception as e:
            logging_error(f"Building DataMessage: {e}")
            asyncio.run_coroutine_threadsafe(
                message.ack(), loop=loop
            )  # ACK -> ignore the message (do not redeliver)

        return download_locations

    async def publish_file(
        self,
        exchange: AbstractExchange,
        file_path: Path,
        routing_key: str,
        loop: AbstractEventLoop,
        max_chunk_size: int = 2**20,  # 1MB default chunk size
        content_type: str = "application/octet-stream",  # Default content type
        delivery_mode: DeliveryMode = DeliveryMode.PERSISTENT,
    ) -> tuple[str, int]:
        """
        Reads a file from the file system and publishes its content to a RabbitMQ exchange
        in chunks.  It declares a queue with a random name, binds it to the specified
        exchange, and uses the queue name as the routing key.

        Args:
            exchange: The aio_pika RobustChannel to use for communication with RabbitMQ.
            file_path: The path to the file to publish.
            max_chunk_size: The maximum size of each chunk (in bytes).
            routing_key: The name of the RabbitMQ exchange routing key to publish to correct queue.
            loop: correct event loop
            max_chunk_size: chunk size to send
            content_type: The content type of the file data.
            delivery_mode: The delivery mode (e.g., PERSISTENT or TRANSIENT).

        Returns:
            The name of the randomly generated queue where the file content was published.
        """
        if not file_path.is_file():
            raise FileNotFoundError(f"File not found: {file_path}")
        _file_size = os.path.getsize(file_path)

        with open(file_path, "rb") as f:
            _offset = 0
            while True:
                _chunk = f.read(max_chunk_size)
                if not _chunk:
                    break  # End of file

                message = Message(
                    body=_chunk,
                    content_type=content_type,
                    delivery_mode=delivery_mode,
                    headers={
                        "file_name": file_path.name,
                        "pos": _offset,
                        "total_size": _file_size,
                        "chunk_size": len(_chunk),
                        "eof": (
                            END_OF_FILE if max_chunk_size + _offset >= _file_size else IN_PROGRESS
                        ),
                    },
                )
                logging.debug(
                    f"Going to Publish chunk {_offset}/{_file_size} to {exchange.name}:{routing_key}"
                )
                _future = asyncio.run_coroutine_threadsafe(
                    exchange.publish(message, routing_key=routing_key), loop
                )
                while not _future.done():
                    # sleep allows coroutine to execute. must wait for completion to ensure all chunks written in right order
                    await asyncio.sleep(0.02)
                    logging_debug(f"Waiting for chunk {_offset}/{_file_size} future= {_future}")

                _offset += len(_chunk)
                logging_debug(
                    f"Published chunk {_offset}/{_file_size} to {exchange.name}:{routing_key}"
                )

        logging_debug(
            f"File {file_path} published to RabbitMQ exchange {exchange.name}, queue {routing_key}"
        )
        return routing_key, _file_size
