import asyncio
from datetime import datetime, timezone


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__("oom_analyzer_resumable")
    task_class = getattr(m, "task_class")()

    # 4 KB limit for aws lambda, for all environment variables associated with the function, in aggregate
    # os.environ["GCP_PRIVATE_KEY"] = base64.b64decode(os.environ["GCP_PRIVATE_KEY"])
    if "resume_context" in event:
        ctx = event["resume_context"]
        print(f"resumed from context {ctx}")
        task = task_class.resume(ctx)
    else:
        related_event_id = event["kwargs"]["activity_id"]
        cluster_id = event["kwargs"]["cluster_id"]
        instance = event["kwargs"]["instance"]
        tenant_id = event["kwargs"]["tenant_id"]
        project_id = event["kwargs"]["project_id"]
        occurrence_ts = event["kwargs"]["occurrence_ts"]
        clinic_org_id = event["kwargs"]["clinic_org_id"]
        provider = event["kwargs"]["provider"]
        region_name = event["kwargs"]["region_name"]
        print(
            datetime.fromtimestamp(occurrence_ts, tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),
            "UTC : related event ",
            related_event_id,
        )

        task = task_class(**event["kwargs"])

    try:
        result = asyncio.run(
            asyncio.wait_for(task.run(), context.get_remaining_time_in_millis() / 1000 - 20)
        )
        return result
    except asyncio.TimeoutError:
        print("task timeout, suspending...")
        ctx = task.suspend()
        print("continuation context: ", ctx)
        return {
            "code": 206,
            "context": ctx,
        }


if __name__ == "__main__":
    event = {
        "kwargs": {
            "activity_id": "3b61bcad-968d-4f8e-9dc5-150d288353a3",
            "cluster_id": "1379661944629574082",
            "instance": "db-tidb-0",
            "occurrence_ts": 1666791974,
            "tenant_id": "1372813089198981282",
            "project_id": "1372813089196881283",
        }
    }
    lambda_handler(event, None)