from typing import  Dict, List
import boto3
from config import configManager

class DataExtractor:
    def __init__(self, spark_session, source):
        self.spark_session = spark_session
        self.options= configManager().get_details(source)
        self.creds= self.options['creds']
        self.source = source
        
    def get_transformed(self,df):
        try:
            logic=configManager().get_config(f"{self.source.split('_')[0]}_transformation")
            logic = ','.join([f" {i[1]} as {i[0]}" for i in logic.items()])
            df.createOrReplaceTempView(f'__df__{self.source}')
            return  self.spark_session.sql(f"select {logic} from __df__{self.source}")
        except Exception as e:
            raise e
            #return logic
            
    def push_to_prod_d(df, source=None):
        ####
        """DataFrame
        Customer Data pushed to PROD D
        """
    
    def extract_data(self,columns :str =None, pushdown_predicate : str =None, limit :int=None ):
        # Connect to Snowflake using the AWS Snowflake Connector
        # Extract data from Snowflake using Glue
        """
        source eg: petco-snowflake, retailscape-aurora, glue etc
        columns : ','.join(list_of_columns)
        pushdown_predicate : source ='xyz' and cast(date as varchar) = '2023-01-08'
        """
        source = self.source
        columns = columns if columns else '*'
        pushdown_predicate= 'where '+pushdown_predicate if pushdown_predicate else ''
        limit = f'limit {str(limit)}' if limit else ''
        
        if 'glue' == source:
            
            query = f"SELECT {columns} FROM {self.options['database']}.{self.options['table']}  {pushdown_predicate} {limit}".strip()
            
            df= self.spark_session.sql(query)
            
        elif 'snowflake' in source:
            
            sfoptions = {
                          "sfURL" : self.creds['account']+'.snowflakecomputing.com',
                          "sfUser" : self.creds['user'],
                          "sfPassword" : self.creds['password'],
                          "sfDatabase" : self.options['database'],
                          "sfSchema" : self.options['schema'],
                          "sfWarehouse" : self.options['warehouse']
                        }
                        
            
            
            query = f"SELECT {columns} FROM {self.options['schema']}.{self.options['table']} {pushdown_predicate} {limit}".strip()
            
            df = self.spark_session.read.format('net.snowflake.spark.snowflake') \
                  .options(**sfoptions) \
                  .option("query",  query) \
                  .load()
        elif 'aurora' in source:
            
            table = self.options['table']
            database= self.options['database']
            try:
                host=self.creds['host']
            except:
                host=self.creds['base_url']
            try:
                port=self.creds['port']
            except:
                port='5432'
            try:
                user=self.creds['user']
            except:
                user='postgres'
            try:    
                password = self.creds['password']
            except:
                raise ValueError(f"Password Not Defined in secret")
                
            url= f'jdbc:postgresql://{host}:{port}/{database}'
       
            query = f"SELECT {columns} FROM {self.options['table']} {pushdown_predicate} {limit}".strip()
            
            
            df = self.spark_session.read.format("jdbc").option("url", url)\
                .option("user", user).option("password", password)\
                    .option("query", query).load()
                
            #df = self.spark_session.read.jdbc(url=AURORA_DATABASE_URL, table=query, properties=properties)
            
        return df