#=========================================
##           Import Dependencies           
#=========================================
import re
import string, tqdm
from time import time, sleep
from statistics import mode
from pytz import timezone
from datetime import datetime, timedelta
from pyspark.sql import SparkSession, functions as F, types as T

##              Pandas Config              
#=========================================
import pandas as pd
pd.options.display.html.table_schema=True
pd.options.display.max_columns=999
pd.options.display.max_rows=999

##              Spark Session              
#=========================================
spark = SparkSession\
  .builder\
  .appName("METADATA TABLES")\
  .config("spark.sql.crossJoin.enabled", "true")\
  .config("spark.dynamicAllocation.enabled", "false")\
  .config("spark.executor.instances", "6")\
  .config("spark.executor.cores", "5")\
  .config("spark.executor.memory", "8g")\
  .config("spark.yarn.executor.memoryOverhead", "4g")\
  .config("spark.ui.showConsoleProgress", "false")\
  .config("spark.network.timeout", 60)\
  .enableHiveSupport()\
  .getOrCreate()

##               Functions                 
#=========================================
def set_timer():
  global START_TIME
  START_TIME = time()

def get_timer(start_time=None):
  if start_time:
    t = datetime(1,1,1)+timedelta(seconds=int(time()-start_time))
  else:
    t = datetime(1,1,1)+timedelta(seconds=int(time()-START_TIME))
  return "{}:{}:{}".format(str(t.hour).zfill(2),str(t.minute).zfill(2),str(t.second).zfill(2))

def get_partition(schema, table):
  try:
    return spark.sql("""
     SHOW PARTITIONS {}.{}
     """.format(schema, table)).sort('partition').collect() # ambil partisi sesuai format
  except:
    return None

def get_partition_column_format(partition, sep="/"):
  list_partition = []
  for value in partition.split('/'):
    split_value = value.split("=")
    list_partition.append(split_value[0])
  return "{}".format(sep).join(list_partition)  

def get_partition_column(partition):
  list_partition = []
  for value in partition.split('/'):
    split_value = value.split("=")
    list_partition.append(split_value[0])
  return list_partition

def try_or(func, default=None, expected_exc=(Exception,)):
  try:
    return func()
  except expected_exc:
    return default

def get_content_summary(source_schema, source_table, partitioned=False, partition=None):

  path = lambda p: spark._jvm.org.apache.hadoop.fs.Path(p)
  fs = spark._jvm.org.apache.hadoop.fs.FileSystem.get(spark._jsc.hadoopConfiguration())

#  loc = "/user/hive/warehouse/{}.db/{}/".format(source_schema, source_table)
  loc = "/dev/{}.db/{}/".format(source_schema, source_table)
  
  list_data = list()
  list_file_count = list()
  list_length = list()
  list_blocksize = list()

  df_describe_formatted = try_or(lambda: spark.sql("""DESCRIBE FORMATTED {}.{}""".format(source_schema, source_table)))
  var_owner = try_or(lambda: df_describe_formatted.where("col_name='Owner'").collect()[0]["data_type"])
  var_created_time = try_or(lambda: datetime.strptime(df_describe_formatted.where("col_name='Created Time'").collect()[0]["data_type"], "%a %b %d %H:%M:%S %Z %Y"))
  var_last_access = try_or(lambda: datetime.strptime(df_describe_formatted.where("col_name='Last Access'").collect()[0]["data_type"], "%a %b %d %H:%M:%S %Z %Y"))
  var_created_by = try_or(lambda: df_describe_formatted.where("col_name='Created By'").collect()[0]["data_type"])
  var_type = try_or(lambda: df_describe_formatted.where("col_name='Type'").collect()[0]["data_type"])
  var_table_properties = try_or(lambda: df_describe_formatted.where("col_name='Table Properties'").collect()[0]["data_type"])
  var_transient_lastddltime = try_or(lambda: datetime.utcfromtimestamp(int({val.split("=")[0].strip():val.split("=")[1].strip() for val in var_table_properties.replace("[","").replace("]","").split(",")}["transient_lastDdlTime"])))
  var_location = try_or(lambda: df_describe_formatted.where("col_name='Location'").collect()[0]["data_type"])

  var_blocksize = 256
  
  if partitioned:

    var_partitioned = "Partitioned table"
    var_partition_col = try_or(lambda: get_partition_column_format(list_partition[0]['partition']))
    var_partition_len = try_or(lambda: len(partition))


    try:
      for row in partition: # melakukan loop atas list partisi
        hdfs_folder = fs.getContentSummary(path(loc+row['partition']))
        list_file_count.append(hdfs_folder.getFileCount())
        length = float(hdfs_folder.getLength())
        list_length.append(length/(1024**2))
        list_blocksize.append(int(1 if length/float(var_blocksize*(1024**2)) < 1 else length/float(var_blocksize*(1024**2))))
    except:
      list_file_count.append(None)
      list_length.append(None)
      list_blocksize.append(None)


    hdfs_folder = try_or(lambda: fs.getContentSummary(path(loc)))
    var_file_count = try_or(lambda: hdfs_folder.getFileCount())
    var_length = try_or(lambda: round(float(hdfs_folder.getLength()/(1024**2)),3))

    var_avg_file_count = try_or(lambda: int(1 if round(sum(list_file_count)/len(list_file_count)) < 1 else round(sum(list_file_count)/len(list_file_count))))
    var_avg_length = try_or(lambda: round(sum(list_length)/len(list_length),3))
    var_repartition = try_or(lambda: int(1 if round(sum(list_blocksize)/len(list_blocksize)) < 1 else round(sum(list_blocksize)/len(list_blocksize))))
    var_optimal = try_or(lambda: mode([True if abs(e1-e2)<4 else False for (e1, e2) in zip(list_file_count, list_blocksize)]))


  else:

    var_partitioned = "Not a partitioned table"
    var_partition_col = None
    var_partition_len = None

    hdfs_folder = try_or(lambda: fs.getContentSummary(path(loc)))
    list_file_count.append(try_or(lambda: hdfs_folder.getFileCount()))
    length = try_or(lambda: float(hdfs_folder.getLength()))
    list_length.append(try_or(lambda: length/(1024**2)))
    list_blocksize.append(try_or(lambda: int(1 if length/float(var_blocksize*(1024**2)) < 1 else length/float(var_blocksize*(1024**2)))))

    var_file_count = try_or(lambda: hdfs_folder.getFileCount())
    var_length = try_or(lambda: round(float(hdfs_folder.getLength()/(1024**2)),3))

    var_avg_file_count = try_or(lambda: int(1 if round(sum(list_file_count)/len(list_file_count)) < 1 else round(sum(list_file_count)/len(list_file_count))))
    var_avg_length = try_or(lambda: round(sum(list_length)/len(list_length),3))
    var_repartition = try_or(lambda: int(1 if round(sum(list_blocksize)/len(list_blocksize)) < 1 else round(sum(list_blocksize)/len(list_blocksize))))
    var_optimal = try_or(lambda: mode([True if abs(e1-e2)<4 else False for (e1, e2) in zip(list_file_count, list_blocksize)]))


  return [source_schema,source_table,
    var_owner,var_created_time,var_last_access,var_transient_lastddltime,var_created_by,
    var_type,var_table_properties,var_location,
    var_partitioned,var_partition_col,var_partition_len,
    var_length,var_avg_length,var_blocksize,var_file_count,var_avg_file_count,var_repartition,var_optimal]



#=========================================# 
#               Parameter                 #
#=========================================#

set_timer()

target_schema = "dev_ddb"
hive_schema, hive_table = "dev_ddb","bribrain_master_fact_metadata_tables_dev_ddb"

curr_date = datetime.now(timezone("Asia/Jakarta"))
ds = curr_date.strftime("%Y%m%d")
#ds = curr_date.strftime("%Y%m%d %H:%M:S")

print("Params `ds` =",ds)

##             Main Process                
#=========================================

df_schema = spark.sql("SHOW TABLES IN {}".format(target_schema))

#list_schema = df_schema.collect()

list_schema = df_schema.collect()

print("Total Record =",len(list_schema))

list_data = list()

for row in list_schema:

  source_schema, source_table = row["database"], row["tableName"]

  list_partition = get_partition(source_schema, source_table) # mendapatkan list partisi tabel

  if list_partition == None or len(list_partition) == 0: # kondisi ketika tabel tidak berpartisi
    data = get_content_summary(
      source_schema, 
      source_table) 
  else: # kondisi ketika tabel berpartisi
    data = get_content_summary(
      source_schema, 
      source_table,
      True,
      list_partition)

  list_data.append(data)


schema = T.StructType([
  T.StructField("schema", T.StringType(), True),
  T.StructField("table", T.StringType(), True),
  T.StructField("owner", T.StringType(), True),
  T.StructField("created_time", T.TimestampType(), True),
  T.StructField("last_access", T.TimestampType(), True),
  T.StructField("transient_lastddltime", T.TimestampType(), True),
  T.StructField("created_by", T.StringType(), True),
  T.StructField("type", T.StringType(), True),
  T.StructField("table_properties", T.StringType(), True),
  T.StructField("location", T.StringType(), True),
  T.StructField("partitioned", T.StringType(), True),
  T.StructField("partition_col", T.StringType(), True),
  T.StructField("partition_count", T.IntegerType(), True),
  T.StructField("length_in_mb_total", T.DoubleType(), True),
  T.StructField("length_in_mb_per_partition", T.DoubleType(), True),
  T.StructField("blocksize_in_mb_recommendation", T.IntegerType(), True),
  T.StructField("file_count_total", T.IntegerType(), True),
  T.StructField("file_count_per_partition", T.IntegerType(), True),
  T.StructField("repartition_recommendation", T.IntegerType(), True),
  T.StructField("is_hdfs_blocksize_optimal", T.BooleanType(), True)])

df_result = spark.createDataFrame(tuple(list_data), schema)\
  .withColumn("modified_date", F.lit(datetime.now(timezone("Asia/Jakarta"))),)\
  .withColumn("partition", F.lit(ds),)\
  .withColumn("batch", F.lit("01"),)

#df_result.toPandas()

df_result\
  .repartition(1)\
  .write.format("parquet")\
  .partitionBy("partition","batch")\
  .mode("overwrite")\
  .saveAsTable("{}.{}".format(hive_schema, hive_table)) # write as table
                  
print("Successfully write\n\t- {}.{}\n\t- target schema `{}`\n\t- {} record\n\t- running time {}".format(hive_schema, hive_table, target_schema, df_result.count(), get_timer()))

