import json
import logging
import tempfile
from time import sleep
from pyinfra.config import CONFIG

from pyinfra.exceptions import AnalysisFailure, DataLoadingFailure
from pyinfra.rabbitmq import make_connection, make_channel, declare_queue
from pyinfra.storage.storages import get_storage
from pyinfra.utils.file import upload_compressed_response


def make_retry_callback(republish, max_attempts):
    def get_n_previous_attempts(props):
        return 0 if props.headers is None else props.headers.get("x-retry-count", 0)

    def attempts_remain(n_attempts):
        return n_attempts < max_attempts

    def callback(channel, method, properties, body):

        n_attempts = get_n_previous_attempts(properties) + 1

        logging.error(f"Message failed to process {n_attempts}/{max_attempts} times: {body}")

        if attempts_remain(n_attempts):
            republish(channel, body, n_attempts)
            channel.basic_ack(delivery_tag=method.delivery_tag)

        else:
            logging.exception(f"Adding to dead letter queue: {body}")
            channel.basic_reject(delivery_tag=method.delivery_tag, requeue=False)

    return callback


def wrap_callback_in_retry_logic(callback, retry_callback):
    def wrapped_callback(channel, method, properties, body):
        try:
            callback(channel, method, properties, body)
        except (AnalysisFailure, DataLoadingFailure):
            sleep(5)
            retry_callback(channel, method, properties, body)

    return wrapped_callback


def json_wrap(body_processor):
    def inner(payload):
        return json.dumps(body_processor(json.loads(payload)))

    return inner


def make_callback_for_output_queue(json_wrapped_body_processor, output_queue_name):

    connection = make_connection()
    channel = make_channel(connection)
    declare_queue(channel, output_queue_name)

    def callback(channel, method, _, body):

        dossier_id, file_id, result = json_wrapped_body_processor(body)

        if not CONFIG.service.response.save:
            channel.basic_publish(exchange="", routing_key=output_queue_name, body=result)
        else:
            upload_compressed_response(get_storage(CONFIG.storage.backend), CONFIG.storage.bucket, dossier_id, file_id, result)
            result = json.dumps({"dossierId" : dossier_id, "fileId" : file_id})
            channel.basic_publish(exchange="", routing_key=output_queue_name, body=result)

        channel.basic_ack(delivery_tag=method.delivery_tag)

    return callback


def make_retry_callback_for_output_queue(json_wrapped_body_processor, output_queue_name, retry_callback):
    callback = make_callback_for_output_queue(json_wrapped_body_processor, output_queue_name)
    callback = wrap_callback_in_retry_logic(callback, retry_callback)

    return callback