from typing import Dict

from pyspark.sql import SparkSession


class SparkDynamicConfig:
    def __init__(self, max_shuffle_partitions: int = None, min_shuffle_partitions: int = None):
        self.max_shuffle_partitions = max_shuffle_partitions
        self.min_shuffle_partitions = min_shuffle_partitions

    def use_min_shuffle_partitions(self, spark: SparkSession) -> 'SparkDynamicConfig':
        assert self.min_shuffle_partitions, 'must provide min_shuffle_partitions to use the conf'
        spark.conf.set("spark.sql.adaptive.enabled", 'false')
        spark.conf.set("spark.sql.shuffle.partitions", str(self.min_shuffle_partitions))
        spark.conf.set("spark.default.parallelism", str(self.min_shuffle_partitions))
        return self

    def use_max_shuffle_partitions(self, spark: SparkSession) -> 'SparkDynamicConfig':
        assert self.max_shuffle_partitions, 'must provide max_shuffle_partitions to use the conf'
        spark.conf.set("spark.sql.adaptive.enabled", 'false')
        spark.conf.set("spark.sql.shuffle.partitions", str(self.max_shuffle_partitions))
        spark.conf.set("spark.default.parallelism", str(self.max_shuffle_partitions))
        return self

    def use_adaptive_shuffle_partitions(self, spark: SparkSession) -> 'SparkDynamicConfig':
        spark.conf.set("spark.sql.adaptive.enabled", 'true')
        return self


def get_spark(app_name: str, conf: Dict = None):
    builder = SparkSession.builder.appName(app_name).enableHiveSupport()
    conf = conf or {}
    for k, v in conf.items():
        builder.config(k, v)

    spark = builder.getOrCreate()
    spark.conf.set('spark.sql.statistics.fallBackToHdfs', 'true')
    # 启用 Adaptive Execution ，从而启用自动设置 Shuffle Reducer 特性
    spark.conf.set("spark.sql.adaptive.enabled", 'true')
    # 设置每个 Reducer 读取的目标数据量，单位为字节。默认64M，一般改成集群块大小
    spark.conf.set("spark.sql.adaptive.shuffle.targetPostShuffleInputSize", '134217728')
    # 允许动态资源分配，配合 spark.dynamicAllocation.minExecutors, spark.dynamicAllocation.maxExecutors 等使用
    # spark 3.0+ 不允许动态设置以下两个参数
    import pyspark
    if str(pyspark.__version__).startswith('2.'):
        spark.conf.set("spark.dynamicAllocation.enabled", 'true')
        spark.conf.set("spark.shuffle.service.enabled", 'true')

    # spark.conf.set("hive.exec.dynamic.partition", "true")
    # default strict. In strict mode, the user must specify at least one static partition in case the user accidentally overwrites all partitions. In nonstrict mode all partitions are allowed to be dynamic.
    spark.conf.set("hive.exec.dynamic.partition.mode", "nonstrict")

    return spark


def clear_temp_views(spark: SparkSession):
    for table in spark.catalog.listTables('default'):
        if table.isTemporary:
            print(f'dropping temp view {table.name}')
            spark.catalog.dropTempView(table.name)

