import logging
from typing import Any, Dict

from agno.utils.log import configure_agno_logging

from LokiLogger import LokiLogger

_AGNO_LOGGING_CONFIGURED = False


class LokiLoggingHandler(logging.Handler):
    """Bridge Python logging records to LokiLogger entries."""

    def __init__(self, loki_logger: LokiLogger) -> None:
        super().__init__()
        self._loki_logger = loki_logger

    def emit(self, record: logging.LogRecord) -> None:
        try:
            message = self.format(record)
        except Exception:
            message = record.getMessage()

        add_fields = _extract_extra_fields(record)
        if add_fields:
            self._loki_logger.log(record.levelname, message, add_fields=add_fields)
        else:
            self._loki_logger.log(record.levelname, message)


def _extract_extra_fields(record: logging.LogRecord) -> Dict[str, Any]:
    """Extract custom fields from the log record while skipping logging internals."""

    skip_keys = {
        "name",
        "msg",
        "args",
        "levelname",
        "levelno",
        "pathname",
        "filename",
        "module",
        "exc_info",
        "exc_text",
        "stack_info",
        "lineno",
        "funcName",
        "created",
        "msecs",
        "relativeCreated",
        "thread",
        "threadName",
        "processName",
        "process",
    }
    extra: Dict[str, Any] = {}
    for key, value in record.__dict__.items():
        if key in skip_keys:
            continue
        extra[key] = value
    return extra


def configure_agno_to_use_loki(loki_logger: LokiLogger) -> None:
    """Configure Agno logging to emit through the provided LokiLogger instance."""

    global _AGNO_LOGGING_CONFIGURED
    if _AGNO_LOGGING_CONFIGURED:
        return

    bridge_logger = logging.getLogger("agno.loki_bridge")
    bridge_logger.setLevel(logging.INFO)
    bridge_logger.propagate = False

    handler = LokiLoggingHandler(loki_logger)
    handler.setFormatter(logging.Formatter("%(name)s - %(levelname)s - %(message)s"))

    # Avoid duplicate handlers if this is invoked multiple times
    if not any(isinstance(existing, LokiLoggingHandler) for existing in bridge_logger.handlers):
        bridge_logger.handlers.clear()
        bridge_logger.addHandler(handler)

    configure_agno_logging(
        custom_default_logger=bridge_logger,
        custom_agent_logger=bridge_logger,
        custom_team_logger=bridge_logger,
    )

    _AGNO_LOGGING_CONFIGURED = True