import asyncio
import json
import sys
from pathlib import Path


def lambda_handler(event, context):
    if context.get_remaining_time_in_millis() < 60000:
        raise ValueError("the lambda function's timeout property should be configured at least 60s")

    m = __import__("app")
    task_class = getattr(m, "task_class")()

    if "resume_context" in event:
        ctx = event["resume_context"]
        print("resumed from context {ctx}")
        task = task_class.resume(ctx)
    else:
        task = task_class(**event["kwargs"])

    try:
        result = asyncio.run(
            asyncio.wait_for(task.run(), context.get_remaining_time_in_millis() / 1000 - 10)
        )
        return result
    except asyncio.TimeoutError:
        print("task timeout, suspending...")
        ctx = task.suspend()
        print("continuation context: ", ctx)
        return {
            "code": 206,
            "context": ctx,
        }