import os
import sys

env = 'tidbprod'
if len(sys.argv) >= 2:
    env = sys.argv[1]

conf_path = os.path.abspath(__file__).split("energy-data-trans")[0] + f"/energy-data-trans/conf/etl_config_{env}.yaml"
os.environ['ETL_CONF'] = conf_path
os.environ['env'] = env

from service.common_connect import trans, plt


def get_all_tables():
    query_sql = f"""
    
    SELECT 
        t.TABLE_NAME
    FROM
        information_schema.`TABLES` t
    WHERE
        t.TABLE_SCHEMA = 'energy_data_prod'
"""

    return trans.execute(query_sql)


def get_all_wind_company():
    query_sql = "SELECT t.field_code,t.field_name FROM wind_field t where t.del_state = 0"
    datas = plt.execute(query_sql)
    result_dict = dict()
    for data in datas:
        result_dict[data['field_code']] = data['field_name']

    return result_dict


if __name__ == '__main__':
    code_name_dict = get_all_wind_company()
    tables = get_all_tables()
    for table in tables:
        table_name = table['TABLE_NAME']

        if table_name.startswith('WOF'):
            field_code = table_name.split('_')[0].split('-')[0]
            if field_code in code_name_dict.keys():
                update_sql = f"ALTER TABLE `{table_name}` COMMENT = '{code_name_dict[field_code]}'"
                trans.execute(update_sql)