import asyncio
import json
import os
from datetime import datetime, timezone

import aiohttp
import boto3
import tuotuo

CE_ENDPOINT = os.environ["CE_ENDPOINT"]
OOM_DETECTOR = os.environ["OOM_DETECTOR"]
TIKV_OOM_DETECTOR = os.environ["TIKV_OOM_DETECTOR"]
TIFLASH_OOM_DETECTOR = os.environ["TIFLASH_OOM_DETECTOR"]
REBOOT_DETECTOR = os.environ["REBOOT_DETECTOR"]
ACTIVITY_BUCKET = os.environ["ACTIVITY_BUCKET"]
LAST_TS_FILE = os.environ["LAST_TS_FILE"]
CAL_INTERVAL = int(os.environ["INTERVAL_SECONDS"])
DELAY_SECONDS = int(os.environ["DELAY_SECONDS"])


async def agent(session, post_args):
    failed = []
    success = []
    async with session.post(
        f"{CE_ENDPOINT}/tasks/{OOM_DETECTOR}/async", data=json.dumps(post_args)
    ) as resp:
        if resp.status != 200:
            msg = await resp.text()
            failed.append(msg)
            print(f"{OOM_DETECTOR} failed, msg: {msg}")
        else:
            res = await resp.json()
            success.append(res["id"])
    async with session.post(
        f"{CE_ENDPOINT}/tasks/{TIKV_OOM_DETECTOR}/async", data=json.dumps(post_args)
    ) as resp:
        if resp.status != 200:
            msg = await resp.text()
            failed.append(msg)
            print(f"{TIKV_OOM_DETECTOR} failed, msg: {msg}")
        else:
            res = await resp.json()
            success.append(res["id"])
    async with session.post(
        f"{CE_ENDPOINT}/tasks/{TIFLASH_OOM_DETECTOR}/async", data=json.dumps(post_args)
    ) as resp:
        if resp.status != 200:
            msg = await resp.text()
            failed.append(msg)
            print(f"{TIFLASH_OOM_DETECTOR} failed, msg: {msg}")
        else:
            res = await resp.json()
            success.append(res["id"])
    async with session.post(
        f"{CE_ENDPOINT}/tasks/{REBOOT_DETECTOR}/async",
        data=json.dumps(post_args),
    ) as resp:
        if resp.status != 200:
            msg = await resp.text()
            failed.append(msg)
            print(f"{REBOOT_DETECTOR} failed, msg: {msg}")
        else:
            res = await resp.json()
            success.append(res["id"])

    return {"failed": failed, "success": success}


async def all():
    # 1 list all normal clusters
    clinic = tuotuo.client("clinic")

    # Clinic canonicalized normal/scaling/recovering to active/modifying/restoring.
    # Keep both generations while compute-engine and Clinic are deployed independently.
    cluster_status = [
        "normal",
        "active",
        "upgrading",
        "importing",
        "scaling",
        "modifying",
        "resuming",
        "recovering",
        "restoring",
    ]
    cluster_args = []
    async for cluster in clinic.list_clusters(
        query="", cluster_status=cluster_status, deploy_type="dedicated"
    ):
        # if cluster.provider == 'gcp', google storage api 'list_blobs' will generate a large number of requests,
        # The cost of cross-cloud platform calls is very high.
        if cluster.provider == "aws":
            post_args = {
                "tags": {"source": "run-detector-agents"},
                "args": {
                    "cluster_id": cluster.id,
                    "tenant_id": cluster.tenant_id,
                    "project_id": cluster.project_id,
                    "clinic_org_id": cluster.clinic_org_id,
                    "provider": cluster.provider,
                    "region_name": cluster.region_name,
                    "shoot_name": cluster.shoot_name,
                },
            }
            cluster_args.append(post_args)

    tasks = []
    # 2 get last end_ts
    s3 = boto3.client("s3")
    try:
        ts_res = s3.get_object(Bucket=ACTIVITY_BUCKET, Key=LAST_TS_FILE)["Body"]
        start_ts = json.loads(ts_res.read())["last_ts"]
    except s3.exceptions.NoSuchKey:
        start_ts = None
    end_ts = int(datetime.now().timestamp()) - DELAY_SECONDS
    if start_ts is None:
        start_ts = end_ts - CAL_INTERVAL
    elif end_ts - start_ts > 3 * CAL_INTERVAL:
        print(f"[Warning] last end_ts {start_ts} is too old, reset to {end_ts - CAL_INTERVAL}")
        start_ts = end_ts - CAL_INTERVAL

    for i in range(len(cluster_args)):
        cluster_args[i]["args"]["start_ts"] = start_ts
        cluster_args[i]["args"]["end_ts"] = end_ts

    print(
        f"scan total {len(cluster_args)} normal clusters between {datetime.fromtimestamp(start_ts, tz=timezone.utc)} and {datetime.fromtimestamp(end_ts, tz=timezone.utc)}"
    )  # without create_time info
    # 3 scan
    task_results = {"success_count": 0, "failed_details": []}
    async with aiohttp.ClientSession(trust_env=True) as session:
        # task_results = await asyncio.gather(*[session.post(f'{CE_ENDPOINT}/tasks/{DETECTOR_AGENT}/async', data=json.dumps(post_args)) for post_args in cluster_args], return_exceptions=True)
        for i in range(0, len(cluster_args), 10):
            responses = await asyncio.gather(
                *[agent(session, post_args) for post_args in cluster_args[i : i + 10]],
                return_exceptions=True,
            )
            for res in responses:
                task_results["success_count"] += len(res["success"])
                task_results["failed_details"].extend(res["failed"])

    # 4 save the latest end_ts
    ts_msg = {"last_ts": end_ts}
    s3.put_object(Body=json.dumps(ts_msg), Bucket=ACTIVITY_BUCKET, Key=LAST_TS_FILE)
    print(f"task_results: {task_results}")
    return task_results


def lambda_handler(event, context):
    tasks = asyncio.run(all())
    return tasks