import logging
from functools import lru_cache
from os import getenv
from typing import Any, Literal, Optional

from rich.logging import RichHandler
from rich.text import Text

LOGGER_NAME = "agno"
TEAM_LOGGER_NAME = f"{LOGGER_NAME}-team"
WORKFLOW_LOGGER_NAME = f"{LOGGER_NAME}-workflow"

# Define custom styles for different log sources
LOG_STYLES = {
    "agent": {
        "debug": "green",
        "info": "blue",
    },
    "team": {
        "debug": "magenta",
        "info": "steel_blue1",
    },
    "workflow": {
        "debug": "sandy_brown",
        "info": "orange3",
    },
}


class ColoredRichHandler(RichHandler):
    def __init__(self, *args, source_type: Optional[str] = None, **kwargs):
        super().__init__(*args, **kwargs)
        self.source_type = source_type

    def get_level_text(self, record: logging.LogRecord) -> Text:
        # Return empty Text if message is empty
        if not record.msg:
            return Text("")

        level_name = record.levelname.lower()
        if self.source_type and self.source_type in LOG_STYLES:
            if level_name in LOG_STYLES[self.source_type]:
                color = LOG_STYLES[self.source_type][level_name]
                return Text(record.levelname, style=color)
        else:
            if level_name in LOG_STYLES["agent"]:
                color = LOG_STYLES["agent"][level_name]
                return Text(record.levelname, style=color)
        return super().get_level_text(record)


class AgnoLogger(logging.Logger):
    def __init__(self, name: str, level: int = logging.NOTSET):
        super().__init__(name, level)

    def debug(self, msg: str, center: bool = False, symbol: str = "*", *args, **kwargs):  # type: ignore
        if center:
            msg = center_header(str(msg), symbol)
        super().debug(msg, *args, **kwargs)

    def info(self, msg: str, center: bool = False, symbol: str = "*", *args, **kwargs):  # type: ignore
        if center:
            msg = center_header(str(msg), symbol)
        super().info(msg, *args, **kwargs)


def build_logger(logger_name: str, source_type: Optional[str] = None) -> Any:
    # If a logger with the name "agno.{source_type}" is already set, we want to use that one
    _logger = logging.getLogger(f"agno.{logger_name}")
    if _logger.handlers or _logger.level != logging.NOTSET:
        return _logger

    # Set the custom logger class as the default for this logger
    logging.setLoggerClass(AgnoLogger)

    # Create logger with custom class
    _logger = logging.getLogger(logger_name)

    # Reset logger class to default to avoid affecting other loggers
    logging.setLoggerClass(logging.Logger)

    # https://rich.readthedocs.io/en/latest/reference/logging.html#rich.logging.RichHandler
    # https://rich.readthedocs.io/en/latest/logging.html#handle-exceptions
    rich_handler = ColoredRichHandler(
        show_time=False,
        rich_tracebacks=False,
        show_path=True if getenv("AGNO_API_RUNTIME") == "dev" else False,
        tracebacks_show_locals=False,
        source_type=source_type or "agent",
    )
    rich_handler.setFormatter(
        logging.Formatter(
            fmt="%(message)s",
            datefmt="[%X]",
        )
    )

    _logger.addHandler(rich_handler)
    _logger.setLevel(logging.INFO)
    _logger.propagate = False
    return _logger


agent_logger: AgnoLogger = build_logger(LOGGER_NAME, source_type="agent")
team_logger: AgnoLogger = build_logger(TEAM_LOGGER_NAME, source_type="team")
workflow_logger: AgnoLogger = build_logger(WORKFLOW_LOGGER_NAME, source_type="workflow")

# Set the default logger to the agent logger
logger: AgnoLogger = agent_logger

debug_on: bool = False
debug_level: Literal[1, 2] = 1


def set_log_level_to_debug(source_type: Optional[str] = None, level: Literal[1, 2] = 1):
    if source_type is None:
        use_agent_logger()

    _logger = logging.getLogger(LOGGER_NAME if source_type is None else f"{LOGGER_NAME}-{source_type}")
    _logger.setLevel(logging.DEBUG)

    global debug_on
    debug_on = True

    global debug_level
    debug_level = level


def set_log_level_to_info(source_type: Optional[str] = None):
    _logger = logging.getLogger(LOGGER_NAME if source_type is None else f"{LOGGER_NAME}-{source_type}")
    _logger.setLevel(logging.INFO)

    global debug_on
    debug_on = False


def set_log_level_to_warning(source_type: Optional[str] = None):
    _logger = logging.getLogger(LOGGER_NAME if source_type is None else f"{LOGGER_NAME}-{source_type}")
    _logger.setLevel(logging.WARNING)

    global debug_on
    debug_on = False


def set_log_level_to_error(source_type: Optional[str] = None):
    _logger = logging.getLogger(LOGGER_NAME if source_type is None else f"{LOGGER_NAME}-{source_type}")
    _logger.setLevel(logging.ERROR)

    global debug_on
    debug_on = False


def center_header(message: str, symbol: str = "*") -> str:
    try:
        import shutil

        terminal_width = shutil.get_terminal_size().columns
    except Exception:
        terminal_width = 80  # fallback width

    header = f" {message} "
    return f"{header.center(terminal_width - 20, symbol)}"


def use_team_logger():
    """Switch the default logger to use team_logger"""
    global logger
    logger = team_logger


def use_agent_logger():
    """Switch the default logger to use the default agent logger"""
    global logger
    logger = agent_logger


def use_workflow_logger():
    """Switch the default logger to use workflow_logger"""
    global logger
    logger = workflow_logger


@lru_cache(maxsize=128)
def _using_default_logger(logger_instance: Any) -> bool:
    """Return True if the currently active logger is our default AgnoLogger"""
    return isinstance(logger_instance, AgnoLogger)


def log_debug(msg, center: bool = False, symbol: str = "*", log_level: Literal[1, 2] = 1, *args, **kwargs):
    global logger
    global debug_on
    global debug_level

    if debug_on:
        if debug_level >= log_level:
            if _using_default_logger(logger):
                logger.debug(msg, center, symbol, *args, **kwargs)
            else:
                logger.debug(msg, *args, **kwargs)


def log_info(msg, center: bool = False, symbol: str = "*", *args, **kwargs):
    global logger
    if _using_default_logger(logger):
        logger.info(msg, center, symbol, *args, **kwargs)
    else:
        logger.info(msg, *args, **kwargs)


def log_warning(msg, *args, **kwargs):
    global logger
    logger.warning(msg, *args, **kwargs)


def log_error(msg, *args, **kwargs):
    global logger
    logger.error(msg, *args, **kwargs)


def log_exception(msg, *args, **kwargs):
    global logger
    logger.exception(msg, *args, **kwargs)


def configure_agno_logging(
    custom_default_logger: Optional[Any] = None,
    custom_agent_logger: Optional[Any] = None,
    custom_team_logger: Optional[Any] = None,
    custom_workflow_logger: Optional[Any] = None,
) -> None:
    """
    Util to set custom loggers. These will be used everywhere across the Agno library.

    Args:
        custom_default_logger: Default logger to use (overrides agent_logger for default)
        custom_agent_logger: Custom logger for agent operations
        custom_team_logger: Custom logger for team operations
        custom_workflow_logger: Custom logger for workflow operations
    """
    if custom_default_logger is not None:
        global logger
        logger = custom_default_logger

    if custom_agent_logger is not None:
        global agent_logger
        agent_logger = custom_agent_logger

    if custom_team_logger is not None:
        global team_logger
        team_logger = custom_team_logger

    if custom_workflow_logger is not None:
        global workflow_logger
        workflow_logger = custom_workflow_logger
