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

import boto3
import pandas as pd
import requests
import tuotuo

S3_BUCKET = os.getenv("S3_BUCKET")


async def sync_qps(date, checkpoint=None):
    s = requests.Session()
    s3 = boto3.client("s3")
    metrics = tuotuo.client("metrics")
    start = int(time.mktime(date.timetuple()))
    end = start + 24 * 60 * 60 - 1
    print(f"start: {start}, end: {end}")
    regions = [
        "us-west-2-f01/aws-us-west-2",
        "us-east-1-f01/aws-us-east-1",
        "ap-northeast-1-f01/aws-ap-northeast-1",
        "eu-central-1-f01/aws-eu-central-1",
        "ap-southeast-1-f01/aws-ap-southeast-1",
    ]
    df = pd.DataFrame(columns=["cluster_id", "value"])
    async with metrics:
        for region in regions:
            qps_res = await metrics._query_range(
                "sum(rate(tidb_executor_statement_total{}[1m])) by (cluster_id)",
                datetime.fromtimestamp(start),
                datetime.fromtimestamp(end),
                "1m",
                headers={
                    "X-OrgType": "cloud",
                    "Authorization": f"Bearer {metrics._clinic_token}",
                    "control_plane_info": f"{region}",
                },
            )
            df = pd.concat(
                [
                    df,
                    pd.DataFrame().from_records(
                        [
                            {"cluster_id": m["metric"]["cluster_id"], "value": m["values"]}
                            for m in qps_res
                            if m["metric"].get("cluster_id") is not None
                        ]
                    ),
                ],
                ignore_index=True,
            )

    buf = df.to_csv(index=False)
    await asyncio.get_running_loop().run_in_executor(
        None,
        functools.partial(
            s3.put_object,
            Body=buf,
            Bucket=S3_BUCKET,
            Key=f"data/metrics/{date}_serverless_qps.csv",
            ContentType="text/csv",
        ),
    )
    print(f"sync data/metrics/{date}_serverless_qps.csv success")
    await metrics.close()

    # make sure this is an async generator
    yield None