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 = 1647620760000
    stop = 1647620820000
    step = 10000
    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(
            {
                "kwargs": {
                    "cluster": "7076471312392201216",
                    "index": "org1_log",
                    "package": "99ddf1bf22478d51-5de03da534e4c8f4-437f2123f46ccb0a",
                    "start_ts": b,
                    "stop_ts": e,
                }
            },
            MockContext(),
        )