import sys
from awsglue.transforms import *
from awsglue.utils import getResolvedOptions
from pyspark.context import SparkContext
from awsglue.context import GlueContext
from awsglue.job import Job
from awsglue.dynamicframe import DynamicFrame
import time
from datetime import datetime, timedelta
import requests
import re
from pyspark.sql.functions import *
import boto3
import datetime
from credentials import CredentialManager
from dataextractor import DataExtractor

args = getResolvedOptions(sys.argv, ['JOB_NAME'])
glueContext = GlueContext(SparkContext.getOrCreate())
spark = glueContext.spark_session
job = Job(glueContext)
job.init(args['JOB_NAME'], args)


spark.conf.set("spark.sql.optimizer.pushedFilters", "true")

def snow():
    options={'table':'CUSTOMER_CATALOG',
             'warehouse':'COMPUTE', 'schema':'BUNGEE', 'database':'PETCO_BUNGEE_SHARE_INT'
             }
    source ="petco-snowflake"
    cred_manager = CredentialManager()
    pushdown_predicate=  "cast(product_id as varchar) in ('2778692', '2254733', '1354906', '3186808')"
    data_extractor = DataExtractor( glueContext,cred_manager)
    df = data_extractor.extract_data(source, options, pushdown_predicate, columns='product_id')
    #{'user': 'us_east_1', 'password': 'ntc*>*tS>3,",HM"', 'account': 'gztlgpj-us_east_1'}
    
    #print('petco count', df.count())
    return df


def aurora():
    source ="retailscape-aurora"
    options={'table':'product_gaps_chewy_20230102', 'database':'assortment_analytics'}
    
    cred_manager = CredentialManager()
    pushdown_predicate=  "cast(sku as varchar) in ('3338349', '3079172', '967750', '159352')"
    
    data_extractor = DataExtractor( glueContext,cred_manager)
    df = data_extractor.extract_data(source, options, pushdown_predicate, columns ='bungee_product_id, source, sku, merch1')
    return df

def glue():
    source ="glue"
    options={'table':'final_ingestion', 'database':'assortment'}
    
    cred_manager = CredentialManager()
    pushdown_predicate=  "load_Date = cast('2023-01-04' as date ) and cast(sku as varchar) in ('3338349', '3079172', '967750', '159352')"
    
    data_extractor = DataExtractor( glueContext,cred_manager)
    df = data_extractor.extract_data(source, options, pushdown_predicate, columns ='bungee_product_id, source, sku, merch1')
    return df




s= datetime.datetime.now()
df= snow()
print(datetime.datetime.now()-s)
df.show()
print(df.select('product_id').collect())

s= datetime.datetime.now()
df= aurora()
print(datetime.datetime.now()-s)
df.show()
print(df.select('bungee_product_id', 'sku').collect())

s= datetime.datetime.now()
df= glue()
print(datetime.datetime.now()-s)
df.show()
print(df.select('source', 'sku').collect())




job.commit()