__all__ = ['aws_athena_query']

from ._imports import *
from ._config import *
from ._amazons3 import s3_convert_file_to_df

class Aws_athena:

    def __init__(self, db, query, env) -> None:
        self.param = self._get_athena_params(db, query, env)
    
    def _get_athena_params(self, db, query, env):
        param = {}
        param ['query'] = query
        param ['athena db'] =  db  
        param['env'] = env
        param ['session'] =  boto3.Session()       
        param ['athena output folder'] = f's3://amica-aws-eds-{env}-datascience/athena_temp/' if env !='dl' else f's3://amica-eds-edsc/athena_temp/'
        param ['athena output bucket'] = param ['athena output folder'].split('/')[2]
        param ['athena output prefix'] = 'athena_temp'
        #param ['athena output bucket S3 path'] = r"s3://" + param ['athena output bucket'] + '/'        
        return param 
    
    def get_athena_query(self):        
        env = self.param['env']
        try:        
            athena_client = self.param ['session'].client('athena', region_name='us-east-2',
                                        aws_access_key_id=config[f'aws_{env}']['aws_access_key_id'],
                                        aws_secret_access_key=config[f'aws_{env}']['aws_secret_access_key'])
            
            response = athena_client.start_query_execution(
                    QueryString = self.param['query'],
                        QueryExecutionContext={
                        'Database': self.param['athena db']
                        },
                        ResultConfiguration={
                        'OutputLocation': self.param['athena output folder'],
                        }
                )            
            self.filename = response['QueryExecutionId'] 
            query_status = None
            while query_status == 'QUEUED' or query_status == 'RUNNING' or query_status is None:
                query_status = athena_client.get_query_execution(QueryExecutionId=self.filename)['QueryExecution']['Status']['State']
                if query_status == 'FAILED' or query_status == 'CANCELLED':
                    raise Exception('Athena query failed or was cancelled')
                time.sleep(3)
            #query complete  
            key= self.param ['athena output prefix'] + '/' + self.filename + '.csv'                  
            df = s3_convert_file_to_df(self.param ['athena output bucket'], key , env)            
            return df
        except Exception as e:
            print(e) 

    
    
    def athena_cleanup(self, env):
        s3 = self.param ['session'].resource('s3',                               
            aws_access_key_id=config[f'aws_{env}']['aws_access_key_id'],
            aws_secret_access_key=config[f'aws_{env}']['aws_secret_access_key'],           
                        )
        my_bucket = s3.Bucket(self.param['athena output bucket'])
        for item in my_bucket.objects.filter(Prefix=self.param['athena output prefix']):
            item.delete()        
        #restore the athena_temp folder
        object = s3.Object(self.param['athena output bucket'], self.param['athena output prefix'])
        object.put()

def aws_athena_query(db, query, env):
    '''
    get data in df from athena query
    exp1
    db = 'rravichandran_test_db'
    query = 'select * FROM "rravichandran_test_db"."documents"'
    env = 'dev'
    df = aws_athena_query(db, query, env)
    exp2
    db = 'eds_ds_ptc'
    query =  'SELECT * FROM "eds_ds_ptc"."prediction_history" limit 10;'
    env = 'dl'
    aws_athena_query(db, query, env)
    '''
    aa = Aws_athena(db, query, env)
    df = aa.get_athena_query()
    aa.athena_cleanup(env)  #clean temp folder    An error occurred (AccessDenied) when calling the DeleteObject operation: Access Denied
    return df