
import os
import sys
import contextlib

@contextlib.contextmanager
def suppress_native_output():
  with open(os.devnull, 'w') as devnull:
    old_stdout_fd = os.dup(1)
    old_stderr_fd = os.dup(2)
    os.dup2(devnull.fileno(), 1)
    os.dup2(devnull.fileno(), 2)
    try:
      yield
    finally:
      os.dup2(old_stdout_fd, 1)
      os.dup2(old_stderr_fd, 2)


from pyspark.sql import SparkSession
    
def sparkSession(
  appname="Bribrain Packages SparkSession", 
  executor="small", 
  instances=2, 
  cores=2, 
  memory=4, 
  overhead=2, 
  verbose=False,
  extra_configs=None):
  
  # Define resource configuration
  if   executor=="small":
    config = [2, 2, 4, 2]
  elif executor=="medium":
    config = [4, 4, 6, 2]
  elif executor=="high":
    config = [6, 6, 8, 4]
  elif executor=="custom":
    config = [instances, cores, memory, overhead]
  else:
    raise ValueError(f"Unknown executor type: {executor}")

  executor_instances, executor_cores, executor_memory, memory_overhead = config
  verbose_str = "true" if verbose else "false"

  # Start building SparkSession
  builder = (
      SparkSession.builder
      .appName(appname)
      .config("spark.sql.crossJoin.enabled", "true")
      .config("spark.dynamicAllocation.enabled", "false")
      .config("spark.executor.instances", f"{executor_instances}")
      .config("spark.executor.cores", f"{executor_cores}")
      .config("spark.executor.memory", f"{executor_memory}g")
      .config("spark.yarn.executor.memoryOverhead", f"{memory_overhead}g")
      .config("spark.sql.broadcastTimeout", "36000")
      .config("spark.ui.showConsoleProgress", verbose_str)
      .config("spark.network.timeout", "300s")
  )
  
  # Add extra configs if provided
  if extra_configs:
      for key, value in extra_configs.items():
          builder = builder.config(key, value)

  with suppress_native_output():
    # Enable Hive and create session
    spark = builder.enableHiveSupport().getOrCreate()

  return spark