import asyncio
import functools
import json
import os
import time
from datetime import datetime

import boto3
import pandas as pd
import tuotuo

S3_BUCKET = os.getenv("S3_BUCKET")


async def pull_gw_access_log(client, region: str, start_ts: int, stop_ts: int):
    logs = []
    prefix = f"1/prod-devtier/{region}-f01/k8s-infra/logs/tidb-gateway"
    async for entity in client.read_with_prefix(
        bucket=f"o11y-prod-shared-{region}", prefix=prefix, start_ts=start_ts, stop_ts=stop_ts
    ):
        try:
            if "gw_access_log" not in entity.content:
                continue
            x = json.loads(entity.content)
            x["gateway_instance"] = entity.instance
            if x.get("msg") == "gw_access_log" and bool(x.get("cluster_id")):
                logs.append(x)
        except Exception as e:
            print(e, entity)
    return logs


def sample(date):
    s3 = boto3.client("s3")
    df = pd.read_csv(
        s3.get_object(Bucket=S3_BUCKET, Key=f"data/gw_access_log/{date}.csv")["Body"],
        low_memory=False,
    )
    df = (
        df.set_index(pd.to_datetime(df.ts, unit="s", utc=True))
        .groupby([pd.Grouper(freq="H"), "cluster_id", "status", "gateway_instance"])
        .sample(1)
    )
    buf = df.to_csv(index=False)
    s3.put_object(
        Body=buf,
        Bucket=S3_BUCKET,
        Key=f"data/gw_access_log/{date}.sample.csv",
        ContentType="text/csv",
    )


async def sync_gw_access_log(date, checkpoint=None):
    s3 = boto3.client("s3")
    client = tuotuo.client("logs")

    if checkpoint is None:
        start = int(time.mktime(date.timetuple()))
        checkpoint = start
        df = pd.DataFrame()
    else:
        date = datetime.fromtimestamp(checkpoint).date()
        start = int(time.mktime(date.timetuple()))
        df = pd.read_csv(
            s3.get_object(Bucket=S3_BUCKET, Key=f"data/gw_access_log/{date}.csv")["Body"],
            low_memory=False,
        )
    end = start + 24 * 60 * 60

    regions = ["ap-southeast-1", "ap-northeast-1", "us-west-2", "us-east-1", "eu-central-1"]
    for checkpoint in range(checkpoint, end, 60 * 60):
        yield checkpoint
        xs = await asyncio.gather(
            *[
                pull_gw_access_log(client, region, checkpoint, checkpoint + 60 * 60)
                for region in regions
            ]
        )
        df = pd.concat([df] + [pd.DataFrame.from_records(x) for x in xs])
        buf = df.to_csv(index=False)
        s3.put_object(
            Body=buf, Bucket=S3_BUCKET, Key=f"data/gw_access_log/{date}.csv", ContentType="text/csv"
        )
        print(f"sync data/gw_access_log/{date}.csv for checkpoint {checkpoint} success")

    sample(date)
    # make sure this is an async generator
    yield None