from pyspark.sql import functions as F
from datetime import datetime
from pytz import timezone
import math
from .utils import set_timer,get_timer,get_time_now_wib
from .df_to_bridrive import download_hdfs_parts_and_cleanup

def delete_partition(spark, hive_schema, hive_table, col_partition, val_partition):
    """
    Deletes a partition from a Hive table if it exists.

    Parameters:
    --------
    - spark: Spark Session
    - hive_schema: Name of the Hive schema/database.
    - hive_table: Name of the Hive table.
    - col_partition: Name of the partition column.
    - val_partition: Value of the partition column to drop.

    Example:
    --------
    >>>delete_partition(spark, "datamart", "cust_fact", "ds", "20250807")
    >>>delete_partition(spark, "datamart", "cust_fact", "year", 2022)
    """
    
    tables_df = spark.sql(f"SHOW TABLES IN {hive_schema} LIKE '{hive_table}'")
    if tables_df.count() == 0:
        print(f"Table {hive_schema}.{hive_table} does not exist.")
        return
      
    if isinstance(val_partition, str):
        val_partition = f"'{val_partition}'"      
    partition_spec = f"{col_partition}={val_partition}"

    partitions_df = spark.sql(f"SHOW PARTITIONS {hive_schema}.{hive_table}")
    if partitions_df.filter(partitions_df.partition == partition_spec).count() == 0:
        print(f"Partition {partition_spec} does not exist in table {hive_schema}.{hive_table}.")
        return

    print(f"Dropping partition {partition_spec}...")
    spark.sql(f"ALTER TABLE {hive_schema}.{hive_table} DROP IF EXISTS PARTITION({partition_spec})")
    print(f"Partition {partition_spec} dropped successfully.")


def write_periodic(spark, df, hive_schema, hive_table, ds, sample_fraction=0.01, target_file_size_mb=256, col_partition = "ds"):
    """
    Write a DataFrame to a Hive table partitioned by 'ds' for periodic data ingestion.

    This function:
    - Adds metadata columns 'modified_date' and 'ds'
    - Estimates DataFrame size by sampling to determine optimal repartitioning
    - Drops existing partition for 'ds' if exists before appending new data
    - Writes data partitioned by 'ds' using parquet format
    - Refreshes and analyzes Hive table statistics for the updated partition

    Parameters:
    -----------
    spark : SparkSession
        Active Spark session object.
    df : DataFrame
        Spark DataFrame to write.
    hive_schema : str
        Hive schema/database name.
    hive_table : str
        Hive table name.
    ds : str
        Partition value (typically a date string) for the 'ds' partition column.
    sample_fraction : float, optional (default=0.01)
        Fraction of rows to sample when estimating average row size.
    target_file_size_mb : int, optional (default=256)
        Target file size in megabytes for repartitioning.
    col_partition : str
        Name for partition column
    """
    print(f"{get_time_now_wib()} - Start save data to {hive_schema}.{hive_table}")
    start = set_timer()
    
    if df.rdd.isEmpty():
        raise ValueError("dataframe empty")
        return

    # Add metadata columns
    df = df.withColumn("modified_date", F.lit(datetime.now(timezone("Asia/Jakarta")))) \
           .withColumn(col_partition, F.lit(ds))

    try:
        total_rows = df.count()
        sample_df = df.sample(False, sample_fraction, seed=42).cache()

        sample_data = sample_df.rdd.map(lambda row: len(str(row)))
        sampled_rows_and_size = sample_data.aggregate(
            (0, 0),  # (row_count, total_size_bytes)
            lambda acc, x: (acc[0] + 1, acc[1] + x),
            lambda acc1, acc2: (acc1[0] + acc2[0], acc1[1] + acc2[1])
        )

        sampled_rows, total_sample_size = sampled_rows_and_size

        avg_row_size = total_sample_size / sampled_rows if sampled_rows > 0 else 1024  # 1KB fallback

        estimated_size = total_rows * avg_row_size
        target_file_size_bytes = target_file_size_mb * 1024 * 1024
        repartition_number = max(1, int(math.ceil(estimated_size / target_file_size_bytes)))

        print(f"Estimated size: {estimated_size/1024/1024:.2f} MB, Partitions: {repartition_number}")

        # Repartition
        df = df.repartition(repartition_number)

        # Drop existing partition if exists
        if spark.sql(f"SHOW TABLES IN {hive_schema} LIKE '{hive_table}'").count() > 0:
            print(f"Dropping partition ds='{ds}' if exists...")
            spark.sql(f"ALTER TABLE {hive_schema}.{hive_table} DROP IF EXISTS PARTITION({col_partition}='{ds}')")

        # Write partitioned data
        df.write \
          .format("parquet") \
          .partitionBy(col_partition) \
          .mode("append") \
          .saveAsTable(f"{hive_schema}.{hive_table}")

        # Refresh table metadata and stats
        spark.sql(f"REFRESH TABLE {hive_schema}.{hive_table}")
        spark.sql(f"ANALYZE TABLE {hive_schema}.{hive_table} PARTITION ({col_partition}='{ds}') COMPUTE STATISTICS")
        print("Total Record = ", total_rows)
        print(f"{get_time_now_wib()} - Done save data to {hive_schema}.{hive_table} with time {get_timer(start)} for {col_partition} = {ds}")
    except Exception as e:
        print(f"Exception occurred: {e}")

def write_non_periodic(spark, df, hive_schema, hive_table, sample_fraction=0.01, target_file_size_mb=256):
    """
    Write a DataFrame to a Hive table without partitioning, overwriting existing data.

    This function:
    - Adds metadata column 'modified_date'
    - Estimates DataFrame size by sampling to determine optimal repartitioning
    - Overwrites the entire Hive table with the new data
    - Refreshes Hive table metadata after writing

    Parameters:
    -----------
    spark : SparkSession
        Active Spark session object.
    df : DataFrame
        Spark DataFrame to write.
    hive_schema : str
        Hive schema/database name.
    hive_table : str
        Hive table name.
    sample_fraction : float, optional (default=0.01)
        Fraction of rows to sample when estimating average row size.
    target_file_size_mb : int, optional (default=256)
        Target file size in megabytes for repartitioning.
    """
    print(f"{get_time_now_wib()} - Start save data to {hive_schema}.{hive_table}")
    start = set_timer()
    
    if df.rdd.isEmpty():
        raise ValueError("dataframe empty")
        return

    # Add metadata columns
    df = df.withColumn("modified_date", F.lit(datetime.now(timezone("Asia/Jakarta"))))
    try:
        total_rows = df.count()
        sample_df = df.sample(False, sample_fraction, seed=42).cache()

        sample_data = sample_df.rdd.map(lambda row: len(str(row)))
        sampled_rows_and_size = sample_data.aggregate(
            (0, 0),  # (row_count, total_size_bytes)
            lambda acc, x: (acc[0] + 1, acc[1] + x),
            lambda acc1, acc2: (acc1[0] + acc2[0], acc1[1] + acc2[1])
        )

        sampled_rows, total_sample_size = sampled_rows_and_size

        avg_row_size = total_sample_size / sampled_rows if sampled_rows > 0 else 1024

        estimated_size = total_rows * avg_row_size
        target_file_size_bytes = target_file_size_mb * 1024 * 1024
        repartition_number = max(1, int(math.ceil(estimated_size / target_file_size_bytes)))

        print(f"Estimated size: {estimated_size/1024/1024:.2f} MB, Partitions: {repartition_number}")

        # Repartition
        df = df.repartition(repartition_number)

        # Write partitioned data
        df.write \
          .format("parquet") \
          .mode("overwrite") \
          .saveAsTable(f"{hive_schema}.{hive_table}")

        # Refresh table metadata and stats
        spark.sql(f"REFRESH TABLE {hive_schema}.{hive_table}")
        print("Total Record = ", total_rows)
        print(f"{get_time_now_wib()} - Done save data to {hive_schema}.{hive_table} with time {get_timer(start)}")
    except Exception as e:
        print(f"Exception occurred: {e}")

def write_temp_periodic_more_than_one_partition(
    spark, df, hive_schema, hive_table,
    val_partition,
    col_partition=["ds"],
    repartition_number=1
):
    print(f"{get_time_now_wib()} - Start save data to {hive_schema}.{hive_table}")
    start = set_timer()

    if len(col_partition) != len(val_partition):
        raise ValueError("Length of col_partition and val_partition must be the same")

    for i, j in zip(col_partition, val_partition):
        df = df.withColumn(i, F.lit(j))

    df = df.repartition(repartition_number)

    if spark.sql(f"SHOW TABLES IN {hive_schema} LIKE '{hive_table}'").count() != 0:
        partition_spec = ", ".join([f"{col}='{val}'" for col, val in zip(col_partition, val_partition)])
        print(f"Dropping partition ({partition_spec}) if exists...")
        query_drop_partition = f"ALTER TABLE {hive_schema}.{hive_table} DROP IF EXISTS PARTITION({partition_spec})"
        spark.sql(query_drop_partition)

    df\
      .write\
      .format("parquet") \
      .partitionBy(*col_partition) \
      .mode("append") \
      .saveAsTable(f"{hive_schema}.{hive_table}")

    spark.sql(f"REFRESH TABLE {hive_schema}.{hive_table}")

    partition_str = ", ".join([f"{col}={val}" for col, val in zip(col_partition, val_partition)])
    df = spark.read.table(f"{hive_schema}.{hive_table}")
    
    for i, j in zip(col_partition, val_partition):
        df = df.filter(F.col(i) == j)
    
    count = df.count()
    print("Count Row  = {}".format(count))
    print(f"{get_time_now_wib()} - Done save data to {hive_schema}.{hive_table} with time {get_timer(start)} for partitions {partition_str}\n")
    
def write_temp_periodic(spark,df, hive_schema, hive_table, ds, repartition_number = 1, col_partition = "ds", add_modified_date = True, del_partition = True):
  print(f"{get_time_now_wib()} - Start save data to {hive_schema}.{hive_table}")
  start = set_timer()
  if add_modified_date:
    df = df\
      .withColumn("modified_date", F.lit(datetime.now(timezone("Asia/Jakarta"))))\
  
  df = df\
    .withColumn(col_partition,F.lit(ds))

  df = df.repartition(repartition_number)
  
  if del_partition == True:
    if spark.sql("SHOW TABLES IN {} LIKE '{}'".format(hive_schema, hive_table)).count() != 0:
      print(f"Dropping partition {col_partition}='{ds}' if exists...")
      query_drop_partition = "ALTER TABLE {}.{} DROP IF EXISTS PARTITION({}='{}')".format(hive_schema, hive_table, col_partition, ds)
      spark.sql(query_drop_partition) # menghapus partisi di spark

  df\
    .write\
    .format("parquet")\
    .partitionBy(col_partition)\
    .mode("append")\
    .saveAsTable("{}.{}".format(hive_schema,hive_table))

  spark.sql("ANALYZE TABLE {}.{} PARTITION ({}={}) COMPUTE STATISTICS".format(hive_schema, hive_table, col_partition, ds)) # kalkulasi data
  spark.sql("REFRESH TABLE {}.{}".format(hive_schema, hive_table)) # refresh tabel di spark
  df = spark.read.table(f"{hive_schema}.{hive_table}")
  df = df.filter(F.col(col_partition) == ds)
    
  count = df.count()
  print("Count Row  = {}".format(count))
  print(f"{get_time_now_wib()} - Done save data to {hive_schema}.{hive_table} with time {get_timer(start)} for {col_partition} = {ds}\n")
  #print("Count Row  = {}".format(count))

def write_temp(spark, df, hive_schema, hive_table, repartition_number=1):
  print(f"{get_time_now_wib()} - Start save data to {hive_schema}.{hive_table}")
  start = set_timer()
  
  #count = df.count()
  #if count!=0:
  df = df.withColumn("modified_date", F.lit(datetime.now(timezone("Asia/Jakarta"))))
  df = df.repartition(repartition_number)

  df\
    .write\
    .format("parquet")\
    .mode("overwrite")\
    .saveAsTable("{}.{}".format(hive_schema,hive_table))

  spark.sql("REFRESH TABLE {}.{}".format(hive_schema, hive_table)) # refresh tabel di spark
  df = spark.read.table(f"{hive_schema}.{hive_table}")
  count = df.count()
  print("Count Row  = {}".format(count))
  print(f"{get_time_now_wib()} - Done save data to {hive_schema}.{hive_table} with time {get_timer(start)}\n")
  #print("Count Row  = {}".format(count))
  
def quick_write_periodic(spark, df, hive_schema, hive_table, ds):
  df = df\
    .withColumn("ds", F.lit(ds))\
    .withColumn("modified_date", F.lit(datetime.now(timezone("Asia/Jakarta"))))

  if spark.sql("SHOW TABLES IN {} LIKE '{}'".format(hive_schema, hive_table)).count() != 0:
    query_drop_partition = "ALTER TABLE {}.{} DROP IF EXISTS PARTITION(ds='{}')".format(hive_schema, hive_table, ds)
    spark.sql(query_drop_partition) # menghapus partisi di spark

  df\
    .write\
    .format("parquet")\
    .partitionBy("ds")\
    .mode("append")\
    .saveAsTable("{}.{}".format(hive_schema,hive_table))

  df = spark.read.table(f"{hive_schema}.{hive_table}")\
    .filter(F.col("ds") == ds)
  
  count = df.count()
  
  print("Total Record =", count)
  
def write_to_csv(df, file_name, repartition_number = 1, option_delimiter = ","):
  """
  Write a Spark DataFrame to HDFS as CSV files and download the parts locally.

  Parameters:
  -----------
  df : pyspark.sql.DataFrame
      The Spark DataFrame to be saved and uploaded.
  base_file_name : str
      The base name used for naming downloaded CSV files.
  repartition_number : int, optional (default=1)
      The number of partitions to repartition the DataFrame before writing CSV files.
  option_delimiter : string, optional (default=";")
      Delimiter for every column

  Returns:
  --------
  None

  Example:
  --------
  >>> write_to_csv(df,"dataframe_daily")
  >>> write_to_csv(df,"dataframe_monthly", repartition_number=2, option_delimiter = ";")
  """
  if " " in file_name:
    raise ValueError(f"Invalid file_name: contains space. '{file_name}'")
  print("Generate file in HDFS.....")
  df\
    .repartition(repartition_number)\
    .write\
    .option("delimiter", option_delimiter)\
    .option("nullValue", "NULL") \
    .option("emptyValue", "NULL") \
    .csv(file_name, header=True, mode='overwrite')
  print("csv file success generated")
  download_hdfs_parts_and_cleanup(file_name)
  
def write_temp_periodic_with_partition(spark, df, hive_schema, hive_table, col_partition, repartition_number = 1,  add_modified_date = True):
  print(f"{get_time_now_wib()} - Start save data to {hive_schema}.{hive_table}")
  start = set_timer()
  if add_modified_date:
    df = df\
      .withColumn("modified_date", F.lit(datetime.now(timezone("Asia/Jakarta"))))\
  
  df = df.repartition(repartition_number)

  df\
    .write\
    .format("parquet")\
    .partitionBy(col_partition)\
    .mode("append")\
    .saveAsTable("{}.{}".format(hive_schema,hive_table))

  spark.sql("REFRESH TABLE {}.{}".format(hive_schema, hive_table)) # refresh tabel di spark
  df = spark.read.table(f"{hive_schema}.{hive_table}")
    
  count = df.count()
  print("Count Row  = {}".format(count))
  print(f"{get_time_now_wib()} - Done save data to {hive_schema}.{hive_table} with time {get_timer(start)}\n")
