import json
import base64
import asyncio
from mixpanel import sync_mixpanel
from metrics import sync_metrics
from sentry import sync_sentry_events
from qps import sync_qps
from datetime import datetime, timedelta

checkpoints = {}


async def swallow_exception(f, p):
    global checkpoints
    try:
        async for checkpoint in f(p, checkpoint=checkpoints.get(f.__name__)):
            if checkpoint is None:
                checkpoints.pop(f.__name__, None)
            else:
                checkpoints[f.__name__] = checkpoint
    except Exception as e:
        raise e
        print(f'got exception in {f.__name__}: {e}')


async def run(date):
    funcs = [
        sync_mixpanel,
        sync_sentry_events,
        # sync_metrics,
        # sync_qps,
    ]
    if len(checkpoints) > 0:
        funcs = [f for f in funcs if f.__name__ in checkpoints]

    await asyncio.gather(*[swallow_exception(f, date) for f in funcs])


def lambda_handler(event, context):
    global checkpoints

    date = datetime.now().date() - timedelta(days=1)
    if 'kwargs' in event:
        date = datetime.utcfromtimestamp(event['kwargs']['time']).date()
    if 'resume_context' in event:
        checkpoints = json.loads(base64.b64decode(event['resume_context']))
    try:
        asyncio.run(asyncio.wait_for(run(date), context.get_remaining_time_in_millis() / 1000 - 10))
        return {
            'success': True,
        }
    except asyncio.TimeoutError:
        return {
            'code': 206,
            'context': base64.b64encode(json.dumps(checkpoints).encode()).decode(),
        } if len(checkpoints) > 0 else {
            'success': False,
            'error': 'timeout',
        }