# -*- coding: utf-8 -*-
# !/usr/bin/python3
# author: raoyi

import pymongo
import os
import argparse
import sys
import mlog

import os

parser = argparse.ArgumentParser(description='mongodb export import util')
parser.add_argument('-t', "--type", type=int, default=0)
parser.add_argument('-p', "--path", type=str, default='')
parser.add_argument('-d', "--dbname", type=str, default='xxxxxx')
parser.add_argument('-host', "--host", type=str, default='localhost')
parser.add_argument('-port', "--port", type=int, default=27017)
parser.add_argument('-user', '--user', type=str, default='')
parser.add_argument('-pwd', '--pwd', type=str, default='')

mongo_conf = "mongodb://localhost:27017/"
mongo_def_db = "xxxxxx"
mongo_host = "localhost"
mongo_port = 27017
mongo_user = ""
mongo_pwd = ""
mongo_auth = ""

log = mlog.Logger('./tool_log.log', level='debug')


# 1.设置mongodb环境变量使 cmd在任意目录下都可以调用mongodb自带工具 C:\Program Files\MongoDB\Server\4.0\bin

def listdir(_path):
    mongo_json = []
    for file in os.listdir(_path):
        file_path = os.path.join(_path, file)
        if os.path.splitext(file_path)[1] == '.json':
            mongo_json.append(file)
    return mongo_json


def connect_local_mongodb():
    _mongo_client = pymongo.MongoClient(mongo_conf)
    _map4j_db = _mongo_client[mongo_def_db]
    return _map4j_db


def get_mapdb_collects():
    _map4j_db = connect_local_mongodb()
    _collist_names = _map4j_db.list_collection_names()
    return _collist_names


def _do_export(_path):
    _collist_names = get_mapdb_collects()
    for collist_name in _collist_names:
        cmd_str = "mongoexport" \
                  " --host " + mongo_host + \
                  " --port " + str(mongo_port) + \
                  mongo_auth + \
                  " -d " + mongo_def_db + \
                  " -c " + collist_name + \
                  " -o " + _path + "/" + mongo_def_db + "." + collist_name + ".json"
        log.logger.info(cmd_str)
        log.logger.info("export collist_name to " + _path + "/" + mongo_def_db + "." + collist_name + ".json")
        try:
            exec_cmd(cmd_str)
        except Exception as e1:
            log.logger.error("导出操作执行异常：" + str(e1))


def _do_import(_path):
    _collist_names = listdir(_path)
    for file_name in _collist_names:
        # if not file_name.startswith(mongo_def_db):
        #     continue
        collist_name = file_name.replace(mongo_def_db + ".", "").replace(".json", "")
        # 先删除 再 导入
        cmd_str = "mongoimport" \
                  " --host " + mongo_host + \
                  " --port " + str(mongo_port) + \
                  mongo_auth + \
                  " -d " + mongo_def_db + \
                  " -c " + collist_name + \
                  " --drop " + _path + "/" + file_name
        log.logger.info(cmd_str)
        try:
            exec_cmd(cmd_str)
        except Exception as e2:
            log.logger.error("导入操作执行异常：" + str(e2))


def exec_cmd(cmd_str):
    cmd_ret = os.popen(cmd_str)
    str_ret = cmd_ret.read()
    cmd_ret.close()
    log.logger.info(str_ret)


if __name__ == '__main__':
    args = parser.parse_args()
    mongo_user = args.user
    mongo_pwd = args.pwd
    if mongo_user == '' or mongo_pwd == '':
        mongo_auth = ''
    else:
        mongo_auth = ' -u ' + mongo_user + " -p " + mongo_pwd
    path = args.path
    if path == "":
        path = "."
    mongo_host = args.host
    mongo_port = args.port
    mongo_conf = "mongodb://" + mongo_host + ":" + str(mongo_port) + "/"
    mongo_def_db = args.dbname
    try:
        if args.type == 0:
            log.logger.info("执行导出操作开始...")
            _do_export(path)
        elif args.type == 1:
            log.logger.info("执行导入操作开始...")
            _do_import(path)
        else:
            log.logger.warn("不支持的操作类型:", args.type)
        log.logger.info("操作完成")
    except Exception as e:
        log.logger.error("操作执行异常：" + str(e))