import time
import base64
import pickle
import logging
import asyncio
import functools
import aiobotocore
from .pool import Pool
from .worker import BatchWorker
loop = asyncio.get_event_loop()
logger = logging.getLogger(__name__)


class SQSWriter:
    service_name = 'sqs'
    MAX_RETRY = 10

    def __init__(self, queue_name, region_name,
                 encode_func=None, **pool_kwargs):
        self.session = aiobotocore.get_session()
        self.queue_name = queue_name
        self.region_name = region_name
        self.encode_func = encode_func or (
            lambda o: base64.b64encode(pickle.dumps(o, 2)).decode('ascii')
        )
        self.jobs = set()
        self.pool = Pool(self.create_client, **pool_kwargs)
        self.QueueUrl = None
        self.lock = asyncio.Lock()
        self.loop = asyncio.get_event_loop()
        self.worker = BatchWorker(self.handle)

    def __repr__(self):
        return '<{}: queue={}, region={}, worker={!r}, pool={!r}>'.format(
                self.__class__.__name__, self.queue_name, self.region_name,
                self.worker, self.pool)

    async def put(self, obj):
        return await self.worker.put(obj)

    async def _get_queue_url(self):
        if self.QueueUrl is None:
            async with self.lock:
                async with self.pool.get() as client:
                    resp = await client.get_queue_url(
                            QueueName=self.queue_name
                        )
                    self.QueueUrl = resp['QueueUrl']
        return self.QueueUrl

    def create_client(self):
        return self.session.create_client(
            self.service_name, region_name=self.region_name
        )

    async def stop(self):
        timestamp = time.time()
        await self.worker.stop()
        if self.jobs:
            await asyncio.wait(self.jobs)
        await self.pool.close()
        cost = time.time() - timestamp
        logger.info('{0!r} stopped in {1:.1f} seconds'.format(self, cost))

    def _encode(self, obj):
        return self.encode_func(obj)

    async def handle(self, obj_list):
        client = await self.pool.acquire()
        fut = asyncio.ensure_future(
            self.pool.auto_release(
                client, self.write_batch(client, obj_list)
            )
        )
        self.jobs.add(fut)
        fut.add_done_callback(self.jobs.remove)

    async def write_one(self, record, delay_seconds=0, queued=False, _seq=0,
                        **attributes):
        if _seq > 0:
            await asyncio.sleep(0.1 * (2 ** _seq))
        if _seq > self.MAX_RETRY - 1:
            raise Exception('write_one error: SQS send_message failed')
        message = {
            'record': record,
            'delay_seconds': delay_seconds,
            'attributes': attributes,
        }
        if queued:
            await self.worker.put(message)
            return
        try:
            async with self.pool.get() as client:
                return await client.send_message(
                    QueueUrl=(await self._get_queue_url()),
                    MessageBody=self._encode(record),
                    DelaySeconds=delay_seconds,
                    MessageAttributes=attributes,
                )
        except Exception as e:
            logger.exception(e)
            return await self.write_one(
                    record, delay_seconds=delay_seconds,
                    queued=False, _seq=_seq+1, **attributes)

    async def write_batch(self, client, obj_list, _seq=0):
        if _seq > 0:
            await asyncio.sleep(0.1 * (2 ** _seq))
        if _seq > self.MAX_RETRY - 1:
            raise Exception('write_batch error: SQS send_message_batch failed',
                            obj_list)
        Entries = [
            {
                'Id': str(i),
                'MessageBody': self._encode(message['record']),
                'DelaySeconds': message['delay_seconds'],
                'MessageAttributes': message['attributes'],
            } for i, message in enumerate(obj_list)
        ]
        try:
            resp = await client.send_message_batch(
                QueueUrl=(await self._get_queue_url()),
                Entries=Entries
            )
        except Exception as e:
            logger.exception(e)
            return obj_list
        if 'Failed' in resp:
            logger.error('Send failed: {}, {}'.format(
                obj_list, resp['Failed']))
            failed_obj_list = [
                obj_list[int(d['Id'])] for d in resp['Failed']
                if not d['SenderFault']
            ]
            return await self.write_batch(
                client, failed_obj_list, _seq=_seq + 1
            )

    def async_rpc(self, func=None, *, delay_seconds=0, queued=True,
                  **attributes):
        if func is None:
            return functools.partial(self.async_rpc,
                                     delay_seconds=delay_seconds,
                                     queued=queued, **attributes)
        if type(func) == str:
            async def async_func(*args, **kw):
                record = {
                    'fpath': func,
                    'args': args,
                    'kw': kw,
                }
                await self.write_one(record, delay_seconds=delay_seconds,
                                     queued=queued, **attributes)
        else:
            @functools.wraps(func)
            async def async_func(*args, **kw):
                fpath = func.__module__ + '.' + func.__name__
                record = {
                    'fpath': fpath,
                    'args': args,
                    'kw': kw,
                }
                await self.write_one(record, delay_seconds=delay_seconds,
                                     queued=queued, **attributes)
        return async_func

    # compatible with old async wrapper
    def old_async(self, func=None, *,
                  module=None, max_retry=0, retry_delay=False):
        if func is None:
            return functools.partial(self.old_async,
                                     module=module,
                                     max_retry=max_retry,
                                     retry_delay=retry_delay)

        @functools.wraps(func)
        async def async_func(*args, **kw):
            task_name = func.__name__
            if task_name.endswith('_async'):
                task_name = task_name[:-6]
            attrs = {
                'module': module,
                'task_name': task_name,
                'max_retry': max_retry,
                'retry_delay': retry_delay,
                'retry': 0,
                'args': args,
                'kwargs': kw,
            }
            await self.write_one(attrs, queued=True)
        async_func._real = func
        return async_func


class StringAttribute(dict):
    def __init__(self, value):
        self['DataType'] = 'String'
        self['StringValue'] = str(value)


class BinaryAttribute(dict):
    def __init__(self, value):
        self['DataType'] = 'Binary'
        self['BinaryValue'] = bytes(value)


class NumberAttribute(dict):
    def __init__(self, value):
        self['DataType'] = 'Number'
        self['StringValue'] = str(value)
