import asyncio
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 {
            "code": 200,
            "status": "task finished",
            "message": result,
        }
    except asyncio.TimeoutError:
        print("task timeout, suspending...")
        ctx = task.suspend()
        print("continuation context: ", ctx)
        return {
            "code": 206,
            "status": "task suspended, should resume in another call",
            "message": ctx,
        }


class MockContext:
    def get_remaining_time_in_millis(self):
        return 900000


if __name__ == "__main__":
    start = 1649923440000
    stop = 1649923500000
    step = 60000
    begin = max([(start // step) * step, start])

    for b, e in [(max([x, start]), min([x + step, stop])) for x in range(begin, stop, step)]:
        print(f"========={b}==========")
        lambda_handler(
            {
                "task_id": "1",
                "kwargs": {
                    "cluster": "6862533764151471111",
                    "start_ts": 1646262480000,
                    "stop_ts": 1646262490000,
                    "org_id": "1372813089193071289",
                    "org_type": "cloud",
                    "package_id": "f3d0601e4e5c32c3-4cb7b3785aebead5-4687514cb3853801",
                    "index": "org1372813089193071289_log",
                },
            },
            MockContext(),
        )

        import time

        # time.sleep(3)