import logging
from fastapi import FastAPI
from tortoise.contrib.fastapi import register_tortoise
from src.settings.config import settings

DATABASE_URL = str(settings.DATABASE_URL)
DATABASE_CREATE_URL = str(settings.DATABASE_CREATE_URL)


# 定义不同日志级别的颜色
class ColoredFormatter(logging.Formatter):
    COLORS = {
        "DEBUG": "\033[94m",  # 蓝色
        "INFO": "\033[92m",  # 绿色
        "WARNING": "\033[93m",  # 黄色
        "ERROR": "\033[91m",  # 红色
        "CRITICAL": "\033[41m",  # 红色背景
    }
    RESET = "\033[0m"  # 重置颜色

    def format(self, record):
        color = self.COLORS.get(record.levelname, self.RESET)
        log_message = super().format(record)
        return f"{color}{log_message}{self.RESET}"


# 配置日志
logger_db_client = logging.getLogger("tortoise.db_client")
logger_db_client.setLevel(logging.DEBUG)

# 创建一个控制台处理程序
ch = logging.StreamHandler()
ch.setLevel(logging.DEBUG)

# 创建并设置格式化器
formatter = ColoredFormatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
ch.setFormatter(formatter)

# 将处理程序添加到日志记录器中
logger_db_client.addHandler(ch)

# 定义模型
models = [
    "src.models.test_table",
    "src.models.stock",
    "src.models.security_account",
    "src.models.order",
    "src.models.snowball",
    "src.models.backtest",
    "src.models.strategy",
    "src.models.stock_details",
    "src.models.transaction",
    "src.models.position",
    "src.models.trand_info",
    "src.models.tran_observer_data",
    "src.models.tran_orders",
    "src.models.tran_return",
    "src.models.tran_trade_info",
    "src.models.back_observed_data",
    "src.models.back_observed_data_detail",
    "src.models.back_position",
    "src.models.back_result_indicator",
    "src.models.back_trand_info",
    "src.models.tran_position",
    "src.models.stock_bt_history",
    "src.models.stock_history",
    "src.models.stock_data_processing",
    "src.models.wance_data_stock",
    "src.models.wance_data_storage_backtest",
    "src.models.user_combination_history",
    "src.models.financial_reports",
    "src.models.stock_hu_shen300",
    "src.models.stock_zhong_zheng_500",
    "src.models.stock_guo_zheng_2000",
    "src.models.stock_hu_shen_jing_a",
    "src.models.user_strategy"

    
]

aerich_models = models
aerich_models.append("aerich.models")

# Tortoise ORM 配置
TORTOISE_ORM = {
    "connections": {"default": DATABASE_CREATE_URL,
                    },
    "apps": {
        "models": {
            "models": aerich_models,
            "default_connection": "default",
        },
    },
}

config = {
    'connections': {
        'default': DATABASE_URL
    },
    'apps': {
        'models': {
            'models': models,
            'default_connection': 'default',
        }
    },
    'use_tz': False,
    'timezone': 'Asia/Shanghai'
}


def register_tortoise_orm(app: FastAPI):
    # 注册 Tortoise ORM 配置
    register_tortoise(
        app,
        config=config,
        generate_schemas=False,
        add_exception_handlers=True,
    )


"""
**********
以下内容对原有无改色配置文件进行注释，上述内容为修改后日志为蓝色
**********
"""

# import logging
#
# from fastapi import FastAPI
# from tortoise.contrib.fastapi import register_tortoise
#
# from src.settings.config import settings
#
# DATABASE_URL = str(settings.DATABASE_URL)
# DATABASE_CREATE_URL = str(settings.DATABASE_CREATE_URL)
#
# # will print debug sql
# logging.basicConfig()
# logger_db_client = logging.getLogger("tortoise.db_client")
# logger_db_client.setLevel(logging.DEBUG)
#
# models = [
#     "src.models.test_table",
#     "src.models.stock",
#     "src.models.security_account",
#     "src.models.order",
#     "src.models.snowball",
#     "src.models.backtest",
#     "src.models.strategy",
#     "src.models.stock_details",
#     "src.models.transaction",
# ]
#
# aerich_models = models
# aerich_models.append("aerich.models")
#
# TORTOISE_ORM = {
#
#     "connections": {"default": DATABASE_CREATE_URL},
#     "apps": {
#         "models": {
#             "models": aerich_models,
#             "default_connection": "default",
#         },
#     },
# }
#
# config = {
#     'connections': {
#         # Using a DB_URL string
#         'default': DATABASE_URL
#     },
#     'apps': {
#         'models': {
#             'models': models,
#             'default_connection': 'default',
#         }
#     },
#     'use_tz': False,
#     'timezone': 'Asia/Shanghai'
# }
#
#
# def register_tortoise_orm(app: FastAPI):
#     register_tortoise(
#         app,
#         config=config,
#         # db_url=DATABASE_URL,
#         # db_url="mysql://adams:adams@127.0.0.1:3306/adams",
#         # modules={"models": models},
#         generate_schemas=False,
#         add_exception_handlers=True,
#     )