import json
import logging
from colorama import Fore, Style, init
from datetime import datetime
import numpy as np
import sys
import inspect
import os

# 初始化 colorama（自动重置颜色）
init(autoreset=True)

class CustomLogger:
    def __init__(self, log_name=None, log_file=None, propagate=False, precise_time=True):
        """初始化日志记录器
        
        :param log_name: 日志名称标识
        :param propagate: 是否传递日志到父记录器
        :param precise_time: 是否启用毫秒级时间戳
        """
        self.logger = logging.getLogger(f"custom_logger_{log_name}")
        self.logger.setLevel(logging.INFO)
        self.logger.propagate = propagate
        self.log_name = log_name
        self.precise_time = precise_time
        # 配置日志格式器
        self.log_formatter = logging.Formatter(
            '%(asctime)s - %(log_name)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s'
        )
        
        # 默认添加控制台处理器
        if not self.logger.handlers:
            console_handler = logging.StreamHandler()
            console_handler.setFormatter(self.log_formatter)
            self.logger.addHandler(console_handler)
        # 添加文件处理器
        if log_file:
            # 确保日志目录存在
            os.makedirs(os.path.dirname(log_file), exist_ok=True)
            file_handler = logging.FileHandler(log_file, encoding='utf-8')
            file_handler.setFormatter(self.log_formatter)
            self.logger.addHandler(file_handler)

    def __enter__(self):
        """上下文管理器入口"""
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        """上下文管理器退出时清理资源"""
        for handler in self.logger.handlers[:]:
            self.logger.removeHandler(handler)
            handler.close()

    # 彩色日志方法
    def log_info(self, message, is_print=True):
        self._log(message, logging.INFO, Fore.GREEN, is_print)

    def log_yellow(self, message, is_print=True):
        self._log(message, logging.INFO, Fore.YELLOW, is_print)
    
    def log_blue(self, message, is_print=True):
        self._log(message, logging.INFO, Fore.CYAN, is_print)

    def log_warning(self, message, is_print=True):
        self._log(message, logging.WARNING, Fore.YELLOW, is_print)

    def log_error(self, message, is_print=True):
        self._log(message, logging.ERROR, Fore.RED, is_print)

    def _log(self, message, level, color, is_print=True):
        """内部日志处理核心方法
        
        :param message: 日志内容（支持字符串/字典/数组等）
        :param level: 日志级别（如 logging.INFO）
        :param color: colorama 颜色控制符
        :param is_print: 是否打印到控制台
        """
        # 1. 时间戳格式化
        current_time = datetime.now()
        if self.precise_time:
            formatted_time = current_time.strftime('%Y-%m-%d %H:%M:%S.') + f"{current_time.microsecond // 1000:03d}"
        else:
            formatted_time = current_time.strftime('%Y-%m-%d %H:%M:%S')

        # 2. 获取调用者信息（跳过两层栈帧）
        caller_frame = inspect.stack()[2]
        filename = os.path.basename(caller_frame.filename)
        lineno = caller_frame.lineno

        # 3. 序列化复杂数据类型
        if isinstance(message, (dict, list, tuple, np.ndarray)):
            try:
                message = json.dumps(message if not isinstance(message, np.ndarray) else message.tolist(),
                                   ensure_ascii=False)
            except (TypeError, ValueError):
                message = str(message)
        else:
            message = str(message)

        # 4. 控制台彩色输出
        if is_print:
            level_name = logging.getLevelName(level)
            print(f"{formatted_time} - {self.log_name} - {level_name} - {filename}:{lineno} - {color}{message}{Style.RESET_ALL}")

        # 5. 写入日志系统
        self.logger.log(level, message, extra={'log_name': self.log_name})

if __name__ == "__main__":
    # 示例用法
    with CustomLogger(
        log_name="测试日志", 
        log_file="logs/MassageRobot_nova5.log", 
        precise_time=True
        ) as logger:
        logger.log_info("普通信息（绿色）")
        logger.log_yellow("黄色提示")
        logger.log_blue("蓝色消息")
        logger.log_warning("警告信息（黄色）")
        logger.log_error("错误信息（红色）")
        # 测试复杂数据类型
        # logger.log_info({"key": "值", "数组": [1, 2, 3]})
        # logger.log_info(np.random.rand(3, 3))