# -*- coding:utf-8 -*-

"""
Author: qiaoxinjiu
Create Data: 2020/11/6 17:34
"""
import pymysql
import pymongo
# from DBUtils.PooledDB import PooledDB
from dbutils.pooled_db import PooledDB
import psycopg2
from psycopg2 import pool
from psycopg2.extras import RealDictCursor
from base_framework.public_tools.read_config import InitConfig
from base_framework.public_tools.read_config import ReadConfig, get_current_config
# from base_framework.public_tools import db_config as config
from base_framework.base_config.current_pth import *

"""
@功能：创建数据库连接池
"""
as_db = ['ZZYY']


class PgConnectionPool(InitConfig):
    pool = None
    pool_cache = dict()

    def __init__(self):
        try:
            super().__init__()
            self.DB_SSL = False
        except Exception as e:
            print(e)

    # 创建数据库连接conn和游标cursor
    def __enter__(self):
        self.conn = self.__getconn()
        self.cursor = self.conn.cursor()

    def __getconn(self, choose_db=None):
        # 如果未指定数据库，使用默认配置
        if not choose_db:
            choose_db = 'default'

        try:
            self.pool = self.pool_cache[choose_db]
        except KeyError:
            rc = ReadConfig(config_file_path)

            # 根据choose_db值获取对应的PostgreSQL配置
            if choose_db == 'default':
                db_host = rc.get_value(sections='PostgreSQL', options='db_test_host')
                db_port = rc.get_value(sections='PostgreSQL', options='db_test_port')
                db_name = rc.get_value(sections='PostgreSQL', options='db_test_dbname')
                db_user = rc.get_value(sections='PostgreSQL', options='db_test_user')
                db_password = rc.get_value(sections='PostgreSQL', options='db_test_password')
                db_min_cached = rc.get_value(sections='PostgreSQL', options='db_min_cached')
                db_max_cached = rc.get_value(sections='PostgreSQL', options='db_max_cached')
                db_max_shared = rc.get_value(sections='PostgreSQL', options='db_max_shared')
                db_max_connecyions = rc.get_value(sections='PostgreSQL', options='db_max_connecyions')
                db_max_usage = rc.get_value(sections='PostgreSQL', options='db_max_usage')
            else:
                # 可以扩展其他数据库连接
                db_host = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_host')
                db_port = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_port')
                db_name = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_name')
                db_user = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_user')
                db_password = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_password')
                db_min_cached = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_min_cached')
                db_max_cached = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_max_cached')
                db_max_shared = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_max_shared')
                db_max_connecyions = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_max_connecyions')
                db_max_usage = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_max_usage')

            # PostgreSQL连接池配置
            try:
                print("=" * 80)
                print("PostgreSQL连接池配置信息:")
                print("  主机(Host): {}".format(db_host))
                print("  端口(Port): {}".format(db_port))
                print("  数据库名(Database): {}".format(db_name))
                print("  用户名(User): {}".format(db_user))
                print("  密码(Password): {} (已隐藏)".format('*' * len(db_password) if db_password else 'None'))
                print("  最小缓存连接数(MinCached): {}".format(db_min_cached))
                print("  最大缓存连接数(MaxCached): {}".format(db_max_cached))
                print("  最大共享连接数(MaxShared): {}".format(db_max_shared))
                print("  最大连接数(MaxConnections): {}".format(db_max_connecyions))
                print("  最大使用次数(MaxUsage): {}".format(db_max_usage))
                print("  SSL模式(SSLMode): {}".format('require' if self.DB_SSL else 'disable'))
                print("  连接超时(ConnectTimeout): 30秒")
                print("=" * 80)
                
                self.pool = PooledDB(
                    creator=psycopg2,
                    host=db_host,
                    port=int(db_port),
                    user=db_user,
                    password=db_password,
                    database=db_name,
                    mincached=int(db_min_cached),
                    maxcached=int(db_max_cached),
                    maxshared=int(db_max_shared),
                    maxconnections=int(db_max_connecyions),
                    blocking=True,
                    maxusage=int(db_max_usage),
                    setsession=None,
                    # PostgreSQL特定参数
                    sslmode='require' if self.DB_SSL else 'disable',
                    connect_timeout=30,
                    keepalives=1,
                    keepalives_idle=30,
                    keepalives_interval=10,
                    keepalives_count=5
                )
                self.pool_cache[choose_db] = self.pool
                print("PostgreSQL连接池创建成功")
            except Exception as e:
                error_msg = """
                    PostgreSQL连接池创建失败！
                    连接配置信息:
                      主机(Host): {}
                      端口(Port): {}
                      数据库名(Database): {}
                      用户名(User): {}
                      密码(Password): {} (已隐藏)
                      SSL模式(SSLMode): {}
                      连接超时(ConnectTimeout): 30秒
                    错误详情: {}
                """.format(
                    db_host, db_port, db_name, db_user,
                    '*' * len(db_password) if db_password else 'None',
                    'require' if self.DB_SSL else 'disable',
                    str(e)
                )
                print(error_msg)
                raise Exception(error_msg)

        try:
            return self.pool.connection()
        except Exception as e:
            # 尝试获取连接配置信息用于错误提示
            try:
                rc = ReadConfig(config_file_path)
                db_host = rc.get_value(sections='PostgreSQL', options='db_test_host')
                db_port = rc.get_value(sections='PostgreSQL', options='db_test_port')
                db_name = rc.get_value(sections='PostgreSQL', options='db_test_dbname')
                db_user = rc.get_value(sections='PostgreSQL', options='db_test_user')
            except:
                db_host = db_port = db_name = db_user = 'unknown'
            error_msg = """
            PostgreSQL连接获取失败！
            连接配置信息:
              主机(Host): {}
              端口(Port): {}
              数据库名(Database): {}
              用户名(User): {}
            错误详情: {}
            """.format(db_host, db_port, db_name, db_user, str(e))
            print(error_msg)
            raise Exception(error_msg)

    # 释放连接池资源
    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.cursor:
            self.cursor.close()
        if self.conn:
            self.conn.close()

    # 获取连接和游标（返回字典形式的结果）
    def getconn(self, choose_db=None):
        conn = self.__getconn(choose_db=choose_db)
        # 使用RealDictCursor返回字典形式的游标
        cursor = conn.cursor(cursor_factory=RealDictCursor)
        return cursor, conn


class MyConnectionPool(InitConfig):
    pool = None
    pool_cache = dict()

    def __init__(self):
        try:
            super().__init__()
        except Exception as e:
            print(e)
        self.current_business = get_current_config(section='run_evn_name', key='current_business')

    # 创建数据库连接conn和游标cursor
    def __enter__(self):
        self.conn = self.__getconn()
        self.cursor = self.conn.cursor()

    def __getconn(self, choose_db=None):
        current_team = ReadConfig(env_choose_path).get_value(sections='run_evn_name', options='current_team')
        if not choose_db:  # 没有指定，则按小组默认设置
            if current_team.upper() in as_db and choose_db is None:
                choose_db = 'as'
            elif current_team.upper() == "SE" and choose_db is None:
                choose_db = 'se'
            elif current_team.upper() == "XUEDAU" and choose_db is None:
                choose_db = 'xdu'
            elif current_team.upper() not in as_db and choose_db is None and self.current_business == 'hh':
                choose_db = 'hh'
            elif current_team.upper() not in as_db and choose_db is None and self.current_business == 'hhi':
                choose_db = 'hhi'
        try:
            self.pool = self.pool_cache[choose_db]
        except Exception as e:
            rc = ReadConfig(config_file_path)
            if choose_db == 'as':
                db_host = rc.get_value(sections='Mysql', options='db_as_svr')
            elif choose_db == 'se':
                db_host = rc.get_value(sections='Mysql', options='db_se_svr')
            elif choose_db == 'xdu':
                db_host = rc.get_value(sections='Mysql', options='db_xdu_svr')
                self.DB_TEST_USER = rc.get_value(sections='Mysql', options='db_xdu_user')
                self.DB_TEST_PASSWORD = rc.get_value(sections='Mysql', options='db_xdu_password')
            elif choose_db == 'hh' or choose_db == 'huohua':
                db_host = rc.get_value(sections='Mysql', options='db_hh_svr')
            elif choose_db == 'hhi':
                db_host = rc.get_value(sections='Mysql', options='db_hhi_svr')
            elif choose_db == 'hh.qa':  # 自动化和信息化的数据都走huohua
                db_host = 'mysql.qa.huohua.cn'
            elif not choose_db:  # 没有传入则默认走huohua
                choose_db = 'hh.qa'
                db_host = 'mysql.qa.huohua.cn'
            else:
                raise Exception("当前仅支持hh,hhi,as,se四个数据库服务器，而你选择的是：{}".format(choose_db))
            default_db = 'sys'
            self.pool = PooledDB(
                creator=pymysql,
                host=db_host,
                port=int(self.DB_TEST_PORT),
                user=self.DB_TEST_USER,
                passwd=self.DB_TEST_PASSWORD,
                db=default_db,
                mincached=int(self.DB_MIN_CACHED),
                maxcached=int(self.DB_MAX_CACHED),
                maxshared=int(self.DB_MAX_SHARED),
                maxconnections=int(self.DB_MAX_CONNECYIONS),
                blocking=True,
                maxusage=int(self.DB_MAX_USAGE),
                setsession=None,
                use_unicode=True,
                charset=self.DB_CHARSET
            )
            self.pool_cache[choose_db] = self.pool
        return self.pool.connection()

    # 释放连接池资源
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.cursor.close()
        self.conn.close()

    # 关闭连接归还给链接池
    # def close(self):
    #     self.cursor.close()
    #     self.conn.close()

    # 从连接池中取出一个连接
    def getconn(self, choose_db=None):
        conn = self.__getconn(choose_db=choose_db)
        # 字典形式返回
        cursor = conn.cursor(pymysql.cursors.DictCursor)
        return cursor, conn


# 获取连接池,实例化
def get_my_connection():
    return MyConnectionPool()


def get_pg_connection():
    return PgConnectionPool()


class MongoConnectionPool(InitConfig):
    def __init__(self):
        try:
            # super().__init__()
            super(MongoConnectionPool, self).__init__()
        except Exception as e:
            print(e)

    def mongo_connect(self):
        try:
            self.connect_ = pymongo.MongoClient(host=self.MONGO_HOST,
                                                port=int(self.MONGO_PORT),
                                                username=self.MONGO_USER,
                                                password=self.MONGO_PASSWORD,
                                                authSource="hulk_teach_marketing"
                                                )
        except Exception as e:
            raise Exception("mongdb连接失败：{}".format(e))
        return self.connect_


def get_my_mongo_connection():
    return MongoConnectionPool()
