import asyncio


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,
        }


class MockContext:
    def get_remaining_time_in_millis(self):
        return 900000


if __name__ == "__main__":
    import time

    t_s = time.perf_counter()

    res = lambda_handler(
        {
            "task_id": "1",
            "kwargs": {
                "cluster": "1311285575352848384",
                "start_ts": 1656594000,
                "stop_ts": 1656594180,
                "org_id": "1372813089193071289",
                "org_type": "op",
                "packages": ["f3d0601e4e5c32c3-4cb7b3785aebead5-4687514cb3853801"],
                "index": "org1372813089193071289_log",
            },
        },
        MockContext(),
    )
    print(f"Time cost：{(time.perf_counter() - t_s):.2f}s Return：{res}")