import asyncio
import base64
import json
from datetime import datetime, timedelta

from gw_access_log import sync_gw_access_log
from metrics import sync_metrics
from mixpanel import sync_mixpanel
from qps import sync_qps
from sentry import sync_sentry_events

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_gw_access_log,
        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",
            }
        )