import os
import requests
import yaml
import pymysql
from datetime import datetime, timedelta
import schedule
import time


cur_dir = os.path.split(os.path.realpath(__file__))[0]
config_file="/etc/dbaas/config.yaml"

env_param = os.getenv("ENV", default="dev")

def read_account_config():
    import yaml
    f = open(config_file, 'r', encoding='utf-8')
    cfg = f.read()
    d = yaml.safe_load(cfg)
    return d

if env_param == "local-test":
    config_file = "/etc/dbaas/config-test.yaml"

print("config_file is " + config_file)

def iso_to_millis(iso_string):
    dt = datetime.strptime(iso_string, "%Y-%m-%dT%H:%M:%SZ")
    return int(dt.timestamp() * 1000)

def query_grafana_item(api_key, expr, start_time, end_time):
    url = 'https://grafana.oauth.prod.aws.tidbcloud.com/api/ds/query'
    headers = {
        'Authorization': f'Bearer {api_key}',
        'sec-ch-ua': '"Not/A)Brand";v="8", "Chromium";v="126", "Google Chrome";v="126"',
        'content-type': 'application/json',
        'accept': 'application/json, text/plain, */*'
    }
    
    body = {
        "queries": [
            {
                "datasource": {
                    "uid": "4BCVu7r4z",
                    "type": "prometheus"
                },
                "editorMode": "code",
                "expr": expr,
                "interval": "7d",
                "legendFormat": "__auto",
                "range": True,
                "refId": "A",
                "queryType": "timeSeriesQuery",
                "exemplar": False,
                "requestId": "95A",
                "utcOffsetSec": 28800,
                "datasourceId": 42,
                "intervalMs": 21600000,
                "maxDataPoints": 1060
            }
        ],
        "range": {
            "from": start_time,
            "to": end_time,
            "raw": {
                "from": start_time,
                "to": end_time
            }
        },
        "from": str(iso_to_millis(start_time)),
        "to": str(iso_to_millis(end_time))
    }
    
    response = requests.post(url, headers=headers, json=body)
    return response.json()

def parse_result(result):
    res_list = []
    try:
        data = result['results']['A']['frames'][0]['data']['values']
        timestamps = data[0]
        values = data[1]
        for ts, val in zip(timestamps, values):
            res_list.append((ts, val))
    except KeyError as e:
        print(f"Error parsing result: {e}")
    return res_list

def insert_into_db(legend, data, db_params):
    connection = pymysql.connect(**db_params)
    cursor = connection.cursor()
    
    insert_query = """
    INSERT INTO prometheus_convert_metrics (legend, metrics_datetime, value)
    VALUES (%s, FROM_UNIXTIME(%s / 1000), %s)
    ON DUPLICATE KEY UPDATE value = VALUES(value)
    """
    
    for ts, val in data:
        cursor.execute(insert_query, (legend, ts, val))
    
    connection.commit()
    cursor.close()
    connection.close()

def job():
    config_info = read_account_config()
    api_key = config_info["apikey"]  
    query_dict = {
        "access_token_call_num": "sum(increase(grpc_server_handled_total{namespace=\"prod-ms\", grpc_service=~\"provider.AuthServerService\", grpc_method=~\"GetIdentityByAccessToken\", grpc_type=\"unary\"}[7d]))",
    }

#    db_params = {
#        'host': 'your_mysql_host',  
#        'user': 'your_mysql_user', 
#        'password': 'your_mysql_password', 
#        'database': 'your_mysql_database', 
#        'charset': 'utf8mb3'
#    }

    now = datetime.utcnow()

    end_time_midnight = datetime(now.year, now.month, now.day).isoformat() + 'Z'

    one_week_ago = now - timedelta(days=7)
    start_time_midnight = datetime(one_week_ago.year, one_week_ago.month, one_week_ago.day).isoformat() + 'Z'

    print(f"start_time:{start_time_midnight} end_time:{end_time_midnight}")

    for legend, query in query_dict.items():
        result = query_grafana_item(api_key, query, start_time_midnight, end_time_midnight)
        parsed_data = parse_result(result)
        insert_into_db(legend, parsed_data, config_info["mysql"])
        print(f"{now} {legend} has been inserted into db")

schedule.every(12).hours.do(job)

job()

while True:
    schedule.run_pending()
    time.sleep(1)