#! /usr/bin/env python

########################################################################
# Copyright (c) 2020 Robert Bosch GmbH
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0
########################################################################



import threading
import configparser
import os, sys, json, signal
import csv
import time
import queue
from pyarrow import parquet
from pyarrow import fs
import boto3
import pyarrow as pa
import numpy as np 
from datetime import date
from datetime import datetime
import pyarrow_mapping


scriptDir= os.path.dirname(os.path.realpath(__file__))
sys.path.append(os.path.join(scriptDir, "../../"))
from kuksa_viss_client import KuksaClientThread

class S3_Client():

    # Constructor
    def __init__(self, config):
        print("Init S3 client...")
        if "s3" not in config:
            print("s3 section missing from configuration, exiting")
            sys.exit(-1)
        s3_config=config['s3']
        endpoint_url = s3_config.get("endpoint_url", None)
        access_key = s3_config.get("access_key_id", None)
        secret_key = s3_config.get("secret_access_key", None)
        verify = s3_config.getboolean("verify", False)
        region_name = s3_config.get("region_name", None)
        api_version = s3_config.get("api_version", None)
        use_ssl = s3_config.getboolean("use_ssl", True)
        token = s3_config.get("session_token", None)
        self.client = boto3.client('s3',
            endpoint_url = endpoint_url,
            aws_access_key_id = access_key,
            aws_secret_access_key = secret_key,
            verify = verify,
            region_name = region_name,
            api_version = api_version,
            use_ssl = use_ssl,
            aws_session_token = token
            )

        # Create bucket if does not exist
        self.bucket = s3_config.get('bucket', "testbucket")
        bucket_exists = False
        response = self.client.list_buckets()
        print('Existing buckets:')
        for b in response['Buckets']:
            print(f'  {b["Name"]}')
            if self.bucket == b["Name"]:
                bucket_exists = True
        if not bucket_exists:
            print("create bucket " + self.bucket)
            self.client.create_bucket(Bucket= self.bucket)


    def upload(self, srcFile, dstFile):

        # Output the bucket names
        print(f'update file {srcFile} to s3://{self.bucket}/{dstFile}')
        self.client.upload_file(srcFile, self.bucket, dstFile)
        time.sleep(1)

        # read uploaded file for debugging
        try:
            self.client.download_file(self.bucket, dstFile, srcFile + ".downloaded")
        except Exception as e:
            print("The object does not exist." + str(e))



class Kuksa_Client():

    # Constructor
    def __init__(self, config):
        print("Init kuksa client...")
        if "kuksa_val" not in config:
            print("kuksa_val section missing from configuration, exiting")
            sys.exit(-1)
        provider_config=config['kuksa_val']
        self.client = KuksaClientThread(provider_config)
        self.client.start()
        self.client.authorize()
        self.dataset = {}
        
    def getMetaData(self, path):
        metadata = (json.loads(self.client.getMetaData(path))["metadata"])

        return metadata

    def getValue(self, path):
        dataset = json.loads(self.client.getValue(path))["data"]["dp"]["value"]
        # convert to strings to Arrays, which is required by pyarrow.from_pydict
        if not isinstance(dataset, dict):
            dataset= {path: [dataset]}
        else:
            for key in dataset:
                dataset[key] = [dataset[key]]
        return (dataset)


    def shutdown(self):
        self.client.stop()

        
class Parquet_Packer():
    def __init__(self, config):
        print("Init parquet packer...")
        if "parquet" not in config:
            print("parquet section missing from configuration, exiting")
            sys.exit(-1)
        
        self.dataprovider = Kuksa_Client(config)
        self.uploader = S3_Client(config)
        config=config['parquet']
        self.interval = config.get('interval', 1)
        self.path = config.get('path', "")
        self.path = self.path.replace(" ", "").split(",")
        self.max_num_rows = config.getint('max_num_rows', -1)
        
        self.schema = self.genSchema(self.path)
        self.createNewParquet()
        self.running = True

        self.thread = threading.Thread(target=self.loop, args=())
        self.thread.start()

    def createNewParquet(self):
        currTime = datetime.now().strftime("%Y-%b-%d_%H_%M_%S.%f")
        self.pqfile = 'kuksa_' + currTime +'.parquet'
        self.pqwriter = parquet.ParquetWriter(self.pqfile, self.schema)
        self.num_rows = 0


    def genFields(self, metadata, prefix = ""):
        fields = []
        # Note: only first children is considered, which mapped the vss tree structure
        for key in metadata:
            #print("key is " + key)
            if "children" in metadata[key]:
                prefix += key + "."
                return pa.struct(self.genFields(metadata[key]["children"], prefix))
            else:
                #print("prefix is " + prefix)
                datatype = pyarrow_mapping.KUKSA_TO_PYARROW_MAPPING[metadata[key]["datatype"]]()
                #print("datatype is " + str(datatype))
                fields.append(pa.field(prefix+key, datatype))
            

        return fields

    def genSchema(self, path):
        fields = [pa.field("ts", pa.timestamp('us'))]
        for p in path:
            metadata = self.dataprovider.getMetaData(p)
            fields += (self.genFields(metadata))

        return pa.schema(fields)


    def writeTable(self, data):
        table = pa.Table.from_pydict(data)
        table = table.add_column(0, "ts", [[datetime.now()]])
        try:
            table = table.cast(self.schema)
            self.pqwriter.write_table(table)
            self.num_rows += 1
        except pa.lib.ArrowInvalid as e:
            print("ERROR: " + str(e))
            print("The following data can not be written into parquet file:")
            print(data)

        if self.max_num_rows > 0 and self.num_rows >= self.max_num_rows:
            self.pqwriter.close()
            self.uploader.upload(self.pqfile, self.pqfile )
            self.createNewParquet()


    def loop(self):
        print("receive loop started")
        while self.running:

            data = {}
            for p in self.path:
                data.update(self.dataprovider.getValue(p))

            self.writeTable(data)
            
            time.sleep(int(self.interval))

    def shutdown(self):

        self.running=False
        self.dataprovider.shutdown()
        self.pqwriter.close()
        self.thread.join()

        
def main():
    config_candidates=[os.path.join(scriptDir, 'config.ini')]
    for candidate in config_candidates:
        if os.path.isfile(candidate):
            configfile=candidate
            break
    if configfile is None:
        print("No configuration file found. Exiting")
        sys.exit(-1)
    config = configparser.ConfigParser()
    config.read(configfile)
    
    client = Parquet_Packer(config)

    def terminationSignalreceived(signalNumber, frame):
        print("Received termination signal. Shutting down")
        client.shutdown()
    signal.signal(signal.SIGINT, terminationSignalreceived)
    signal.signal(signal.SIGQUIT, terminationSignalreceived)
    signal.signal(signal.SIGTERM, terminationSignalreceived)

if __name__ == "__main__":
    sys.exit(main())
