import os
from typing import Optional
from loguru import logger


class SessionLogger:
    def __init__(self):
        self.is_main_process = os.environ.get("UVICORN_WID", "0") == "0"

        # 创建日志目录
        os.makedirs("logs", exist_ok=True)

        # 配置loguru
        logger.remove()  # 移除默认handler

        # 控制台输出
        logger.add(
            lambda msg: print(msg, end=""),
            format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> - <level>{level}</level> - [设备: {extra[device_id]}] [会话: {extra[session_id]}] {message}",
            level="INFO"
        )

        # 文件输出 - 持久化保存
        logger.add(
            "logs/talkingq_{time:YYYY-MM-DD}.log",
            format="{time:YYYY-MM-DD HH:mm:ss} - {level} - [设备: {extra[device_id]}] [会话: {extra[session_id]}] {message}",
            rotation="00:00",  # 每天轮转
            retention="30 days",  # 保留30天
            compression="zip",  # 压缩旧日志
            level="INFO"
        )

        self.logger = logger

    def log(
        self,
        level: str,
        device_id: Optional[str],
        session_id: Optional[str],
        message: str,
        **kwargs
    ):
        if isinstance(device_id, str) and device_id.startswith("{"):
            try:
                import json
                device_info = json.loads(device_id)
                device_id = device_info.get("device_id", device_id)
            except:
                pass  # 如果解析失败，保持原样

        # 绑定额外信息到logger
        bound_logger = self.logger.bind(
            device_id=device_id or "unknown",
            session_id=session_id or "unknown"
        )

        # 根据level调用对应方法
        if level.upper() == "INFO":
            bound_logger.info(message)
        elif level.upper() == "ERROR":
            bound_logger.error(message)
        elif level.upper() == "WARNING":
            bound_logger.warning(message)
        else:
            bound_logger.log(level, message)

    def system_log(self, level: str, session_id: str, message: str):
        if self.is_main_process:
            self.log(level, "system", session_id, message)

    def info(self, device_id: Optional[str], session_id: Optional[str], message: str):
        self.log("INFO", device_id, session_id, message)

    def system_info(self, session_id: str, message: str):
        self.system_log("INFO", session_id, message)

    def error(
        self,
        device_id: Optional[str],
        session_id: Optional[str],
        message: str,
        **kwargs
    ):
        self.log("ERROR", device_id, session_id, message, **kwargs)

    def warning(
        self, device_id: Optional[str], session_id: Optional[str], message: str
    ):
        self.log("WARNING", device_id, session_id, message)


session_logger = SessionLogger()