import base64
import json
import os
from datetime import datetime, timedelta, timezone
from uuid import uuid4

import aiohttp
import boto3
import tuotuo
from sinks import FullActivityClient
from tuotuo.activities import Activity, ActivityType, Sevirity

ACTIVITY_BUCKET = os.environ["ACTIVITY_BUCKET"]  # "activities"
PATH = "oom_risk_detector"
ACTIVITY_HUB_ENDPOINT = os.environ["ACTIVITY_HUB_ENDPOINT"]
SOURCE_KEY = os.environ["SOURCE_KEY"]
SOURCE = os.environ["SOURCE"]
CE_ENDPOINT = os.environ["CE_ENDPOINT"]
OOM_ANALYSIS_TASK = os.environ["OOM_ANALYSIS_TASK"]


async def write_activity_to_hub(activity: Activity):
    async with FullActivityClient(
        source=SOURCE, source_key=SOURCE_KEY, endpoint=ACTIVITY_HUB_ENDPOINT
    ) as client:
        return await client.write_activity(activity)


async def run_oom_analysis(activity: Activity, clinic_org_id: str):
    args = {
        "tags": {"source": "oom_risk_detector"},
        "args": {
            "activity_id": activity.activity_id,
            "cluster_id": activity.target_id,
            "instance": activity.payload["instance"],
            "occurrence_ts": activity.calibration_time,
            "tenant_id": activity.application_id,
            "project_id": activity.concept_id,
            "clinic_org_id": clinic_org_id,
        },
    }
    url = f"{CE_ENDPOINT}/tasks/{OOM_ANALYSIS_TASK}/async"
    data = json.dumps(args)
    async with aiohttp.ClientSession(trust_env=True) as session:
        res = await session.post(url, data=data)
        return await res.text()


def dict_to_s3(dict, bucket, key):
    s3 = boto3.resource("s3")
    res = s3.Bucket(bucket).put_object(Key=key, Body=json.dumps(dict))
    return res


async def find_oom_risk(
    cluster_id: str,
    start: datetime,
    end: datetime,
    tenant_id: str = None,
    project_id: str = None,
    region_name: str = None,
    clinic_org_id: str = None,
    provider: str = None,
):
    # get gcp key
    try:
        s3 = boto3.client("s3")
        key = s3.get_object(Bucket=ACTIVITY_BUCKET, Key="secrets/gcp_secret")["Body"]
        os.environ["GCP_PRIVATE_KEY"] = bytes.decode(base64.b64decode(key.read()))
    except Exception as e:
        print(e)

    client = tuotuo.client("logs")
    oom_risk_events = []
    start_ts = int(start.timestamp())
    end_ts = int(end.timestamp())
    async for entity in client.read_single_cluster(
        start_ts,
        end_ts,
        cluster_id=cluster_id,
        project_id=project_id,
        tenant_id=tenant_id,
        region_name=region_name,
        components=["tidb"],
        provider=provider,
    ):
        if "risk of OOM" in entity.content:
            activity_id = str(uuid4())
            instance = entity.instance
            occurrence_time = int(entity.time.timestamp())
            detect_time = int(datetime.now().timestamp())
            summary = f"{instance} of cluster {cluster_id}: an OOM risk occurred."
            details = {
                "cluster_id": cluster_id,
                "instance": instance,
                "tenant_id": tenant_id,
                "project_id": project_id,
                "occurrence_time": occurrence_time,
                "detect_time": detect_time,
                "evidence": {
                    "source": "tidb.log",
                    "content": [entity.content],
                },
            }
            res = dict_to_s3(details, ACTIVITY_BUCKET, f"{PATH}/{activity_id}.json")
            print(res)
            payload = {
                "summary": summary,
                "cluster_id": cluster_id,
                "instance": instance,
                "occurrence_time": occurrence_time,
                "detect_time": detect_time,
                "details": {
                    "type": "s3",
                    "bucket": ACTIVITY_BUCKET,
                    "key": f"{PATH}/{activity_id}.json",
                },
            }
            oom_risk_event = Activity(
                activity_type=ActivityType.EVENT,
                activity_name="OOM Risk",
                target_id=cluster_id,
                calibration_time=occurrence_time,
                source="compute-engine",
                source_key="source_token",
                catalog="OOM",
                severity=Sevirity.INFO,
                activity_id=activity_id,
                application_id=tenant_id,
                concept_id=project_id,
                report_time=detect_time,
                author="oom-risk-detector",
                context_id=activity_id,
                # context=[],
                payload=payload,
                occurrence_time=occurrence_time,
                detect_time=detect_time,
            )
            oom_risk_events.append(oom_risk_event)
        # elif "Out Of Memory Quota" in entity.content:
        #     pass
        # elif "expensive_query" in entity.content:
        #     #TODO: parse expensive query
        #     pass
        # elif "memory exceeds quota" in entity.content:
        #     pass

    if len(oom_risk_events) > 0:
        for activity in oom_risk_events:
            res = await write_activity_to_hub(activity)
            analyze_task = await run_oom_analysis(activity, clinic_org_id)
            print(analyze_task)
            if res.get("status") == "ok":
                print(f"[Info][Write event Successful] {cluster_id} OOM risk activity: {activity}")
            else:
                print(f"[Error][Write event Failed] {cluster_id} OOM risk activity: {activity}")
    return oom_risk_events


if __name__ == "__main__":
    import asyncio
    import time

    """
    input: cluster_id, start, end, tenant_id, project_id
    output: oom_events
    excute: write details to s3 and write oom_events to activityhub
    non-goal: alert, invoke consumer
    non-goals are implemented in agent
    """
    cluster_id = "1379661944629574082"
    tenant_id = "1372813089198981282"
    project_id = "1372813089196881283"
    # start = datetime.fromtimestamp(1668714324) # 1666785600
    # end = start + timedelta(0, 0, 0, 1, 30, 0) # 1666799999
    start = datetime.fromtimestamp(1666785600 + 3600)  # 1666785600
    end = datetime.fromtimestamp(1666785600 + 7200)  # 1666799999
    proc_start = time.perf_counter()
    oom_events = asyncio.run(find_oom_risk(cluster_id, start, end, tenant_id, project_id))
    proc_end = time.perf_counter()
    print(f"Elapsed time: {proc_end - proc_start} seconds")
    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"),
        )