import asyncio
from datetime import datetime, timedelta, timezone

from tikv_oom_detector import find_oom_evicted


def lambda_handler(event, context):
    cluster_id = event["kwargs"]["cluster_id"]
    start = datetime.fromtimestamp(event["kwargs"]["start_ts"])
    end = datetime.fromtimestamp(event["kwargs"]["end_ts"])
    tenant_id = event["kwargs"]["tenant_id"]
    project_id = event["kwargs"]["project_id"]
    region_name = event["kwargs"]["region_name"]
    clinic_org_id = event["kwargs"]["clinic_org_id"]
    provider = event["kwargs"]["provider"]
    oom_events = asyncio.run(
        find_oom_evicted(
            cluster_id, start, end, tenant_id, project_id, region_name, clinic_org_id, provider
        )
    )
    return {
        "time_range": f"{start.strftime('%Y-%m-%d, %H:%M:%S')} to {end.strftime('%Y-%m-%d, %H:%M:%S')}",
        "oom_event_count": len(oom_events),
    }


if __name__ == "__main__":
    import asyncio

    cluster_id = "1379661944636804070"
    tenant_id = "1372813089208281359"
    project_id = "1372813089206481407"
    event = {
        "kwargs": {
            "cluster_id": cluster_id,
            "start_ts": 1675756800,
            "end_ts": 1675760400,
            "tenant_id": tenant_id,
            "project_id": project_id,
            "region_name": "",
            "clinic_org_id": "",
            "provider": "",
        }
    }

    async def lambda_handler(event, context):
        cluster_id = event["kwargs"]["cluster_id"]
        start = datetime.fromtimestamp(event["kwargs"]["start_ts"])
        end = datetime.fromtimestamp(event["kwargs"]["end_ts"])
        tenant_id = event["kwargs"]["tenant_id"]
        project_id = event["kwargs"]["project_id"]
        region_name = event["kwargs"]["region_name"]
        clinic_org_id = event["kwargs"]["clinic_org_id"]
        provider = event["kwargs"]["provider"]
        oom_events = await find_oom_evicted(
            cluster_id, start, end, tenant_id, project_id, region_name, clinic_org_id, provider
        )
        return oom_events

    oom_events = asyncio.run(lambda_handler(event, None))
    for oom_event in oom_events:
        print(
            datetime.fromtimestamp(oom_event.calibration_time)
            .replace(tzinfo=timezone.utc)
            .astimezone(timezone(timedelta(hours=8))),
            ": ",
            oom_event.payload.get("summary"),
        )