import logging
import signal
import time
from concurrent.futures import ThreadPoolExecutor, wait
from typing import Any, Callable, Dict

from pydantic import BaseModel, Field, PrivateAttr

from sqsx.helper import backoff_calculator_seconds, base64_to_dict, dict_to_base64

logger = logging.getLogger(__name__)
queue_url_regex = r"(http|https)[:][\/]{2}[a-zA-Z0-9-_:.]+[\/][0-9]{12}[\/]{1}[a-zA-Z0-9-_]{0,80}"


class BaseQueueMixin:
    def consume_messages(
        self, max_messages: int = 1, max_threads: int = 1, wait_seconds: int = 10, run_forever: bool = True
    ) -> None:
        logger.info(f"Starting consuming tasks, queue_url={self.url}")
        signal.signal(signal.SIGINT, self._exit_gracefully)
        signal.signal(signal.SIGTERM, self._exit_gracefully)

        while True:
            if self._should_consume_tasks_stop:
                logger.info(f"Stopping consuming tasks, queue_url={self.url}")
                break

            response = self.sqs_client.receive_message(
                QueueUrl=self.url,
                AttributeNames=["All"],
                MaxNumberOfMessages=min(max_messages, 10),
                MessageAttributeNames=["All"],
            )

            sqs_messages = response.get("Messages", [])
            if not sqs_messages:
                logger.debug(
                    f"Waiting some seconds because no message was received, seconds={wait_seconds}, queue_url={self.url}"
                )
                time.sleep(wait_seconds)
                continue

            with ThreadPoolExecutor(max_workers=max_threads) as executor:
                futures = []
                for sqs_message in sqs_messages:
                    futures.append(executor.submit(self._consume_message, sqs_message))
                wait(futures)

            if not run_forever:
                break

    def _exit_gracefully(self, signal_num, current_stack_frame) -> None:
        logger.info("Starting graceful shutdown process")
        self._should_consume_tasks_stop = True

    def _message_ack(self, sqs_message: dict) -> None:
        receipt_handle = sqs_message["ReceiptHandle"]
        self.sqs_client.delete_message(QueueUrl=self.url, ReceiptHandle=receipt_handle)

    def _message_nack(self, sqs_message: dict) -> None:
        receipt_handle = sqs_message["ReceiptHandle"]
        receive_count = int(sqs_message["Attributes"]["ApproximateReceiveCount"]) - 1
        timeout = backoff_calculator_seconds(
            receive_count, self.min_backoff_seconds, self.max_backoff_seconds
        )
        self.sqs_client.change_message_visibility(
            QueueUrl=self.url, ReceiptHandle=receipt_handle, VisibilityTimeout=timeout
        )


class Queue(BaseModel, BaseQueueMixin):
    url: str = Field(pattern=queue_url_regex)
    sqs_client: Any
    min_backoff_seconds: int = Field(default=30)
    max_backoff_seconds: int = Field(default=900)
    _handlers: Dict[str, Callable] = PrivateAttr(default={})
    _should_consume_tasks_stop: bool = PrivateAttr(default=False)

    def add_task(self, task_name: str, **task_kwargs) -> dict:
        return self.sqs_client.send_message(
            QueueUrl=self.url,
            MessageAttributes={"TaskName": {"DataType": "String", "StringValue": task_name}},
            MessageBody=dict_to_base64({"kwargs": task_kwargs}),
        )

    def add_task_handler(self, task_name: str, task_handler_function: Callable) -> None:
        self._handlers.update({task_name: task_handler_function})

    def _consume_message(self, sqs_message: dict) -> None:
        message_id = sqs_message["MessageId"]
        task_name_attribute = sqs_message["MessageAttributes"].get("TaskName")
        if task_name_attribute is None:
            logger.warning(f"Message without TaskName attribute, message_id={message_id}")
            return self._message_nack(sqs_message)

        task_name = task_name_attribute["StringValue"]
        task_handler_function = self._handlers.get(task_name)
        if task_handler_function is None:
            logger.warning(f"Task handler not found, message_id={message_id}, task_name={task_name}")
            return self._message_nack(sqs_message)

        try:
            message_data = base64_to_dict(sqs_message["Body"])
        except Exception:
            logger.exception(f"Invalid message body, message_id={message_id}, task_name={task_name}")
            return self._message_nack(sqs_message)

        kwargs = message_data["kwargs"]
        context = {
            "queue_url": self.url,
            "task_name": task_name,
            "sqs_message": sqs_message,
        }

        try:
            task_handler_function(context, **kwargs)
        except Exception:
            logger.exception(f"Error while processing, message_id={message_id}, task_name={task_name}")
            return self._message_nack(sqs_message)

        self._message_ack(sqs_message)


class RawQueue(BaseModel, BaseQueueMixin):
    url: str = Field(pattern=queue_url_regex)
    message_handler_function: Callable
    sqs_client: Any
    min_backoff_seconds: int = Field(default=30)
    max_backoff_seconds: int = Field(default=900)
    _should_consume_tasks_stop: bool = PrivateAttr(default=False)

    def add_message(self, message_body: str, message_attributes: dict = {}) -> dict:
        return self.sqs_client.send_message(
            QueueUrl=self.url,
            MessageAttributes=message_attributes,
            MessageBody=message_body,
        )

    def _consume_message(self, sqs_message: dict) -> None:
        try:
            self.message_handler_function(self.url, sqs_message)
        except Exception:
            message_id = sqs_message["MessageId"]
            logger.exception(f"Error while processing, message_id={message_id}")
            return self._message_nack(sqs_message)

        self._message_ack(sqs_message)
