import logging
from pathlib import Path
from gradio_log import Log

AGENT_LOG_FILES = {
    "ensemble_monitor": "./agent_logs/ensemble_monitor.log",
    "weight_optimization": "./agent_logs/weight_optimization.log",
    "system_health": "./agent_logs/system_health.log",
    "context_intelligence": "./agent_logs/context_intelligence.log",
    "forensic_anomaly_detection": "./agent_logs/forensic_anomaly_detection.log",
}

class CustomFormatter(logging.Formatter):

    green = "\x1b[32;20m"
    blue = "\x1b[34;20m"
    yellow = "\x1b[33;20m"
    red = "\x1b[31;20m"
    bold_red = "\x1b[31;1m"
    reset = "\x1b[0m"
    format = "%(asctime)s - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)"

    FORMATS = {
        logging.DEBUG: blue + format + reset,
        logging.INFO: green + format + reset,
        logging.WARNING: yellow + format + reset,
        logging.ERROR: red + format + reset,
        logging.CRITICAL: bold_red + format + reset,
    }

    def format(self, record):
        log_fmt = self.FORMATS.get(record.levelno)
        formatter = logging.Formatter(log_fmt)
        return formatter.format(record)


class AgentLogger:
    def __init__(self, agent_log_files=None):
        self.agent_log_files = agent_log_files or AGENT_LOG_FILES
        self.loggers = {}
        self._setup_loggers()

    def _setup_loggers(self):
        # Import CustomFormatter from app.py or define it here if needed
        try:
            custom_formatter = CustomFormatter()
        except ImportError:
            # Fallback: use default formatter if CustomFormatter is not available
            custom_formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)")
        for agent, log_file in self.agent_log_files.items():
            Path(log_file).parent.mkdir(parents=True, exist_ok=True)
            Path(log_file).touch(exist_ok=True)
            logger = logging.getLogger(f"agent_log.{agent}")
            logger.setLevel(logging.DEBUG)
            # Remove existing handlers to avoid duplicate logs
            for handler in logger.handlers:
                logger.removeHandler(handler)
            handler = logging.FileHandler(log_file)
            handler.setLevel(logging.DEBUG)
            handler.setFormatter(custom_formatter)
            logger.addHandler(handler)
            self.loggers[agent] = logger

    def log(self, agent, level, message):
        logger = self.loggers.get(agent)
        if logger:
            getattr(logger, level)(message)
        else:
            raise ValueError(f"No logger found for agent: {agent}")

    def get_log_file(self, agent):
        return self.agent_log_files.get(agent)

# Usage Example:
# agent_logger = AgentLogger()
# agent_logger.log("ensemble_monitor", "info", "Ensemble monitoring started.")
# log_file_path = agent_logger.get_log_file("ensemble_monitor")