import re
import os
import json

import pymongo

from urllib import parse
from bson import ObjectId


class MongoHelper:

    def __init__(self, param):
        """
        param:
            type：str
            desc: 选择连接哪个MongoDB数据库
        """
        with open(os.path.abspath(os.path.dirname(__file__) + '/DBConfig.json')) as f:
            db_configs = json.load(f)
            this_mongo_cfg = db_configs['MongoDB'][param]
            m = re.match(r'([\s\S].*?):([\s\S].*)@([\s\S].*)', this_mongo_cfg)
            parsed_mongo_config = "{}:{}@{}".format(parse.quote_plus(m.group(1)), parse.quote_plus(m.group(2)), m.group(3))
        self.client = pymongo.MongoClient('mongodb://{}'.format(parsed_mongo_config))

    def insert_data(self, dbname: str, sheet: str, data: dict):
        collection = self.client[dbname][sheet]
        item = collection.insert_one(data)
        return item.inserted_id.__str__()

    def delete_data_by_id(self, dbname: str, sheet: str, _id: str):
        collection = self.client[dbname][sheet]
        collection.delete_one({'_id': ObjectId(_id)})
        return True

    def find_data_by_id(self, dbname: str, sheet: str, _id: str):
        collection = self.client[dbname][sheet]
        return collection.find_one({'_id': ObjectId(_id)}, {"_id": False})

    def update_data_by_id(self, dbname: str, sheet: str, _id: str, data: dict):
        collection = self.client[dbname][sheet]
        collection.update_one({'_id': ObjectId(_id)}, {"$set": data})
        return True


def get_mongodb():
    try:
        db = MongoHelper("test")
        yield db
    finally:
        db.client.close()