# AUTOGENERATED! DO NOT EDIT! File to edit: s3bz.ipynb (unless otherwise specified).

__all__ = ['S3', 'ExtraArgs', 'generalSave', 'generalLoad', 'saveZl', 'loadZl', 'savePklZl', 'loadPklZl', 'saveRaw',
           'loadRaw', 'Requests', 'createBucket', 'deleteBucket']

# Cell
from botocore.config import Config
from nicHelper.wrappers import add_method, add_class_method, add_static_method
from botocore.errorfactory import ClientError
import bz2, json, boto3, logging, requests, zlib, pickle

# Cell
class S3:
  @staticmethod
  def s3(region = 'ap-southeast-1', user = None, pw = None, accelerate = True, **kwargs):
    '''
    create and return s3 client
    '''
    logging.info(f'using {("standard","accelerate")[accelerate]} endpoint')
    config = Config(s3={"use_accelerate_endpoint": accelerate,
                        "addressing_style": "virtual"})
    s3 = boto3.client(
        's3',
        aws_access_key_id= user,
        aws_secret_access_key= pw,
        region_name = region,
        config = config
      )
    return s3
  @classmethod
  def saveFile(cls, key, path, bucket = '', **kwargs):
    '''save a file to s3'''
    s3 = cls.s3(**kwargs)
    result = s3.upload_file(path, bucket, key, **kwargs)
    return result
  @classmethod
  def loadFile(cls, key, path, bucket = '', useUrl = False, **kwargs):
    '''load file from s3'''
    if useUrl:
      print('using url')
      url = cls.presign(key=key, bucket=bucket, checkExist = False)
      r = requests.get(url)
      if r.status_code == 200:
        print('presign success')
        with open(path, 'wb') as f:
          f.write(r.content)
        return True
      else:
        print('presign failed')
    s3 = cls.s3(**kwargs)
    result = s3.download_file(bucket,key, path )
    return result
  @classmethod
  def deleteFile(cls, key, bucket, **kwargs):
    s3 = cls.s3(**kwargs)
    result = s3.delete_object(Bucket=bucket, Key=key)
    return result

  @classmethod
  def save(cls,  key, objectToSave, bucket = '',**kwargs):
    '''
    save an object to s3
    '''
    s3 = cls.s3(**kwargs)
    compressedData = bz2.compress(json.dumps(objectToSave).encode())
    result = s3.put_object(Body=compressedData, Bucket=bucket, Key=key)
    success = result['ResponseMetadata']['HTTPStatusCode'] ==  200
    logging.info('data was saved to s3')
    if not success: raise Error(success)
    else: return True
  @classmethod
  def exist(cls, key, bucket, **kwargs):
    s3 = cls.s3(**kwargs)
    res = s3.list_objects_v2(Bucket=bucket, Prefix=key, MaxKeys=1)
    return 'Contents' in res

  @classmethod
  def load(cls, key, bucket='',fileName = '/tmp/tempFile.bz', useUrl = False, **kwargs):
#     if not cls.exist(key, bucket, **kwargs):
#       logging.info('object doesnt exist')
#       return {}
#     logging.info('object exists, loading')
    s3 = cls.s3(**kwargs)
    try:
      cls.loadFile(key=key, path=fileName, bucket=bucket, useUrl=useUrl)
#       s3.download_file(bucket,key, fileName )
    except Exception as e:
      print(f'error downloading file {e}')
    with open (fileName , 'rb') as f:
      allItemsByte = f.read()
    if not allItemsByte: raise ValueError('all data does not exist in the database')
    allItems = json.loads(bz2.decompress(allItemsByte).decode())
    return allItems

  @classmethod
  def presign(cls, key, expiry = 1000, bucket = '', checkExist = False,**kwargs):
    if checkExist:
      if not cls.exist(key,bucket=bucket,**kwargs): return 'object doesnt exist'
    s3 = cls.s3(**kwargs)
    result = s3.generate_presigned_url(
        'get_object',
          Params={'Bucket': bucket,
                  'Key': key},
        ExpiresIn=expiry)
    return result
  @classmethod
  def presignUpload(cls, bucket, key, expiry = 1000, **kwargs):
    '''
    # usage of the presigned url
    with open(object_name, 'rb') as f:
        files = {'file': (object_name, f)}
        http_response = requests.post(response['url'], data=response['fields'], files=files)
    # If successful, returns HTTP status code 204
    '''
    s3 = cls.s3(**kwargs)
    return s3.generate_presigned_post(bucket, key, ExpiresIn = expiry)

  @classmethod
  def loadDataFrame(cls, bucket, key,path='/tmp/tmpfile.csv',**kwargs):
    import pandas as pd
    cls.loadFile(key=key, path=path,bucket=bucket, **kwargs)
    return pd.read_csv(path)
  @classmethod
  def saveDataFrame(cls,bucket,key,df,path='/tmp/tmpfile.csv', **kwargs):
    df.to_csv(path)
    return cls.saveFile(key,path,bucket=bucket)

# Cell
class ExtraArgs:
  gzip = {'ContentType': 'application/json', 'ContentEncoding':'gzip'}
  publicRead = {'ACL':'public-read'}

# Cell
@add_class_method(S3)
def generalSave(cls, key, objectToSave:dict, bucket = '',
                compressor=lambda x: zlib.compress(x),
                encoder=lambda x: json.dumps(x).encode() ,**kwargs):
  '''save a file to s3'''
  s3 = cls.s3(**kwargs)
  compressedData = compressor(encoder(objectToSave))
  result = s3.put_object(Body=compressedData, Bucket=bucket, Key=key)
  success = result['ResponseMetadata']['HTTPStatusCode'] ==  200
  logging.info('data was saved to s3')
  if not success: raise Error(success)
  else: return True
@add_class_method(S3)
def generalLoad(cls, key, bucket = '',fileName = '/tmp/tempFile.bz',
                decompressor=lambda x: zlib.decompress(x),
                decoder=lambda x: json.loads(x.decode()), useUrl=False, **kwargs):
  '''load file from s3'''
  ### check object exist
#   if not cls.exist(key, bucket, **kwargs):
#     logging.info('object doesnt exist')
#     return {}
#   logging.info('object exists, loading')
  ### download file
  try:
    s3 = cls.s3(**kwargs)
#     s3.download_file(bucket,key, fileName ,useUrl=useUrl)
    cls.loadFile(key=key, path=fileName, bucket=bucket, useUrl=useUrl)
  except Exception as e:
    print(f'downlaod failed {e}')
  ### extract
  with open (fileName , 'rb') as f:
    allItemsByte = f.read()
  if not allItemsByte: raise ValueError('all data does not exist in the database')
  allItems = decoder(decompressor(allItemsByte))
  return allItems

@add_class_method(S3)
def saveZl(cls, key, objectToSave:dict, bucket = '', **kwargs):
  '''save a file to s3'''
  return cls.generalSave(key,objectToSave, bucket )
@add_class_method(S3)
def loadZl(cls, key, bucket = '',fileName = '/tmp/tempFile.bz', **kwargs):
  '''load file from s3'''
  return cls.generalLoad(key,bucket,fileName , **kwargs)

@add_class_method(S3)
def savePklZl(cls, key, objectToSave:dict, bucket = '', **kwargs):
  '''save a file to s3'''
  return cls.generalSave(key,objectToSave, bucket,
                         compressor=lambda x: zlib.compress(x),
                         encoder=lambda x: pickle.dumps(x))


@add_class_method(S3)
def loadPklZl(cls, key, bucket = '',fileName = '/tmp/tempFile.bz', **kwargs):
  '''load file from s3'''
  return cls.generalLoad(key,bucket,fileName,
                         decompressor=lambda x: zlib.decompress(x),
                         decoder = lambda x: pickle.loads(x), **kwargs)

@add_class_method(S3)
def saveRaw(cls, key, objectToSave, bucket = '', **kwargs):
  '''save a file to s3'''
  return cls.generalSave(key,objectToSave, bucket,
                         compressor=lambda x: x,
                         encoder=lambda x: json.dumps(x).encode())
@add_class_method(S3)
def loadRaw(cls, key, bucket = '',fileName = '/tmp/tempFile.bz', **kwargs):
  '''load file from s3'''
  return cls.generalLoad(key,bucket,fileName,
                         decompressor=lambda x: x,
                         decoder = lambda x: json.loads(x.decode()), **kwargs)



# Cell
class Requests:
    '''
      for uploading and downloading contents from url
    '''
    @staticmethod
    def getContentFromUrl( url):
      result = requests.get(url)
      if not result.ok:
        print('error downloading')
        return result.content
      content = result.content
      decompressedContent = bz2.decompress(content)
      contentDict = json.loads(decompressedContent)
      return contentDict

# Cell
from botocore.exceptions import ClientError
@add_class_method(S3)
def createBucket(cls, bucket:str, **kwargs):
  s3 = cls.s3(**kwargs)
  try:
    s3.create_bucket(
      Bucket=bucket,
      CreateBucketConfiguration={'LocationConstraint':'ap-southeast-1'})
  except ClientError as e:
    print(e)
  response = s3.put_bucket_accelerate_configuration(
      Bucket=bucket ,
      AccelerateConfiguration={
          'Status': 'Enabled'
      }
  )
  return response
@add_class_method(S3)
def deleteBucket(cls, bucket:str, **kwargs):
  s3 = cls.s3(**kwargs)
  response = s3.delete_bucket(
    Bucket=bucket)
  return response