import json
import os
from datetime import datetime, timedelta
from pydoc import cli

import boto3
import requests

API_KEY = os.environ["CLINIC_API_KEY"]
CLINIC_ENDPOINT = os.environ["CLINIC_ENDPOINT"]

orgs = [
    "1311265642040528896",
    "1326455380288475136",
    "1372813089189201287",
    "1372813089191061283",
    "1372813089193221283",
    "1372813089189951285",
    "1372813089192171283",
    "1372813089188871292",
    "1372813089193851282",
    "1372813089187041280",
    "1372813089195471292",
]

clusters = []


def lambda_handler(event, context):
    client = boto3.client("lambda")
    if "kwargs" in event:
        # tirggered by scheduler call
        on_scheduler(client, event)
    else:
        # triggered by cron
        on_cron(client)


def on_scheduler(client, event):
    print(event)

    task_id = event["task_id"]
    cluster = event["kwargs"]["tidb_cluster"]
    org_type = event["kwargs"]["org_type"]
    index = event["kwargs"]["index"]
    packages = event["kwargs"]["packages"]
    start = event["kwargs"]["start_ts"] * 1000
    stop = event["kwargs"]["end_ts"] * 1000
    step = 60000
    begin = min([(start // step) * step, start])

    if stop - start > 14400000:  # not allowed log more than 4 hours
        print("log time range too large")
        raise Exception("log time range too large, the maximum value is 4h")

    for b, e in [(max([x, start]), min([x + step, stop])) for x in range(begin, stop, step)]:
        print("process", b, e)
        for package in packages:
            print(f"processing package {package}")
            invoke_lambda(
                client,
                cluster,
                b,
                e,
                task_id=task_id,
                org_type=org_type,
                index=index,
                package_id=package,
            )


def on_cron(client):
    step = 60000
    start_ts = int((datetime.now() - timedelta(minutes=10)).timestamp() * 1000) // step * step
    stop_ts = start_ts + step
    print("process", start_ts, stop_ts)

    for cluster in clusters:
        invoke_lambda(
            client,
            cluster,
            start_ts,
            stop_ts,
            f"{cluster}_{start_ts}",
            org_type="cloud",
        )
    for org in orgs:
        try:
            r = requests.get(
                f"{CLINIC_ENDPOINT}/clinic/api/v1/orgs/{org}/clusters",
                headers={"Authorization": f"Bearer {API_KEY}"},
            )
            if not r.ok:
                continue
            content = json.loads(r.content)
            for cluster in map(lambda x: x["id"], content["items"]):
                invoke_lambda(
                    client,
                    cluster,
                    start_ts,
                    stop_ts,
                    f"{cluster}_{start_ts}",
                    org_type="cloud",
                )
        except Exception as e:
            print(f"error on getting clusters for org {org}: {e}")
            continue


def invoke_lambda(
    client,
    cluster,
    start_ts,
    stop_ts,
    task_id="",
    org_type="",
    index="",
    package_id="",
):
    print(f"processing cluster {cluster}")
    payload = json.dumps(
        {
            "task_id": task_id,
            "kwargs": {
                "cluster": cluster,
                "org_type": org_type,
                "start_ts": start_ts,
                "stop_ts": stop_ts,
                "package_id": package_id,
                "index": index,
            },
        }
    )
    print(payload)
    client.invoke(
        FunctionName="clinic-logs-prod",
        InvocationType="Event",
        Payload=payload,
    )


if __name__ == "__main__":
    lambda_handler(None, None)