import os
import sys

sys.path.append(os.path.dirname(sys.path[0]))

from utils.ssh_util import SSHUtil
from utils.common_util import CommonUtil
from utils.db_util import DBUtil
from datetime import date

if __name__ == '__main__':
    site_name = CommonUtil.get_sys_arg(1, None)
    date_type = CommonUtil.get_sys_arg(2, None)
    date_info = CommonUtil.get_sys_arg(3, None)
    print(f"执行参数为{sys.argv}")

    CommonUtil.judge_is_work_hours(
        site_name=site_name, date_type=date_type, date_info=date_info,
        principal='chenyuanjie', priority=3, export_tools_type=1, belonging_to_process='关联流量'
    )

    db_type = 'postgresql_cluster'
    engine = DBUtil.get_db_engine(db_type, site_name)
    export_cols = [
        'asin',
        'related_asin',
        'related_type',
        'related_time',
        'related_count'
    ]

    if date_type == 'month':
        dt = str(date_info).replace("-", "_")
        export_tb = f"{site_name}_asin_related_{dt}"
        report_type = date_type
        report_date = date_info
    elif date_type == 'month_week':
        export_tb = f"{site_name}_asin_related_30_day_copy"
        export_tb_target = f"{site_name}_asin_related_30_day"
        report_type = '30_day'
        report_date = date.today()
    else:
        raise Exception('date_type有误！')

    sql = f"""
        DROP TABLE IF EXISTS {export_tb};
        CREATE TABLE {export_tb} (
          id serial8,
          asin varchar(10) NOT NULL,
          related_asin varchar(10)[] NOT NULL,
          related_type int2[] NOT NULL,
          related_time varchar(10) NOT NULL,
          related_count int4 NOT NULL
        );
        ALTER TABLE {export_tb} ALTER COLUMN related_asin TYPE text;
        ALTER TABLE {export_tb} ALTER COLUMN related_type TYPE text;
    """

    DBUtil.engine_exec_sql(engine, sql)

    partition_dict = {
        "site_name": site_name,
        "date_type": date_type,
        "date_info": date_info
    }
    sh = CommonUtil.build_export_sh(
        site_name=site_name,
        db_type=db_type,
        hive_tb="dwt_asin_related_traffic",
        export_tb=export_tb,
        col=export_cols,
        partition_dict=partition_dict
    )
    client = SSHUtil.get_ssh_client()
    SSHUtil.exec_command_async(client, sh, ignore_err=False)
    client.close()
    print("数据导出完成，准备修改数据类型！")

    sql = f"""
        ALTER TABLE {export_tb}
        ALTER COLUMN related_asin TYPE VARCHAR(10)[]
        USING string_to_array(related_asin, ',');

        ALTER TABLE {export_tb}
        ALTER COLUMN related_type TYPE INTEGER[]
        USING string_to_array(related_type, ',')::int[];
    """
    DBUtil.engine_exec_sql(engine, sql)
    print("数据类型修改完成，准备创建索引！")

    if date_type == 'month':
        sql = f"""
            CREATE INDEX {export_tb}_asin_idx ON {export_tb} USING btree (
              "asin" COLLATE "pg_catalog"."default" "pg_catalog"."text_ops" ASC NULLS LAST
            );

            CREATE INDEX {export_tb}_related_asin_idx ON {export_tb} USING gin (
              "related_asin" COLLATE "pg_catalog"."default" "pg_catalog"."array_ops"
            );
        """
        DBUtil.engine_exec_sql(engine, sql)
        print("索引创建完成，准备插入流程记录表！")
    elif date_type == 'month_week':
        # 构建索引并交换表名
        DBUtil.exchange_tb(
            engine=engine,
            source_tb_name=export_tb,
            target_tb_name=export_tb_target,
            cp_index_flag=True
        )
        print("索引创建完成，准备插入流程记录表！")

    sql = f"""
        REPLACE INTO selection.workflow_everyday 
        (site_name, report_date, status, status_val, table_name, date_type, page, is_end, remark, export_db_type)
        VALUES
        ('{site_name}', '{report_date}', '导出PG集群完成', 14, '{export_tb}', '{report_type}', '关联流量', '是', '关联流量模块数据', 'postgresql_cluster');
    """
    DBUtil.engine_exec_sql(DBUtil.get_db_engine('mysql', 'us'), sql)
    print("success！")
