import os
import shutil
from typing import Any, Dict

from pyspark.sql import SparkSession


class LocalSpark:
    spark: SparkSession = None
    __conf: Dict = {}

    @staticmethod
    def stop():
        if LocalSpark.spark:
            LocalSpark.spark.stop()
            LocalSpark.spark = None

    @staticmethod
    def get(conf: Dict[str, Any] = None, clean_existing_data: bool = True) -> SparkSession:
        conf = conf or {}
        if LocalSpark.spark is None:
            default_conf = {
                'spark.default.parallelism': 4,
                'hive.exec.dynamic.partition.mode': 'nonstrict',
                'spark.sql.warehouse.dir': '/tmp/spark-warehouse-localdw-ut',
                'spark.driver.extraJavaOptions': "-Dderby.system.home=/tmp/spark-warehouse-metastore-ut "
                                                 "-Dderby.stream.error.file=/tmp/spark-warehouse-metastore-ut.log"
            }
            default_conf.update(conf)
            conf = default_conf

            if clean_existing_data:
                # delete old spark warehouse/metastore dir
                print(f"removing dir {conf['spark.sql.warehouse.dir']}")
                shutil.rmtree(conf['spark.sql.warehouse.dir'], ignore_errors=True)
                if '-Dderby.system.home' in conf['spark.driver.extraJavaOptions']:
                    import re
                    java_options = re.sub(r'\s*=\s*', '=', conf['spark.driver.extraJavaOptions'].strip()).split()
                    for op in java_options:
                        if '-Dderby.system.home' == op.split('=')[0].strip():
                            print(f"removing dir {op.split('=')[1].strip()}")
                            shutil.rmtree(op.split('=')[1].strip(), ignore_errors=True)

            # ensure a local spark with default config
            os.environ['SPARK_CONF_DIR'] = '/tmp/local-spark-conf-ut'
            spark_builder = SparkSession.builder.appName('UnitTest').master('local[4]')

            print('using conf: ', conf)
            for k, v in conf.items():
                spark_builder.config(k, v)
            LocalSpark.spark = spark_builder.enableHiveSupport().getOrCreate()

        spark = LocalSpark.spark
        spark.catalog.clearCache()
        for table in spark.catalog.listTables('default'):
            if table.isTemporary:
                print(f'dropping temp view {table.name}')
                spark.catalog.dropTempView(table.name)

        return LocalSpark.spark
