from functools import reduce
from pyspark.sql import functions as F

def remove_blank_space(x):
  return F.when(F.trim(F.col(x)) != '', F.trim(F.col(x))).otherwise(F.lit(None))

def cleanse_blank_space(df_to_clean):
  col_to_clean = []
  for val in df_to_clean.dtypes:
    if val[1] == "string":
      col_to_clean.append(val[0])
  return reduce(lambda df, x: df.withColumn(x, remove_blank_space(x)), set(col_to_clean), df_to_clean)

def lower_columns_name(df):
  list_column = []
  for col in df.columns:
    list_column.append(col.lower())
  return df.select(list_column)

def read(spark,*table_names):
  """
  Reads one or multiple tables into Spark DataFrames.

  Parameters:
  - spark: SparkSession object
  - table_names: one or more string names of tables to read

  Returns:
  A dictionary of {table_name: DataFrame} if multiple tables,
  or a single DataFrame if only one table is passed.

  Example usage:
  >>> df = read(spark, table1)
  >>> df1, df2 = read(spark, table1, table2)
  """
  read_dataframe = []
  for table_name in table_names:
    df = spark.read.table(table_name)
    df = lower_columns_name(cleanse_blank_space(df))
    read_dataframe.append(df)
  if len(read_dataframe) == 1:
      return read_dataframe[0]
  return tuple(read_dataframe)

def get_table_dict(): 
  """ Return dictionary of code → table name mapping. """
  table_dict = {}
  with open("/home/cdsw/lib/list_table.txt", "r") as f:
      for line in f:
          line = line.strip()
          if not line or line.startswith("#"):
              continue
          if "=" not in line:
              continue  # skip malformed lines
          code, table_name = line.split("=", 1)
          table_dict[code.strip()] = table_name.strip()
  return table_dict

def read_default_tables(spark, *table_code_list):
    """
    Load default tables based on provided table codes from a mapping text file.

    The mapping file should contain lines of the form:
        code=table_name
    Lines starting with '#' or blank lines are ignored.

    Parameters:
    -----------
    spark : SparkSession
        Active Spark session used for reading tables.
    *table_code_list : str
        One or more table codes to load.
    mapping_file : str, optional (default="table_mapping.txt")
        Path to the table mapping text file.

    Returns:
    --------
    DataFrame or tuple of DataFrames
        DataFrame(s) corresponding to the provided table codes, 
        read via `read(spark, *table_names)`.

    Raises:
    -------
    KeyError
        If any provided table code is not found in the mapping file.

    Example:
    --------
    >>> df_savingmaster = read_default_tables(spark, "savingmaster")
    >>> (df_savingmaster, df_fitur_holding) = read_default_tables(spark, "savingmaster", "fitur_holding")
    """
    table_dict = get_table_dict()

    try:
        table_name_list = [table_dict[code] for code in table_code_list]
    except KeyError as e:
        raise KeyError(f"Invalid table code: {e.args[0]}. "
                       f"Available codes: {list(table_dict.keys())}")

    return read(spark, *table_name_list)
