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

__all__ = ['NeedUpdateIndex', 'Database', 'DATABASE_TABLE_NAME', 'INVENTORY_BUCKET_NAME', 'INPUT_BUCKET_NAME',
           'ACCESS_KEY_ID', 'SECRET_ACCESS_KEY', 'REGION']

# Cell
from helper import *
import pandas as pd
from datetime import datetime
from pynamodb.models import Model
from pynamodb.attributes import UnicodeAttribute, NumberAttribute, JSONAttribute, BooleanAttribute, BinaryAttribute
from pynamodb.indexes import GlobalSecondaryIndex, AllProjection
from botocore.config import Config
import pickle, json, boto3, bz2, requests, validators

DATABASE_TABLE_NAME = os.environ['DATABASE_TABLE_NAME']
INVENTORY_BUCKET_NAME = os.environ['INVENTORY_BUCKET_NAME']
INPUT_BUCKET_NAME = os.environ['INPUT_BUCKET_NAME']
ACCESS_KEY_ID = None
SECRET_ACCESS_KEY = None
REGION = 'ap-southeast-1'

class NeedUpdateIndex(GlobalSecondaryIndex):
  class Meta:
    index_name = 'need-update'
    projection = AllProjection()
    read_capacity_units = 1
    write_capacity_units = 1
  needUpdate = UnicodeAttribute(default=False, hash_key=True)

class Database(Model):
  class Meta:
    aws_access_key_id = ACCESS_KEY_ID
    aws_secret_access_key = SECRET_ACCESS_KEY
    table_name = DATABASE_TABLE_NAME
    region = REGION
    billing_mode='PAY_PER_REQUEST'
    dax_read_endpoints = [os.environ['DAX_ENDPOINT']]
    dax_write_endpoints = [os.environ['DAX_ENDPOINT']]

  ib_prcode = UnicodeAttribute(hash_key=True, default = '')
  lastUpdate = NumberAttribute( default = 0)
  inventory = JSONAttribute(default={})
  needUpdate = UnicodeAttribute(default= False)
  needUpdateIndex = NeedUpdateIndex()
  TRUE = 'Y'
  FALSE = 'N'

  # def __init__ (self, **kwargs):
    # super().__init__( **kwargs)
  def __repr__(self):
    invDict = self.inventory
    invDict['ib_prcode'] = self.ib_prcode
    return json.dumps(invDict)
  def setNoUpdate(self):
    self.needUpdate = self.FALSE
    return self.save()
  def setUpdate(self):
    self.needUpdate = self.TRUE
    return self.save()

  class S3:
    @staticmethod
    def s3():
      '''
      create and return s3 client
      '''
      config = Config(s3={"use_accelerate_endpoint": True,
                          "addressing_style": "virtual"})
      s3 = boto3.client(
          's3',
          aws_access_key_id= ACCESS_KEY_ID,
          aws_secret_access_key= SECRET_ACCESS_KEY,
          region_name = REGION,
          config = config
        )
      return s3
    @classmethod
    def save(cls,  key, objectToSave, bucket = INVENTORY_BUCKET_NAME):
      '''
      save an object to s3
      '''
      s3 = cls.s3()
      compressedData = bz2.compress(json.dumps(objectToSave).encode())
      result = s3.put_object(Body=compressedData, Bucket=bucket, Key=key)
      success = result['ResponseMetadata']['HTTPStatusCode'] ==  200
      if not success: raise Error(success)
      else: return True
    @classmethod
    def exist(cls, key, bucket):
      return 'Contents' in cls.s3().list_objects(
          Bucket=bucket , Prefix=key )
    @classmethod
    def load(cls, key, bucket=INVENTORY_BUCKET_NAME, fileName = '/tmp/temp.bz'):
      if not cls.exist(key, bucket):
        return {}
      s3 = cls.s3()
      print('downloading object from s3')
      s3.download_file(bucket,key, fileName )
      # requestResult =  s3.get_object(
                    # Bucket = bucket,
                    # Key = key
                  # )
      # allItemsByte = next(requestResult.get('Body',None))
      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 = INVENTORY_BUCKET_NAME):
      if not cls.checkIfExist(key,bucket=bucket): return 'object doesnt exist'
      s3 = cls.s3()
      result = s3.generate_presigned_url(
          'get_object',
            Params={'Bucket': bucket,
                    'Key': key},
          ExpiresIn=expiry)
      return result

    @classmethod
    def checkIfExist(cls, key, bucket = INVENTORY_BUCKET_NAME):
      results = cls.s3().list_objects(Bucket=bucket , Prefix= key)
      return 'Contents' in results

  @classmethod
  def update(cls, input):
    '''
    update products in the database, first use the helper.groupbyprcode to treat the data
    input structure is the foloowing
    - iprcode:
        ibrcode:
          ib_cf_qty: Int
          lastUpdate: Float
          new_ib_bs_stock_cv: Int
    '''
    itemsUpdated = {'success':0, 'failure': 0, 'failureMessage':[]}

    print(f'there are {len(list(input.keys()))} products to update')

    # loop through each product
    for prcode, value in input.items():

      incumbentBr = None
      # check if product is in the database, if not, create an empty class with the product code
      incumbentBr = next(cls.query(prcode),cls(ib_prcode = prcode, inventory = {}))
      print(f'incumbentBr is {incumbentBr}\n, prcode is {prcode}')

      # drop the data with no update
      filteredBr = {
          brcode : v2
          for brcode, v2 in value.items()
          if v2['ib_cf_qty'] != incumbentBr.inventory.get(brcode ,{}).get('ib_cf_qty',None)
      }
      print(f'filteredBr is {filteredBr}')


      # check if there are data to update
      if filteredBr:
        # put the latest update into the lastupdate index
        incumbentBr.lastUpdate = max( v.get('lastUpdate') for v in filteredBr.values() )
        incumbentBr.inventory = (lambda d: d.update(filteredBr) or d)(incumbentBr.inventory)
        incumbentBr.inventory['lastUpdate'] = incumbentBr.lastUpdate
        incumbentBr.inventory['ib_prcode'] = incumbentBr.ib_prcode
        incumbentBr.needUpdate = cls.TRUE

        print(f'saving result {incumbentBr}')
        # try to save the result
        saveResult = incumbentBr.save()
        # saveResult = batch.save(incumbentBr)

        # record the saving success
        if saveResult.get('ConsumedCapacity'):
          itemsUpdated['success'] += 1
        else:
          itemsUpdated['failure'] += 1
          itemsUpdated['failureMessage'].append(saveResult)

    return itemsUpdated

  @classmethod
  def bulkUpdate(cls,input):
    '''
    update products in the database, first use the helper.groupbyprcode to treat the data
    input structure is the foloowing
    - iprcode:
        ibrcode:
          ib_cf_qty: Int
          lastUpdate: Float
          new_ib_bs_stock_cv: Int
    '''
    itemsUpdated = {'success':0, 'failure': 0, 'failureMessage':[]}

    print(f'there are {len(list(input.keys()))} products to update')

    db = cls.loadFromS3()

    with cls.batch_write() as batch:

      # loop through each product
      for prcode, value in input.items():

        incumbentBr = None
        # check if product is in the database, if not, create an empty class with the product code
        incumbentBr = cls(inventory = db.get(prcode) or {}, ib_prcode = prcode)
        # print(incumbentBr)
        # incumbentBr = next(cls.query(prcode),cls(ib_prcode = prcode, inventory = {}))
        # print(f'incumbentBr is {incumbentBr}\n, prcode is {prcode}')

        # drop the data with no update
        filteredBr = {
            brcode : v2
            for brcode, v2 in value.items()
            if v2['ib_cf_qty'] != incumbentBr.inventory.get(brcode ,{}).get('ib_cf_qty',None)
        }
        # print(f'filteredBr is {filteredBr}')


        # check if there are data to update
        if filteredBr:
          # put the latest update into the lastupdate index
          incumbentBr.lastUpdate = max( v.get('lastUpdate') for v in filteredBr.values() )
          incumbentBr.inventory = (lambda d: d.update(filteredBr) or d)(incumbentBr.inventory)
          incumbentBr.inventory['lastUpdate'] = incumbentBr.lastUpdate
          incumbentBr.inventory['ib_prcode'] = incumbentBr.ib_prcode
          incumbentBr.needUpdate = cls.TRUE

          try:
            saveResult = batch.save(incumbentBr)
            itemsUpdated['success'] += 1
          except Error as e:
            print(f'error saving {e}')
            itemsUpdated['failure'] += 1
            itemsUpdated['failureMessage'].append(e)

    return itemsUpdated

  @classmethod
  def updateLambdaInput(cls, input):
    '''
    update products in the database by first grouping the data from lambda
    input
    - ib_prcode: String
      ib_brcode: String
      ib_cf_qty: Int
      new_ib_vs_stock_cv: int
    '''
    groupedInput = Helper.groupByProduct(input)
    return cls.update(groupedInput)

  @classmethod
  def updateS3Input(cls, inputBucketName = INPUT_BUCKET_NAME, key = ''):
    s3Result = cls.loadFromS3(bucketName= inputBucketName, key = key)
    transformedS3Result = [{
        'ib_prcode': item.get('ib_prcode') or item.get('prcode'),
        'ib_brcode': item.get('ib_brcode') or item.get('brcode'),
        'ib_cf_qty': item.get('ib_cf_qty'),
        'new_ib_vs_stock_cv': item.get('new_ib_vs_stock_cv')
        } for item in s3Result]

    print(f' s3 result is {s3Result}')
    groupedInput = Helper.groupByProduct(transformedS3Result)
    updateResult = cls.bulkUpdate(groupedInput)
    return updateResult

  @classmethod
  def singleProductQuery(cls, input):
    if not cls.validateInputQuery(['ib_prcode'] , input): return f"error input {input}"
    return next(cls.query(input.get('ib_prcode')),{})

  @classmethod
  def branchQuery(cls, branchId:str, bucketName = INVENTORY_BUCKET_NAME):
    key = branchId
    result = cls.S3.presign(key, bucket = bucketName)
    return result

  @staticmethod
  def validateInputQuery(keys: list, input:dict):
    '''
      check if input query contains the valid key
      data should have the following structure
      key is a list of keys to check

      ib_prcode: String?
      ib_brcode: String?

      option, one of or both of the ib_procde must be present
    '''
    for key in keys:
      if key not in input.keys():
        raise ValueError(f"key {key} is missing from the input")
      if not input.get(key).isdigit():
        raise ValueError(f'key is not convertable to in {input.get(key)}')
    return True

  @classmethod
  def loadFromS3(cls, bucketName= INVENTORY_BUCKET_NAME, key = 'allData'):
    return cls.S3.load(bucket = bucketName, key=key)

  @classmethod
  def dumpToS3(cls, bucketName= INVENTORY_BUCKET_NAME, key = 'allData'):
    ''' upload changes to s3'''
    # get original data
    allData = cls.S3.load(bucket = bucketName, key = key)
    print(f'all data is {allData}')
    changes = {
                item.ib_prcode:(
                  item.inventory,
                  item.setNoUpdate()
                )[0]
                for item in cls.needUpdateIndex.query(cls.TRUE)
              }

    print(f'all data is {allData}\n changes are {changes}')
    if changes:
      print(f'updating')
      allData.update(changes)
      print(cls.S3.save(key = 'allData',
                  objectToSave = allData,
                  bucket = INVENTORY_BUCKET_NAME)
      )
    else:
      print('no changes to update')

    return allData


  @classmethod
  def splitBranches(cls):
    '''
    split database items from s3 into individual branch
    '''
    data = cls.S3.load('allData')
    # check if object is blank
    if not data: return 'items are empty'
    df = pd.DataFrame(data).drop(['lastUpdate', 'ib_prcode'])
    splitDict = df.T.fillna('').to_dict()
    cleanedSplitDict = {brcode:
      { k:v for k,v in inv.items()  if v}
      for brcode, inv in splitDict.items()}
    cleanedSplitDict

    saveResults = {'success':0, 'failure':0,'errorMessage':[]}
    for k,v in cleanedSplitDict.items():
      # print(k,v)
      saveResult = cls.S3.save(k,v)
      # print(saveResult)
      if saveResult: saveResults['success'] += 1
      else: saveResults['failure'] += 1
    return saveResults