import os
import sys

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

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

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)
    CommonUtil.judge_is_work_hours(site_name=site_name, date_type=date_type, date_info=date_info,
                                   principal='chenyuanjie', priority=1, export_tools_type=1)

    db_type = "postgresql_cluster"
    print("导出到PG集群")
    master_tb = f"{site_name}_amazon_report"
    export_tb = f"{site_name}_amazon_report_copy"

    # 获取数据库连接
    engine = DBUtil.get_db_engine(db_type, site_name)

    with engine.connect() as connection:
        # 不再用 like {master_tb} including all 复制结构：显式建表，新字段直接写进结构里，
        # 不需要对线上分布式表 master_tb 做 ALTER TABLE ADD COLUMN（容易触发 Citus 分布式死锁）；
        # 新字段随着导出完成后的换表流程一起生效。建表时不建索引（导出完成后由 exchange_tb
        # 从 master_tb 复制现有索引重建，整表建索引比逐行插入维护索引快很多）
        sql = f"""
            DROP TABLE IF EXISTS "public"."{export_tb}";
            CREATE TABLE "public"."{export_tb}" (
              "asin" varchar(20) COLLATE "pg_catalog"."default" NOT NULL,
              "monthly_sales" text COLLATE "pg_catalog"."default" NOT NULL,
              "zr_count" text COLLATE "pg_catalog"."default" NOT NULL,
              "sp_count" text COLLATE "pg_catalog"."default" NOT NULL,
              "total_count" text COLLATE "pg_catalog"."default" NOT NULL,
              "date_info_list" text COLLATE "pg_catalog"."default" NOT NULL,
              "created_at" timestamp(6) NOT NULL DEFAULT CURRENT_TIMESTAMP,
              "updated_at" timestamp(6) NOT NULL DEFAULT CURRENT_TIMESTAMP,
              "weekly_sales" text COLLATE "pg_catalog"."default",
              "weekly_views" text COLLATE "pg_catalog"."default",
              "monthly_views" text COLLATE "pg_catalog"."default",
              "variation_monthly_sales" text COLLATE "pg_catalog"."default" NOT NULL
            );
            COMMENT ON COLUMN "public"."{export_tb}"."asin" IS 'asin';
            COMMENT ON COLUMN "public"."{export_tb}"."monthly_sales" IS 'Amazon月销';
            COMMENT ON COLUMN "public"."{export_tb}"."zr_count" IS 'zr词数';
            COMMENT ON COLUMN "public"."{export_tb}"."sp_count" IS 'sp词数';
            COMMENT ON COLUMN "public"."{export_tb}"."total_count" IS '总词数';
            COMMENT ON COLUMN "public"."{export_tb}"."date_info_list" IS 'date_info';
            COMMENT ON COLUMN "public"."{export_tb}"."variation_monthly_sales" IS '变体销量月度历史(逗号分隔,对齐date_info_list)';

            SELECT create_distributed_table('{export_tb}', 'asin');
        """
        connection.execute(sql)

    # 导出脚本
    sh = CommonUtil.build_export_sh(
        site_name=site_name,
        db_type=db_type,
        hive_tb="dwt_amazon_report",
        export_tb=export_tb,
        col=[
            "asin",
            "monthly_sales",
            "zr_count",
            "sp_count",
            "total_count",
            "date_info_list",
            "variation_monthly_sales"
        ],
        partition_dict={
            "site_name": site_name,
            "date_type": date_type,
            "date_info": date_info
        }
    )

    client = SSHUtil.get_ssh_client()
    SSHUtil.exec_command_async(client, sh, ignore_err=False)
    client.close()

    print("导出完成，准备创建索引和交换表名")
    DBUtil.exchange_tb(engine, export_tb, master_tb, True)

    engine.dispose()
    print("success")

pass
