import contextlib
import pandas as pd
import shutil, os
import logging
from time import time
import psutil
import sys
from pyspark.sql import SparkSession
from pyspark.sql import DataFrame as SparkDataFrame
import pyspark.sql.functions as F
from datetime import datetime, timezone
from functools import reduce
from pyspark.sql.functions import udf
from typing import List, Dict
from pyspark.sql.types import StructType, StructField, StringType, IntegerType, BooleanType, TimestampType, DoubleType, DecimalType, ArrayType, BinaryType, LongType
import boto3
from requests import session

class sparkClient():
    """
        !!! env is currently passed via **kwargs but this is a mandatory field! Should be refactored 
        
        In case we want to read straight from s3 into spark assumed_role kwarg needs to be provided
    
    """

    def __init__(self, s3_client, s3_bucket, s3_target_prefix, partition_start, **kwargs):
        self.s3_bucket = s3_bucket
        self.s3_prefix = s3_target_prefix
        self.partition_date = partition_start
        self.s3_client = s3_client
        self.env = kwargs['env']
        # self.assumed_role = kwargs['assumed_role']
        self.spark = self._create_spark_session()
        return None

    def _get_temp_batch_credentials(self ):
        # we cannot fetch sts temp credentials via an assumed role, hence, we get the current credentials of the assumed role
        # https://stackoverflow.com/questions/36287720/boto3-get-credentials-dynamically
        session = boto3.Session()
        credentials = session.get_credentials()
        credentials = credentials.get_frozen_credentials()
        access_key = credentials[0]
        secret_key = credentials[1]
        session_token = credentials[2]
        return access_key, secret_key, session_token

    def _create_spark_session(self):
        """
           Create a spark session on a local cluster (ie. single node) using nbr_of_cores cores.

           In case spark runs on AWS batch, we need to use temporary AWS credentials in order to read from 
           S3. In case of local dev we can use default AWS profile. 
        """
         # https://stackoverflow.com/questions/50891509/apache-spark-codegen-stage-grows-beyond-64-kb 
        nbr_cores = self._get_nbr_cores()

        try:
            if self.env == 'local':
                logging.info('Building spark session w ProfileCredentialsProvider for S3 access')
                spark = SparkSession.builder.master(f"local[{nbr_cores}]") \
                            .appName(f'Spark_{self.s3_prefix}') \
                            .config('spark.sql.codegen.wholeStage', 'false') \
                            .config("spark.sql.session.timeZone", "UTC") \
                            .config('spark.jars.packages', 'org.apache.hadoop:hadoop-aws:3.2.0')\
                            .config("spark.sql.legacy.parquet.datetimeRebaseModeInRead", "LEGACY") \
                            .config("fs.s3a.aws.credentials.provider","com.amazonaws.auth.profile.ProfileCredentialsProvider")\
                            .getOrCreate()

            else:
                access_key, secret_key, session_token = self._get_temp_batch_credentials()
                # https://stackoverflow.com/questions/54223242/aws-access-s3-from-spark-using-iam-role?noredirect=1&lq=1
                logging.info('Building spark session w TemporaryAWSCredentialsProvider for S3 access')
                spark = SparkSession.builder.master(f"local[{nbr_cores}]") \
                            .appName(f'Spark_{self.s3_prefix}') \
                            .config('spark.sql.codegen.wholeStage', 'false') \
                            .config("spark.sql.session.timeZone", "UTC") \
                            .config('spark.jars.packages', 'org.apache.hadoop:hadoop-aws:3.2.0')\
                            .config("spark.sql.legacy.parquet.datetimeRebaseModeInRead", "LEGACY") \
                            .config('fs.s3a.aws.credentials.provider', 'org.apache.hadoop.fs.s3a.TemporaryAWSCredentialsProvider')\
                            .config('fs.s3a.access.key', access_key)\
                            .config('fs.s3a.secret.key', secret_key)\
                            .config('fs.s3a.session.token', session_token)\
                            .getOrCreate()  

            spark.sparkContext.setLogLevel("ERROR")
            logging.info(f"Spark Session Spark_{self.s3_prefix} created")
        except Exception as e:
            logging.error("Issue creating Spark Session. Exiting...")
            logging.error(e)
            sys.exit(1)
        return spark

    def _get_nbr_cores(self) -> int:
       return psutil.cpu_count(logical = False)

    def read_single_csv_file(self, file:str, schema: Dict[str, StructField] = None, header: str = "true", delimiter: str = ";") -> SparkDataFrame:
        file = file.replace('s3://', 's3a://')
        try: 
            df = self.spark.read.format("csv").option("header",header).option("delimiter", delimiter).load(file)
            logging.info(f"Read {self.s3_prefix} in spark DF")
            if schema:
                new_df = self.enforce_schema(df, schema)
                logging.info("Enforced schema")
            else:
                new_df = df
        except Exception as e:
            logging.error("Issue reading parquet Exiting...")
            logging.error(e)
            sys.exit(1)
        return new_df

    def read_single_parquet_file(self, file: str, schema: Dict[str, StructField] = None) -> SparkDataFrame:
        file = file.replace('s3://', 's3a://')
        try: 
            df = self.spark.read.format("parquet").load(file)
            logging.info(f"Read {self.s3_prefix} in spark DF")
            if schema:
                new_df = self.enforce_schema(df, schema)
                logging.info("Enforced schema")
            else:
                new_df = df
        except Exception as e:
            logging.error("Issue reading parquet Exiting...")
            logging.error(e)
            sys.exit(1)
        return new_df

    def read_multiple_csv_files(self, files: List[str], schema: Dict[str, StructField] = None, header: str = "true", delimiter: str = ";") -> SparkDataFrame:
        try:
            list_of_dfs = []
            for file in files:
                file = file.replace('s3://', 's3a://')
                df = self.spark.read.format("csv").option("header",header).option("delimiter", delimiter).load(file)
                if schema:
                    new_df = self.enforce_schema(df, schema)
                    list_of_dfs.append(new_df)
                else:
                    list_of_dfs.append(df)
            united_df = reduce(self._unite_dfs, list_of_dfs)
            logging.info("Succesfully read all files in a spark DF")
        except Exception as e:
            logging.error("Issue reading parquet files and unioning them in spark DF Exiting...")
            logging.error(e)
            sys.exit(1)
        return united_df
    
    def read_multiple_parquet_files(self, files: List[str], schema: Dict[str, StructField] = None) -> SparkDataFrame:
        try:
            list_of_dfs = []
            for file in files:
                file = file.replace('s3://', 's3a://')
                df = self.spark.read.format("parquet").load(file)
                if schema:
                    new_df = self.enforce_schema(df, schema)
                    list_of_dfs.append(new_df)
                else:
                    list_of_dfs.append(df)
            united_df = reduce(self._unite_dfs, list_of_dfs)
            logging.info("Succesfully read all files in a spark DF")
        except Exception as e:
            logging.error("Issue reading parquet files and unioning them in spark DF Exiting...")
            logging.error(e)
            sys.exit(1)
        return united_df

    def enforce_schema(self, df: SparkDataFrame, schema: Dict[str, StructField]) -> SparkDataFrame:    
        try:
            new_df = df.select([F.col(c).cast(schema[c]).alias(c) for c in df.columns])
            logging.info("Succesfully applied schema")
        except Exception as e:
            logging.error("Issue applying schema. Exiting...")
            logging.error(e)
            sys.exit(1)
        return new_df

    def _unite_dfs(self, df1: SparkDataFrame, df2: SparkDataFrame) -> SparkDataFrame:
        return df1.unionByName(df2)

    def add_metadata(self, df:SparkDataFrame, file_name = None ) -> SparkDataFrame:
        if not file_name: file_name = F.input_file_name()
        print(f'file_name {file_name}')
        try: 
            df = df.withColumn('META_file_name', F.lit(file_name)) \
                        .withColumn('META_partition_date', F.lit(datetime.strptime(self.partition_date, "%Y-%m-%d"))) \
                        .withColumn('META_processing_date_utc', F.lit(datetime.now(timezone.utc)))
            logging.info("Succesfully added metadata")
        except Exception as e:
            logging.error("Issue adding metadata. Exiting...")
            logging.error(e)
            sys.exit(1)
        return df

    @staticmethod
    @udf
    def udf_unicode(x):
        if x is None: 
            res = None
        else:
            try:
                res = x.encode("ascii","ignore")
            except AttributeError:
                res = x
        return res

    def convert_version(self, df: SparkDataFrame) -> SparkDataFrame: 
        if 'version' in df.columns:
            try: 
                new_df = df.select(*[self.udf_unicode(column).alias('version') if column == 'version' else column for column in df.columns])
                logging.info("Succesfully converted lynx version column to ascii.")
            except Exception as e:
                logging.error("Issue converting lynx version column.Exiting...")
                logging.error(e)
                sys.exit(1)
        else: new_df = df
        return new_df

    def reindex_cols(self, df: SparkDataFrame, columns_order: List[str]) -> SparkDataFrame:
        try:
            new_df = df.select(*columns_order)
            logging.info("Succesfully reindexed columns")
        except Exception as e:
            logging.error("Issue reindexing columns. Exiting...")
            logging.error(e)
            sys.exit(1)
        return new_df

    def apply_schema(self, df:SparkDataFrame , schema) -> SparkDataFrame:
        return self.spark.createDataFrame(df.collect(), schema = schema)
        # return self.spark.createDataFrame(df.take(100), schema = schema)

    @staticmethod
    @udf
    def udf_float_to_int(x):
        if x is None:
            res = None
        else:
            try:
                res = int(x)
            except AttributeError:
                res = x
        return res

    @staticmethod
    @udf
    def udf_string_to_int(x):
        if x is None or x =='':
            res = None
        else:
            try:
                # https://stackoverflow.com/questions/1841565/valueerror-invalid-literal-for-int-with-base-10 
                res = int(x)
            except AttributeError:
                res = x
        return res

    @staticmethod
    @udf
    def udf_string_to_timestamp(x):
        if x is None or x =='':
            res = None
        else:
            x = x.split('.')[0]
            date_format = '%Y-%m-%d %H:%M:%S'
            try:
                # https://stackoverflow.com/questions/1841565/valueerror-invalid-literal-for-int-with-base-10 
                res = datetime.strptime(x, date_format)
            except AttributeError:
                res = x
        return str(res)

    def convert_float_to_int_int(self, df: SparkDataFrame, cols: list) -> SparkDataFrame:
        """
            When reading from SQL, columns which are INT but contain only NULLS get converted to parquet float columns.
            We create a pyspark schema based on the original SQL schema, hence INT column gets translated into StructField('col', IntegerType(), True).
            In order to apply this pyspark schema, we need to convert the relevant int columns back to int.
        """
        try: 
            new_df = df.select(*[self.udf_float_to_int(column).cast(IntegerType()).alias(column) if column in cols else column for column in df.columns])
            logging.info("Succesfully converted float columns back to int")
        except Exception as e:
            logging.error("Issue converting float columns to int. Exiting...")
            logging.error(e)
            sys.exit(1)
        return new_df

    def convert_string_to_int_int(self, df: SparkDataFrame, cols: list) -> SparkDataFrame:
        """
            Convert string cols to int
        """
        try: 
            new_df = df.select(*[self.udf_string_to_int(column).cast(IntegerType()).alias(column) if column in cols else column for column in df.columns])
            logging.info("Succesfully converted str columns to int")
        except Exception as e:
            logging.error("Issue converting str columns to int. Exiting...")
            logging.error(e)
            sys.exit(1)
        return new_df

    def convert_string_to_timestamp(self, df: SparkDataFrame, cols: list) -> SparkDataFrame:
        """
            Convert string cols to int
        """
        try: 
            print('cols')
            print(cols)
            new_df = df.select(*[self.udf_string_to_timestamp(column).cast(TimestampType()).alias(column) if column in cols else column for column in df.columns])
            logging.info("Succesfully converted string columns to timestamp")
        except Exception as e:
            logging.error("Issue converting str columns to timestamp. Exiting...")
            logging.error(e)
            sys.exit(1)
        return new_df

    def write_to_parquet(self, spark_df: SparkDataFrame, max_records_per_file = 100000):
        local_path = './data'
        with contextlib.suppress(Exception):
            shutil.rmtree(local_path)
        try:
            spark_df.write.option('maxRecordsPerFile', max_records_per_file).mode('overwrite').parquet(local_path)
            for file in os.listdir(local_path):
                local_file = f'{local_path}/{file}'
                if '.crc' not in file and 'SUCCESS' not in file:
                    self.s3_client.upload_local_file(local_file, self.s3_prefix)
                os.remove(local_file)
            try:
                os.rmdir(local_path)
            except Exception:
                None
            logging.info(f'Succesfully wrote {self.s3_prefix}')
        except Exception as e:
            logging.error(f'Issue creating parquet file: {self.s3_prefix}. Exiting...')
            logging.error(e)
            sys.exit(1)
        return 

    def spark_df_to_pandas_df(self, df: SparkDataFrame) -> pd.DataFrame:
        return df.toPandas()


    def clean_old_dates(self, spark_df:SparkDataFrame, cols_to_clean:list) -> SparkDataFrame:
        """
            Spark 3.0 has difficulties working with dates older than 1900-01-01.

            org.apache.spark.SparkUpgradeException: You may get a different result due to the upgrading of Spark 3.0: reading dates before 1582-10-15 or timestamps before 1900-01-01T00:00:00Z from Parquet files can be ambiguous, as the files may be written by Spark 2.x or legacy versions of Hive, which uses a legacy hybrid calendar that is different from Spark 3.0+'s Proleptic Gregorian calendar. See more details in SPARK-31404. You can set spark.sql.legacy.parquet.datetimeRebaseModeInRead to 'LEGACY' to rebase the datetime values w.r.t. the calendar difference during reading. Or set spark.sql.legacy.parquet.datetimeRebaseModeInRead to 'CORRECTED' to read the datetime values as it is.
        
            Adding datetimeRebaseModeInRead still results in 'mktime argument out of range' errors. Hence, we'll manually clean these dates.
        """
        new_df = spark_df.select([
            F.when(F.col(column) < "1900-01-01 00:00:00", F.lit("1900-01-01 00:00:00"))\
            .otherwise(F.col(column))\
            .cast(TimestampType())\
            .alias(column) if column in cols_to_clean else column for column in spark_df.columns
            ]
        )
        return new_df

    def pandas_df_to_spark_df(self, pandas_df: pd.DataFrame, schema=None) -> SparkDataFrame:
        if schema:
            df_spark = self.spark.createDataFrame(pandas_df, schema=schema)
        else:
            df_spark = self.spark.createDataFrame(pandas_df)
        return df_spark

    def nan_to_null(self, df: SparkDataFrame) -> SparkDataFrame:
        """
            In case you have string columns with NaN strings ( eg after reading pd.read_csv) which you want 
            to load as NULL in redshift, apply this function. 
        """
        for col in df.columns:
            df = df.withColumn(col, F.when(F.col(col) == 'NaN', None).otherwise(F.col(col)))
        return df
        



