import pymongo
import xarray_mongodb
import bson
import builtins

from ToolFunction.ToolFunction import get_date


npTypeDict = {v: getattr(builtins, k) for k, v in np.sctypeDict.items() if k in vars(builtins)}
npArrayType = type(np.array([0]))


class MongoDB:
    
    def __init__(self, mongoClient, mongoDB, date=None) -> None:
        self.mongoClient = mongoClient
        self.mongoDB = mongoDB
        self.xdb = xarray_mongodb.XarrayMongoDB(mongoDB)
        
        if date is None:
            date= get_date()
        self.set_date(date)
        
    def _convert_numpy_type(self, data):
        for key in data:
            typeKey = type(data[key])
            if typeKey in npTypeDict:
                data[key] = data[key].item()
            elif typeKey == npArrayType:
                data[key] = data[key].tolist()
            else:
                try:
                    data[key] = data[key].item()
                except:
                    pass
        return data
        
    def set_date(self, date):
        date = date.split("/")
        self.year = int(date[0])
        self.month = int(date[1])
        self.day = int(date[2])
        
    def create_global(self, shotNum, dataSet=None, date=None):
        if not date is None:
            self.set_date(date)
        
        data = {
            'year': self.year,
            'month': self.month,
            'day': self.day,
            'shotNum': shotNum,
            'global_parameters' : {},
        }
        
        global_parameters = self._convert_numpy_type(dataSet.attrs)
        
        if not dataSet is None:
            data['global_parameters'].update(global_parameters)
          
        data = self._convert_numpy_type(data)
        
        self.mongoDB['global'].insert_one(data)
    
    def _add_data_normal(self, shotNum, data):
        filter = {   
            'year': self.year,
            'month': self.month,
            'day': self.day,
            'shotNum': shotNum,
        }
        
        self.mongoDB['global'].update_one(filter, {"$set": data}, upsert=False)
        
    def _add_data_xarray_dataArray(self, shotNum, dataArray):
        filter = {   
            'year': self.year,
            'month': self.month,
            'day': self.day,
            'shotNum': shotNum,
        }
        
        dataArray.attrs = self._convert_numpy_type(dataArray.attrs)
        
        mongoID, _ = self.xdb.put(dataArray)
        
        data_label = {
            dataArray.name: 
                {
                    'name': dataArray.name,
                    'mongoID': mongoID,
                    'engine': 'xarray',
                    'dtype': 'dataArray',
                }
        }
        
        self.mongoDB['global'].update_one(filter, {"$set": data_label}, upsert=False)
    
    def _add_data_xarray_dataSet(self, shotNum, dataSet, name):
        filter = {   
            'year': self.year,
            'month': self.month,
            'day': self.day,
            'shotNum': shotNum,
        }
        
        dataSet.attrs = self._convert_numpy_type(dataSet.attrs)
        
        for key in list(dataSet.data_vars):
            dataSet[key].attrs = self._convert_numpy_type(dataSet[key].attrs)
            print(key)
        
        mongoID, _ = self.xdb.put(dataSet)
        
        data_label = {
            name: 
                {
                    'name': name,
                    'mongoID': mongoID,
                    'engine': 'xarray',
                    'dtype': 'dataSet',
                }
        }
        
        self.mongoDB['global'].update_one(filter, {"$set": data_label}, upsert=False)
        
    def add_data(self, shotNum, data, date=None, name=None, engine='normal'):
        pass
    
    def read_data(self, shotNum, data=None):
        pass