import os
import json
import time
import boto3
import tuotuo
import asyncio
import functools
import pandas as pd
from datetime import datetime, timedelta
import gc
import base64
from memory_profiler import profile

S3_BUCKET = "pingcap-plg"

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-access-logs'
    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
            start_index = entity.content.find('{')
            end_index = entity.content.rfind('}')
            json_part = entity.content[start_index:end_index + 1]
            x = json.loads(json_part)
            x['gateway_instance'] = entity.instance
            x['region'] = region

            time_str = entity.content[:23]
            timestamp_with_microseconds = datetime.strptime(time_str, "%Y-%m-%dT%H:%M:%S.%f").timestamp()
            x['ts'] = timestamp_with_microseconds
            if bool(x.get('cluster_id')):
                logs.append(x)
        except Exception as e:
            print(e, entity)
    return logs
    
def sample(df, prefix):
    s3 = boto3.client('s3')
    df = df[df.duration > 1]
    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/{prefix}.sample.csv', ContentType='text/csv')
    s3.put_object(Body=buf, Bucket=S3_BUCKET, Key=f'data/gw_access_log_sample/{prefix}.sample.csv', ContentType='text/csv')   
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print("sample log: ", current_time)
    print(f'sync data/gw_access_log/{prefix}.sample.csv success')

def update_time_granularity(n, time_granularity):
    max_time_granularity = 60 * 60 # 1 hour
    standard_log_count = 400000
    if n > standard_log_count:
        return min(max(time_granularity * standard_log_count // n, 1), max_time_granularity)
    else:
        return min(time_granularity * 2, max_time_granularity)

async def sync_gw_access_log(date, checkpoint=None):
    s3 = boto3.client('s3')
    client = tuotuo.client('logs')
    
    start = int(time.mktime(date.timetuple()))
    checkpoint = start
    df = pd.DataFrame()
    end = start + 24 * 60 * 60

    time_granularity = 1 * 60  # 1 minute
    
    regions = ['ap-northeast-1', 'ap-southeast-1', 'us-west-2', 'us-east-1', 'eu-central-1']
    while checkpoint < end:
        st = checkpoint
        et = min(checkpoint + time_granularity, end)
        xs = await asyncio.gather(*[pull_gw_access_log(client, region, st, et) for region in regions])
        df = pd.concat([pd.DataFrame.from_records(x) for x in xs], ignore_index=True)
        time_granularity = update_time_granularity(len(df), time_granularity)
        buf = df.to_csv(index=False)
        prefix = f'{date}-{et}'
        s3.put_object(Body=buf, Bucket=S3_BUCKET, Key=f'data/gw_access_log/{prefix}.csv', ContentType='text/csv')
        print(f'sync data/gw_access_log/{prefix}.csv for checkpoint {checkpoint} success')
        sample(df, prefix)
        del df
        del buf
        gc.collect()
        checkpoint = et


date = datetime.now().date() - timedelta(days=1)
current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
print("start time: ", current_time)
asyncio.run(sync_gw_access_log(date))