import os
import boto3
import time
import json
import tuotuo
import asyncio
import requests
import functools
import numpy as np
import pandas as pd
from datetime import datetime
from tuotuo.metrics import QueryParam, QueryRange

S3_BUCKET = os.getenv('S3_BUCKET')

queries = {
    'connection': 'sum(max_over_time(tidb_server_connections{cluster_id=~".*$tidb_cluster"}[1h])) by (cluster_id)',
    'qps': 'sum(rate(tidb_executor_statement_total{cluster_id=~".*$tidb_cluster"}[1h])) by (cluster_id)',
    #    'tiflash_used': 'sum(rate(tiflash_coprocessor_executor_count{cluster_id=~".*$tidb_cluster"}[1d])) by (cluster_id)',
}


async def sync_metrics(date, checkpoint=None):
    s = requests.Session()
    s3 = boto3.client('s3')
    clinic = tuotuo.client('clinic')
    metrics = tuotuo.client('metrics')
    start = int(time.mktime(date.timetuple()))
    end = start + 24 * 60 * 60 - 1

    clusters = []
    async for cluster in clinic.list_clusters(cluster_status=['normal', 'paused'], deploy_type='dedicated'):
        clusters.append(cluster)
    async for cluster in clinic.list_clusters(sort='clusterDeletedAt', order='descend', deploy_type='dedicated'):
        if cluster.deleted_at / 1000 < start:
            break
        clusters.append(cluster)

    for name, query in queries.items():
        metric_list = []
        for idxs in np.array_split(np.arange(len(clusters)), len(clusters) // 20 + 1):
            metric_list += await metrics.query(*[QueryParam(
                time_arg=QueryRange(start=datetime.fromtimestamp(start), end=datetime.fromtimestamp(end), step='1h'),
                query_template=query,
                cluster_id=str(cluster.id),
                query_name='max_qps',
                query_labels=['cluster_id'],
            ) for cluster in [clusters[i] for i in idxs]])
        for m in metric_list:
            if m.exception:
                print(m.query.cluster_id, m.query.query_str(), m.exception)

        df = pd.concat([m.frame.reset_index().melt(
            id_vars='ds', value_vars=m.frame.columns, var_name='cluster_id', value_name='value'
        ) for m in metric_list if len(m) > 0])
        df['time'] = df.ds.astype(int) // 10 ** 9

        buf = df.drop(columns='ds').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}_{name}.csv',
                                    ContentType='text/csv'),
        )
        print(f'sync data/metrics/{date}_{name}.csv success')
    await metrics.close()

    # make sure this is an async generator
    yield None