#=============================================================
###  CREATED AT     : 12 MARET 2024                         
###  UPDATED AT     : 03 FEBRUARI 2024                          
###  COPYRIGHT      : BRIBRAIN DATA ENGINEER TEAM           
###  DESCRIPTION    : Module untuk kumpulan function        
#=============================================================
import re
import time
from pytz import timezone
from datetime import datetime, timedelta
from pyspark.sql import functions as F
from pyspark.sql import types as T

import pandas as pd
from typing import Optional, List

from pandas import DatetimeTZDtype
from pandas.api.types import (
    is_datetime64_any_dtype,
    is_datetime64tz_dtype,
    is_integer_dtype,
    is_float_dtype,
    is_bool_dtype,
    is_string_dtype,
    is_categorical_dtype,
)



#=============================================================#
def try_or(func, default=None, expected_exc=(Exception,)):
  """Menangkap error dan memberikan keluaran sesuai parameter

  Args:
      func (function): python function
        (notes is ditambahkan lambda pada sebelum nama function, ex: try_or(lambda:func))
      default (object): keluaran yang diharapkan ketika terjadi error
      expected_exc (Exception): Exception yang diharapkan
    
  Returns:
      object: keluaran dari function atau mengembalikan keluaran sesuai input jika terdapat error
  """
  
  try:
    return func()
  except expected_exc:
    return default
  
#=============================================================# 

START_TIME = None

def set_timer():
  """Menetapkan waktu awal untuk penghitungan durasi proses

  Args:
      -
    
  Returns:
      float: waktu dalam format float
  """
  global START_TIME
  START_TIME = time.time()
  # We return it just in case, but usually, we just want to 'set' it.
  return START_TIME

#=============================================================#
def get_timer(start_time=None):
  """Memperoleh durasi proses berdasarkan waktu awal dikurangi waktu sekarang

  Args:
      start_time (float) : timestamp yang menunjukan time
    
  Returns:
      string: durasi proses dengan format HH:MM:SS
  """
    
  if START_TIME is None:
        return "Timer was not started. Call set_timer() first."
    
  if start_time:
    return (datetime(1,1,1)+timedelta(seconds=int(time.time()-start_time))).strftime("%H:%M:%S")
  
  return (datetime(1,1,1)+timedelta(seconds=int(time.time()-START_TIME))).strftime("%H:%M:%S")

#=============================================================#
def show_time(format_time="%Y-%m-%d %H:%M:%S"):
  """Menampilkan waktu saat ini

  Args:
      format_time (str) : format waktu yang ingin ditampilkan 
          (default is %Y-%m-%d %H:%M:%S)

  Returns:
      string: waktu saat ini sesuai format input
  """
      
  # Get the current time as a UNIX timestamp
  timestamp = time.time()

  # Convert the timestamp to a timezone-aware datetime object
  local_time = datetime.fromtimestamp(timestamp, timezone('Asia/Jakarta'))

  # Format the time as Year-Month-Day Hour:Minute:Second with timezone
  formatted_time = local_time.strftime(format_time)

  return formatted_time

#=============================================================#  
def spark_type_from_pandas_dtype(dtype, series: Optional[pd.Series] = None) -> T.DataType:
  """Menentukan tipe data Spark berdasarkan dtype dari pandas

  Args:
      dtype : tipe data kolom pandas (mis. numpy.dtype atau pandas ExtensionDtype)
      series (pandas.Series, optional): data kolom asli, dapat digunakan untuk
          logika tambahan (mis. cek min/max), saat ini opsional dan belum
          dimanfaatkan secara spesifik

  Returns:
      pyspark.sql.types.DataType: tipe data Spark yang paling sesuai dengan dtype pandas
  """

  # datetime (naive atau timezone-aware)
  if is_datetime64_any_dtype(dtype) or is_datetime64tz_dtype(dtype):
      return T.TimestampType()

  # boolean (termasuk pandas BooleanDtype)
  if is_bool_dtype(dtype):
      return T.BooleanType()

  # integer (termasuk nullable Int64/Int32)
  if is_integer_dtype(dtype):
      # LongType paling aman untuk range nilai besar
      return T.LongType()

  # float (termasuk nullable Float64)
  if is_float_dtype(dtype):
      return T.DoubleType()

  # string / kategori → StringType
  if is_string_dtype(dtype) or is_categorical_dtype(dtype):
      return T.StringType()

  # fallback untuk object / mixed type
  return T.StringType()


def pandas_to_spark(pdf: pd.DataFrame, tz_to_utc: bool = True):
  """Mengonversi pandas DataFrame menjadi Spark DataFrame dengan schema otomatis

  Args:
      pdf (pandas.DataFrame): DataFrame pandas yang akan dikonversi ke Spark
      tz_to_utc (bool): jika True, seluruh kolom datetime yang memiliki timezone
          akan dikonversi ke zona waktu UTC dan dibuat naive (tanpa informasi TZ)
          untuk menjaga konsistensi di Spark
          (default is True)

  Returns:
      pyspark.sql.DataFrame: DataFrame Spark dengan schema yang diturunkan dari
          dtype masing-masing kolom di pandas
  """

  # Bekerja pada copy agar tidak mengubah pdf asli di pemanggil
  pdf = pdf.copy()

  # Normalisasi datetime yang timezone-aware → naive UTC
  if tz_to_utc:
      for col in pdf.columns:
          dtype = pdf[col].dtype
          if isinstance(dtype, DatetimeTZDtype) or is_datetime64tz_dtype(dtype):
              pdf[col] = (
                  pdf[col]
                  .dt.tz_convert("UTC")   # konversi ke UTC
                  .dt.tz_localize(None)   # hilangkan informasi timezone
              )

  # Bangun schema Spark berdasarkan dtype tiap kolom
  fields: List[T.StructField] = []
  for col in pdf.columns:
      dtype = pdf[col].dtype
      spark_type = spark_type_from_pandas_dtype(dtype, pdf[col])
      fields.append(T.StructField(col, spark_type, nullable=True))

  schema = T.StructType(fields)

  # Buat Spark DataFrame dengan schema eksplisit
  return spark.createDataFrame(pdf, schema=schema)

#=============================================================#  
def get_list_partition(spark, schema, table):
  """Memperoleh list partisi dari hive table yang diurutkan dari partisi terbaru

  Args:
      spark (pyspark.sql.session.SparkSession): spark session
      schema (str): nama schema dari table di hive
      table (str): nama table di hive
    
  Returns:
      list: list partisi diurutkan dari partisi yang terbaru
  """
  
  try:
    partitions = spark.sql("""
     SHOW PARTITIONS {}.{}
     """.format(schema, table)).sort("partition", ascending=False).collect() # ambil partisi sesuai format
    if len(partitions) != 0: # jika ada partisi
      list_partition = []
      for row in partitions:
        if "__HIVE_DEFAULT_PARTITION__" not in row[0]:
          arrange = []
          dict_partition = {}
          for partition in row[0].split("/"):
            value = partition.split("=")
            arrange.append(value[1].zfill(2))
            dict_partition[value[0]] = value[1]
          dict_partition["__formatted_partition"] = "|".join(arrange)
          list_partition.append(dict_partition)
      list_partition = sorted(list_partition, key=lambda row: row['__formatted_partition'], reverse=True)
      return list_partition
    else: # selain itu
      return None # tidak ada partisi
  except:
    print("is not a partitioned table")
    return None
  
#=============================================================#
def get_first_partition(spark, schema, table):
  """Memperoleh partisi pertama dari hive table

  Args:
      spark (pyspark.sql.session.SparkSession): spark session
      schema (str): nama schema dari table di hive
      table (str): nama table di hive
    
  Returns:
      dict: partisi pertama
  """
    
  partitions = get_list_partition(spark, schema, table)
  if partitions == None:
    return None
  
  return partitions[-1]

#=============================================================#  
def get_last_partition(spark, schema, table):
  """Memperoleh partisi terakhir dari hive table

  Args:
      spark (pyspark.sql.session.SparkSession): spark session
      schema (str): nama schema dari table di hive
      table (str): nama table di hive
    
  Returns:
      dict: partisi terakhir
  """
    
  partitions = get_list_partition(spark, schema, table)
  if partitions == None:
    return None
  
  return partitions[0]


#=============================================================#
def check_partition_exists(spark, target_value, schema, table, mode=None):
    """Mengecek keberadaan partisi pada tabel HDFS dan mengembalikan metadata partisi terkait

    Args:
        spark (pyspark.sql.session.SparkSession): objek spark session yang digunakan
        target_value (str): nilai partisi yang ingin dicek (contoh: '20230307')
            *gunakan pemisah '|' jika partiti kolom lebih dari satu, (contoh: '2023|03|07')
        schema (str): nama database/schema di Hive/HDFS
        table (str): nama table yang akan dicek partisinya
        mode (str): mode pencarian jika target_value tidak ditemukan 
            (optional is None, 'forward', 'backward')
            (default is None)

    Returns:
        dict: dictionary berisi detail kolom partisi (year, month, day, dsb.)
            tanpa kolom __formatted_partition, atau None jika tidak ditemukan
    """
    partition_list = get_list_partition(spark, schema, table)
    
    if not partition_list:
        return None

    # 1. Get a sample format from the metadata
    sample_format = partition_list[0].get('__formatted_partition', '')
    
    # 2. STRICT VALIDATION: Compare structure
    # We build the regex pattern outside the f-string to avoid SyntaxError
    digit_pattern = re.sub(r'\d', r'\\d', sample_format)
    escaped_pattern = digit_pattern.replace('|', r'\|')
    regex_final = "^" + escaped_pattern + "$"
    
    if not re.match(regex_final, target_value):
        raise ValueError(
            "Format Mismatch!\n"
            f"Input target_value: '{target_value}'\n"
            f"HDFS Expected Format: '{sample_format}'\n"
            "Please ensure separators (|, -, /) and padding (01 vs 1) match exactly."
        )

    # 3. Build lookup
    lookup = {}
    for p in partition_list:
        fmt_key = p.get('__formatted_partition')
        if fmt_key:
            # Clean dict for output by removing the helper key
            clean_dict = {k: v for k, v in p.items() if k != '__formatted_partition'}
            lookup[fmt_key] = clean_dict

    # 4. Logic Execution
    sorted_keys = sorted(lookup.keys())

    if target_value in lookup:
        return lookup[target_value]

    if mode == "forward":
        for k in sorted_keys:
            if k > target_value:
                return lookup[k]
                
    elif mode == "backward":
        for k in reversed(sorted_keys):
            if k < target_value:
                return lookup[k]

    return None
  
  
#=============================================================#

import os

#=============================================================#
def put_to_bridrive(file_path, upload_url):
  """Mengirimkan file dari CDSW ke BRIDrive

  Args:
      upload_url (str): upload link dari Seafile (BRIDrive) untuk upload file
      file_path (str): path file yang akan di upload dari CDSW ke BRIDrive
    
  Returns:
      -
  """

  res = os.system("hdfs dfs -get /user/bribrain_sac/packages/seaf-share.py")

  script = "PYTHONHTTPSVERIFY=0 python seaf-share.py put {}".format(" ".join([upload_url, file_path]))
  os.system(script)
  
  os.remove("/home/cdsw/seaf-share.py")
  
  
#=============================================================# 
def put_to_hdfs(file_path, hdfs_path = "/tmp/development/bribrain"):
  """Mengirimkan file dari CDSW ke HDFS

  Args:
      file_path (str): path file yang akan di upload dari CDSW ke HDFS
      hdfs_path (str): path lokasi penyimpanan data di HDFS
          (default is /tmp/development/bribrain)
    
  Returns:
      -
  """
  command = "hdfs dfs -put {0} {1}".format(file_path, hdfs_path)
  res = os.system(command)
  if res == 0:
    print("success - put: `{0}` to `{1}`".format(file_path, hdfs_path))

    
#=============================================================#
def get_from_hdfs(file_path, cdsw_path = "/home/cdsw/"):
  """Mengambil file dari HDFS ke CDSW

  Args:
      file_path (str): path file yang akan dipindahkan dari HDFS ke CDSW
      cdsw_path (str): path penyimpanan file di HDFS yang ingin dipindahkan
          (default is "/home/cdsw/")
    
  Returns:
      -
  """
  command = "hdfs dfs -get {0} {1}".format(file_path, cdsw_path)
  res = os.system(command)
  if res == 0:
    print("success - get: `{0}` to `{1}`".format(file_path, cdsw_path))


#=============================================================#
def list_file_hdfs(hdfs_path = "/user/{}/".format(os.environ["HADOOP_USER_NAME"])):
  """Memperlihatkan list file yang ada pada lokasi HDFS

  Args:
      hdfs_path (str): path HDFS yang ingin di show
          (default is "/user/{username}/")
    
  Returns:
      -
  """
  command = "hdfs dfs -ls {}".format(hdfs_path)
  res = os.system(command)



