import asyncio
import json


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 event.get("kwargs") is not None:
        event["kwargs"]["context"] = context

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

    try:
        result = asyncio.run(
            asyncio.wait_for(task.run(), context.get_remaining_time_in_millis() / 1000 - 10)
        )
        if isinstance(result, dict):
            return json.loads(json.dumps(result, default=str))
        else:
            return {"summary": 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()

    # oss-insight
    cluster = "1379661944623214066"
    # end = int(time.time())
    end = 1660533579  # 2022-08-15 11:19:39
    start = end - 600
    org_type = "cloud"
    org_id = "1372813089187041280"

    # cluster = "1379661944623214066"
    # start = 1650591648
    # end = 1650592648
    # org_type = "cloud"
    # org_id = "1372813089196911678"
    event = {
        "task_id": "914d3aa2-7b21-4ee3-bdb8-6b0f5203a84c",
        "tags": {"source": "83613d40-d446-461f-9920-d2a12f6b7268"},  # predecessor task id
        "kwargs": {
            "cluster": cluster,
            "start_ts": start,
            "stop_ts": end,
            "org_id": org_id,
            "org_type": org_type,
            "packages": [],
            "index": "org1326455380288475136_log",
        },
    }

    res = lambda_handler(event, MockContext())
    print(f"Time cost：{(time.perf_counter() - t_s):.2f}s Return：{res}")