import lib
from pyspark.sql import functions as F

class CampaignManager():
  def __init__(
    self, 
    spark, 
    table_campaign="project_sv.dgv_userpool_campaign_brimo_raw_cifno_daily"
  ):
    self.spark = spark
    self.table_campaign = table_campaign
    self.df_campaign = lib.read(self.spark, self.table_campaign)
    
  def _filter_campaign(self, campaign=None, code=None):
    if campaign and code:
      raise ValueError("Please provide only one: campaign or code.")
    elif code:
      print(f"Processing using campaign with code {code}.")
      df_campaign_filter = self.df_campaign.filter(F.lower(F.col("campaign")).like(f"%{code.lower()}%"))
    elif campaign:
      print(f"Processing using campaign {campaign}.")
      df_campaign_filter = self.df_campaign.filter(F.col("campaign") == campaign)
    else:
      raise ValueError("You must provide either a campaign or a code.")
      
    count = df_campaign_filter.count()
    if count == 0:
      raise ValueError("The campaign you are looking for does not exist.")
    

    return df_campaign_filter
    
  def get_campaign_info(self, campaign=None, code=None):
    df_campaign_filter = self._filter_campaign(campaign, code)
    df_campaign_info = (
      df_campaign_filter
      .select(
        "campaign", "channel", "specific_fiturid", "specific_jenis_fitur", 
        "specific_corpid", "date_start_campaign", "date_end_campaign",
        "blast_time", "group", "ds"
      )
      .distinct()
      .collect()
    )

    campaign_dict = [row.asDict() for row in df_campaign_info][0]

    col_list = ["specific_fiturid", "specific_jenis_fitur", "specific_corpid"]
    for col in col_list:
      try:
        campaign_dict[col] = campaign_dict[col].split(",")
      except Exception as e:
        campaign_dict[col] = []
    
    return campaign_dict
  
  def get_cif(self, campaign=None, code=None):
    df_campaign_filter = self._filter_campaign(campaign, code)
    df_campaign_cif = (
      df_campaign_filter
      .select("cifno")
      .distinct()
    )
    
    return df_campaign_cif