import asyncio
import json
import sys
from pathlib import Path

sys.path = [str(Path(__file__).parent)] + sys.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
        # return json.loads(json.dumps(result, default=str))
    except asyncio.TimeoutError:
        print("task timeout, suspending...")
        ctx = task.suspend()
        print("continuation context: ", ctx)
        return {
            "code": 206,
            "context": ctx,
        }


if __name__ == "__main__":
    from app import task_class

    cluster = "1379661944602244071"
    start = 1650591648
    end = 1650592648
    event = {
        "task_id": "1234",
        "kwargs": {
            "cluster": cluster,
            "start_ts": start,
            "stop_ts": end,
            "org_id": "1326455380288475136",
            "org_type": "cloud",
            "packages": [],
            "index": "org1326455380288475136_log",
        },
    }
    task_class = task_class()
    task = task_class(**event["kwargs"])

    try:
        result = asyncio.run(asyncio.wait_for(task.run(), 100000))
        print(
            {
                "code": 200,
                "status": "task finished",
                "message": result,
            }
        )
    except asyncio.TimeoutError:
        print("task timeout and cannot be suspend")
        print(
            {
                "code": 500,
                "status": "task timeout",
                "message": "task timeout",
            }
        )