__all__ = [
    's3_upload_file', 's3_save_and_upload_file', 's3_save_and_upload_csv', 's3_save_and_upload_json',
    's3_download_file', 's3_download_and_read_file', 's3_download_and_read_csv', 's3_download_and_read_json',
    's3_list_folder_contents', 's3_check_file_exists', 's3_download_folder_contents',
    's3_delete_file', 's3_sql_query_csv', 's3_upload_file_to_silver', 's3_get_folder_list', 's3_create_folder_list', 
    's3_compare_file_dates_in_silver', 's3_get_prediction_file_in_silver', 's3_remove_temp_local_download']

from ._imports import *
from ._config import config
from ._awsbase import Aws_base
from ._datetime import DtFind

class AmazonS3(Aws_base):
    clients = {}

    @classmethod
    def get_client(cls, env):
        return super()._get_client(env, 's3')

def client(env):
    return AmazonS3.get_client(env)

@atexit.register
def s3_close():
    for client in AmazonS3.clients.values():
        client.close()

def s3_upload_file(local_file_path, s3_file_path=None, overwrite=False, env='dev'):
    if s3_file_path == None:
        s3_file_path  = os.path.basename(local_file_path)
    if not overwrite and s3_file_path in s3_list_folder_contents(os.path.dirname(s3_file_path), env=env):
        raise Exception('File Already Exists')
    else:
        client(env).upload_file(local_file_path, config[f'aws_{env}']['bucket'], s3_file_path)

def s3_save_and_upload_file(data, s3_file_path, local_file_name='output', overwrite=False, save_file_kwargs={}, env='dev'):
    _, s3_file_extension = os.path.splitext(s3_file_path)
    local_folder_path    = config['local']['local_folder']
    local_file_name      = os.path.basename(local_file_name) + '.' + s3_file_extension
    local_file_path      = os.path.join(local_folder_path, local_file_name)
    if s3_file_extension == '.csv':
        data.to_csv(local_file_path, index=False, **save_file_kwargs)
    elif s3_file_extension == '.parquet':
        data.to_parquet(local_file_path, index=False, **save_file_kwargs)
    elif s3_file_extension == '.json':
        with open(local_file_path, 'w') as f:
            json.dump(data, f, indent=4, **save_file_kwargs)
    elif s3_file_extension == '.pkl':
        with open(local_file_path, 'wb') as f:
            pickle.dump(data, f)
    s3_upload_file(local_file_path, s3_file_path, overwrite=overwrite, env=env)
    os.remove(local_file_path)

def s3_save_and_upload_csv(df, local_file_name='output.csv', s3_file_path=None, overwrite=False, env='dev'):
    local_folder_path = config['local']['local_folder']
    local_file_name = os.path.basename(local_file_name)
    local_file_path = os.path.join(local_folder_path, local_file_name)
    df.to_csv(local_file_path, index=False)
    s3_upload_file(local_file_path, s3_file_path, overwrite=overwrite, env=env)
    os.remove(local_file_path)

def s3_save_and_upload_json(data, local_file_name='output.json', s3_file_path=None, overwrite=False, env='dev'):
    local_folder_path = config['local']['local_folder']
    local_file_name = os.path.basename(local_file_name)
    local_file_path = os.path.join(local_folder_path, local_file_name)
    with open(local_file_path, 'w') as f:
        json.dump(data, f, indent=4)
    s3_upload_file(local_file_path, s3_file_path, overwrite=overwrite, env=env)
    os.remove(local_file_path)
    
def s3_download_file(s3_file_path, local_file_path=None, s3_bucket=None, env='dev'):
    if s3_bucket is None:
        s3_bucket = config[f'aws_{env}']['bucket']
    if local_file_path is None:
        local_folder_path = config['local']['local_folder']
        local_file_name = os.path.basename(s3_file_path)
        local_file_path = os.path.join(local_folder_path, local_file_name)
    try:
        client(env).download_file(s3_bucket, s3_file_path, local_file_path)
    except botocore.exceptions.ClientError as e:
        if e.response['Error']['Code'] == '404':
            raise Exception('File Does Not Exist')
        else:
            raise

def s3_download_and_read_file(s3_file_path, local_file_path=None, s3_bucket=None, read_file_kwargs={}, env='dev'):
    _, s3_file_extension = os.path.splitext(s3_file_path)
    if local_file_path is None:
        local_folder_path = config['local']['local_folder']
        local_file_name = os.path.basename(s3_file_path)
        local_file_path = os.path.join(local_folder_path, local_file_name)
    s3_download_file(s3_file_path, local_file_path, s3_bucket=s3_bucket, env=env)
    if s3_file_extension == '.csv':
        d = pd.read_csv(local_file_path, **read_file_kwargs)
    elif s3_file_extension == '.xlsx':
        d = pd.read_excel(local_file_path, **read_file_kwargs)
    elif s3_file_extension == '.parquet':
        d = pd.read_parquet(local_file_path, **read_file_kwargs)
    elif s3_file_extension == '.json':
        with open(local_file_path) as f:
            d = json.load(f)
    #os.remove(local_file_path)
    return d

def s3_download_and_read_csv(s3_file_path, local_file_path=None, s3_bucket=None, read_csv_kwargs={}, env='dev'):
    if local_file_path is None:
        local_folder_path = config['local']['local_folder']
        local_file_name = os.path.basename(s3_file_path)
        local_file_path = os.path.join(local_folder_path, local_file_name)
    s3_download_file(s3_file_path, local_file_path, s3_bucket=s3_bucket, env=env)
    df = pd.read_csv(local_file_path, **read_csv_kwargs)
    os.remove(local_file_path)
    return df

def s3_download_and_read_json(s3_file_path, local_file_path=None, s3_bucket=None, env='dev'):
    if local_file_path is None:
        local_folder_path = config['local']['local_folder']
        local_file_name = os.path.basename(s3_file_path)
        local_file_path = os.path.join(local_folder_path, local_file_name)
    s3_download_file(s3_file_path, local_file_path, s3_bucket=s3_bucket, env=env)
    with open(local_file_path) as f:
        d = json.load(f)
    os.remove(local_file_path)
    return d

def list_objects_with_paginator (bucket, prefix, env):
    '''
    list all file/folder within a prefix, can handle over 1000 files with paginator
    This is fix original function which is not able to handle over 1000 files in folder
    Parameters
    ----------
    Bucket : str 
                S3 bucket name.
    Prefix : str 
                S3 prefix.
    env : str 
                the env you are running.
    Returns
    ----------
    None : list 
                erturn list of folders/files path.
    Examples
    ----------
    list_objects_with_paginator (bucket_name, prefix, 'dev')
    '''
    s3_client = client(env)
    paginator = s3_client.get_paginator('list_objects')
    contents = []
    for page in paginator.paginate(Bucket=bucket, Prefix=prefix):        
        if 'Contents' in page:
            for obj in page['Contents']:
                contents.append (obj['Key'])
    return contents

def s3_list_folder_contents(s3_folder_path, s3_bucket=None, pattern=r'^.*$', env='dev'):
    '''
    this funcition able to filter file/folder by regex expression 
    Parameters
    ----------
    s3_folder_path : str 
                s3 prefix.
    s3_bucket : str 
                s3 bucket name.
    pattern : str 
                regex pattern you want to match with file/folder name.
    env : str 
                environment you are running.
    Returns
    ----------
    contents : list 
                list of s3 path meeting the condition.
    Examples
    ----------
    s3_list_folder_contents(r'cash_flow_forecasting/model_assets', s3_bucket=None, pattern=f'^.*cash_flow_model_{'vendor'}_20250407_090547.csv$', env='dev')
    '''
    if s3_bucket is None:
        s3_bucket = config[f'aws_{env}']['bucket']
    contents = []
    for item in list_objects_with_paginator(bucket=s3_bucket, prefix=s3_folder_path, env =env):
        if re.match(pattern, item):
            contents.append(item)
    return contents


# def s3_list_folder_contents(s3_folder_path, s3_bucket=None, pattern=r'^.*$', env='dev'):
#     if s3_bucket is None:
#         s3_bucket = config[f'aws_{env}']['bucket']
#     contents = []
#     for item in client(env).list_objects(Bucket=s3_bucket, Prefix=s3_folder_path)['Contents']:
#         if re.match(pattern, item['Key']):
#             contents.append(item['Key'])
#     return contents

def s3_check_file_exists(s3_file_path, env='dev'):
    s3_folder_path = os.path.dirname(s3_file_path)
    if s3_file_path in s3_list_folder_contents(s3_folder_path, env=env):
        return True
    return False

def s3_download_folder_contents(s3_folder_path, local_folder_path=None, s3_bucket=None, pattern=r'^.*\.csv$', env='dev'):
    if local_folder_path is None:
        local_folder_path = config['local']['local_folder']
    for s3_file_path in s3_list_folder_contents(s3_folder_path, s3_bucket=s3_bucket, pattern=pattern, env=env):
        s3_file_name = os.path.basename(s3_file_path)
        local_file_path = os.path.join(local_folder_path, s3_file_name)
        s3_download_file(s3_file_path, local_file_path, s3_bucket=s3_bucket, env=env)

def s3_delete_file(s3_file_path, s3_bucket=None, env='dev'):
    if s3_bucket is None:
        s3_bucket = config[f'aws_{env}']['bucket']
    try:
        client(env).delete_object(Bucket=s3_bucket, Key=s3_file_path)
    except botocore.exceptions.ClientError as e:
        if e.response['Error']['Code'] == '404':
            raise Exception('File Does Not Exist')
        else:
            raise

def s3_sql_query_csv(s3_file_path, s3_query='select * from s3object', s3_bucket=None, env='dev'):
    if s3_bucket is None:
        s3_bucket = config[f'aws_{env}']['bucket']
    resp = client(env).select_object_content(
        Bucket = s3_bucket,
        Key = s3_file_path,
        Expression=s3_query,
        ExpressionType='SQL',    
        InputSerialization={
            "CSV": {
                'FileHeaderInfo': 'USE',
                'AllowQuotedRecordDelimiter': True,
            },
        },
        OutputSerialization={
            'CSV': {}
        },
    )

    records = []
    for event in resp['Payload']:
        if 'Records' in event:
            records.append(event['Records']['Payload'])
        elif 'Stats' in event:
            stats = event['Stats']['Details']

    file_str = ''.join(r.decode('utf-8') for r in records)
    df = pd.read_csv(io.StringIO(file_str),  header=None)
    return df

def s3_create_single_folder(env= 'dev', s3_bucket = None, prefix = r'code_test/test/'):
    '''
    this is use to programmatically generate folder
    s3_create_single_folder(env= 'dev', s3_bucket = None, prefix = r'code_test/test/')
    '''
    if not s3_bucket:
        s3_bucket = config[f'aws_{env}']['bucket']
    client(env).put_object(Bucket=s3_bucket, Key=(prefix))
    return 

def s3_convert_file_to_df ( bucket, file_path, env):
    
    s3_client = boto3.client('s3',
                            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']
                            )
    obj = s3_client.get_object(Bucket=bucket, Key=file_path)
    output = pd.read_csv(io.BytesIO(obj['Body'].read()),index_col=False)
    return output

def s3_upload_file_to_silver(local_file_path, s3_file_path=None, overwrite=False, env='dev'):
    
    if s3_file_path == None:
        s3_file_path  = os.path.basename(local_file_path)
    if not overwrite and s3_file_path in s3_list_folder_contents(os.path.dirname(s3_file_path), env=env):
        raise Exception('File Already Exists')
    else:
        client(env).upload_file(local_file_path,f'amica-aws-eds-{env}-dl-silver' , s3_file_path)


#copy folder structure
def s3_get_folder_list( env = 'dl', s3_folder_path = "archive/", s3_bucket = None,  pattern = r'^.*$'):
    if not s3_bucket:        
        s3_bucket = config[f'aws_{env}']['bucket']

    contents = set()    
    for item in client(env).list_objects(Bucket=s3_bucket, Prefix=s3_folder_path)['Contents']:
        contents.add(os.path.dirname(item['Key']))  #this only get folder name
    return contents

def s3_create_folder_list (folder_list, env = 'dl', original_project_prefix = 'archive', new_project_prefix = "archive1", s3_bucket = 'amica-eds-edsc'):
    #https://stackoverflow.com/questions/1939743/amazon-s3-boto-how-to-create-a-folder
    # s3 = boto3.resource('s3')
    # bucket = s3.Bucket('amica-eds-edsc')
    #project_path = r"archive1/"
    '''
    archive_list2 = s3_get_folder_list( env = 'dl', s3_folder_path = r"archive/drtest/", s3_bucket = None,  pattern = r'^.*$')
    
    '''
    for k in folder_list:
        #remove the old name
        k = k.replace(original_project_prefix, new_project_prefix, 1)  #only replace the first occurance
        prefix = f'{k}/'        
        client(env).put_object(Bucket=s3_bucket, Key=(prefix))
    return 

def s3_delete_all_objects_from_folder(bucket_name = "amica-eds-edsc", prefix = r"archive/drtest2/", env = 'dl'):
    """
    This function deletes all files in a folder from S3 bucket for troubleshooting purpose
    s3_delete_all_objects_from_folder(bucket_name = "amica-eds-edsc", prefix = r"archive/drtest2/", env = 'dl')
    """
    
    s3_client = client(env)
    # First we list all files in folder
    response = s3_client.list_objects_v2(Bucket=bucket_name, Prefix=prefix)
    files_in_folder = response["Contents"]
    files_to_delete = []
    # We will create Key array to pass to delete_objects function
    for f in files_in_folder:
        files_to_delete.append({"Key": f["Key"]})
    # This will delete all files in a folder
    s3_client.delete_objects(
        Bucket=bucket_name, Delete={"Objects": files_to_delete}
    )
    return 

#this is for checking if we need to rerun on business day 
def s3_compare_file_dates_in_silver(env='dev'):
    '''
    this function is used to avoid repeated run in prod because accounting required using business first day, 
    we need to schedule run prediction multi time to accomodate holidy and weekend. To check if prediction has been run  
    check prediction file in silver latest time vs expected time to determine next step
    actual_file_year_month, expected_file_year_month = compare_file_dates_in_silver(env='dev')
    print(actual_file_year_month, expected_file_year_month)
    '''
    today = dt.datetime.now()
    objects = client(env).list_objects(Bucket=f'amica-aws-eds-{env}-dl-silver', Prefix = r'cashflow_forecast/forecast_stage/') 
    df = pd.DataFrame(objects["Contents"])
    actual_file_year_month =  (DtFind.first_day_last_month(today).strftime('%Y-%m') if not df[~df['Key'].str.endswith(r'/')]['LastModified'].to_list()  #if no file, get last month 
                            else df['LastModified'].dt.tz_convert('US/Eastern').apply(pd.to_datetime).max().strftime("%Y-%m") #if contain file, get max date )

    )
    expected_file_year_month = today.strftime("%Y-%m")
    return actual_file_year_month, expected_file_year_month

def s3_get_prediction_file_in_silver (s3_folder_path = r"cashflow_forecast/forecast_stage/", s3_bucket = f'amica-aws-eds-dev-dl-silver', pattern = r'^.*\.parquet$', env = 'dev'):
    file_list = []
    for item in client(env).list_objects(Bucket=s3_bucket, Prefix=s3_folder_path)['Contents']:
        if item['Key'][-1] != "/"  and  re.match(pattern, item['Key']):     
            file_list.append(item['Key'])
    return file_list

def s3_remove_temp_local_download (s3_file_path):
    '''
    this help remove local file downloaded from s3, based on s3 path    
    '''
    local_folder_path = config['local']['local_folder']
    local_file_path = os.path.join(local_folder_path, os.path.basename(s3_file_path))
    os.remove(local_file_path)
    return