#!/usr/bin/env python
"""
Rich Logger - Enhanced logging with Rich formatting and custom levels.
"""

import logging
import os
import re
import inspect
import shutil
from typing import Optional, Union, Iterable, List
from types import ModuleType

try:
    from rich.logging import FormatTimeCallable
except ImportError:
    # Fallback type if Rich is not available
    from typing import Callable
    FormatTimeCallable = Callable[[float], str]

try:
    from rich.logging import RichHandler
    from rich.text import Text
    from rich.console import Console
    from rich.syntax import Syntax
    from rich import traceback as rich_traceback
    console = Console()
    RICH_AVAILABLE = True
except ImportError:
    RICH_AVAILABLE = False
    
class CustomFormatter(logging.Formatter):
    """Custom formatter with ANSI color codes for different log levels."""
    
    # ANSI color codes
    COLORS = {
        'debug': "\x1b[38;2;255;170;0m",         # #FFAA00
        'info': "\x1b[38;2;0;255;255m",          # #00ffff
        'warning': "\x1b[30;48;2;255;255;0m",    # black on #ffff00
        'error': "\x1b[97;41m",                  # #ffffff on red
        'critical': "\x1b[97;48;2;85;0;0m",      # bright_white on #550000
        'fatal': "\x1b[97;48;2;0;85;255m",       # bright_white on #0055FF
        'emergency': "\x1b[97;48;2;170;0;255m",  # bright_white on #aa00ff
        'alert': "\x1b[97;48;2;0;85;0m",         # bright_white on #005500
        'notice': "\x1b[30;48;2;0;255;255m",     # black on #00ffff
        'reset': "\x1b[0m"
    }
    
    FORMAT_TEMPLATE = "%(asctime)s - %(name)s - %(process)d - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)"

    def __init__(
        self,
        show_background=True,
        format_template=None,
        show_time=True,
        show_name=True,
        show_pid=True,
        show_level=True,
        show_path=True,
    ):
        """
        Initialize CustomFormatter with color formatting options.
        
        Args:
            show_background (bool): Whether to show background colors for log levels
            format_template (str, optional): Custom format template string
            show_time (bool): Whether to show timestamp in logs
            show_name (bool): Whether to show logger name in logs
            show_pid (bool): Whether to show process ID in logs
            show_level (bool): Whether to show log level in logs
            show_path (bool): Whether to show file path and line number in logs
        """
        super().__init__()
        # Build format template based on flags if not provided
        if format_template:
            self.FORMAT_TEMPLATE = format_template
        else:
            parts = []
            if show_time:
                parts.append("%(asctime)s")
            if show_name:
                parts.append("%(name)s")
            if show_pid:
                parts.append("%(process)d")
            if show_level:
                parts.append("%(levelname)s")
            parts.append("%(message)s")
            if show_path:
                parts.append("(%(filename)s:%(lineno)d)")
            self.FORMAT_TEMPLATE = " - ".join(parts)

        # Remove background colors for all levels except those that are foreground only
        if not show_background:
            self.COLORS.update({
                'warning': "\x1b[38;2;255;255;0m",    # yellow
                'error': "\x1b[31m",                  # red
                'critical': "\x1b[38;2;85;0;0m",      # #550000
                'fatal': "\x1b[38;2;0;85;255m",       # #0055FF
                'emergency': "\x1b[38;2;170;0;255m",  # #aa00ff
                'alert': "\x1b[38;2;0;85;0m",         # #005500
                'notice': "\x1b[38;2;0;255;255m",     # #00ffff
            })
        
        self.formatters = {
            logging.DEBUG: logging.Formatter(self.COLORS['debug'] + self.FORMAT_TEMPLATE + self.COLORS['reset']),
            logging.INFO: logging.Formatter(self.COLORS['info'] + self.FORMAT_TEMPLATE + self.COLORS['reset']),
            logging.WARNING: logging.Formatter(self.COLORS['warning'] + self.FORMAT_TEMPLATE + self.COLORS['reset']),
            logging.ERROR: logging.Formatter(self.COLORS['error'] + self.FORMAT_TEMPLATE + self.COLORS['reset']),
            logging.CRITICAL: logging.Formatter(self.COLORS['critical'] + self.FORMAT_TEMPLATE + self.COLORS['reset']),
            CRITICAL_LEVEL: logging.Formatter(self.COLORS['critical'] + self.FORMAT_TEMPLATE + self.COLORS['reset']),
            EMERGENCY_LEVEL: logging.Formatter(self.COLORS['emergency'] + self.FORMAT_TEMPLATE + self.COLORS['reset']),
            FATAL_LEVEL: logging.Formatter(self.COLORS['fatal'] + self.FORMAT_TEMPLATE + self.COLORS['reset']),
            ALERT_LEVEL: logging.Formatter(self.COLORS['alert'] + self.FORMAT_TEMPLATE + self.COLORS['reset']),
            NOTICE_LEVEL: logging.Formatter(self.COLORS['notice'] + self.FORMAT_TEMPLATE + self.COLORS['reset']),
        }

    def format(self, record):
        """
        Format a log record with appropriate color coding.
        
        Args:
            record (logging.LogRecord): The log record to format
            
        Returns:
            str: The formatted log message with ANSI color codes
        """
        formatter = self.formatters.get(record.levelno)
        if formatter:
            return formatter.format(record)
        return super().format(record)

# Define custom log levels
EMERGENCY_LEVEL = logging.CRITICAL + 10
FATAL_LEVEL = EMERGENCY_LEVEL + 1
CRITICAL_LEVEL = FATAL_LEVEL + 1
ALERT_LEVEL = CRITICAL_LEVEL + 1
NOTICE_LEVEL = ALERT_LEVEL + 1

DEBUG = logging.DEBUG
ERROR = logging.ERROR
INFO = logging.INFO
WARNING = logging.WARNING

# Add custom level names
logging.addLevelName(EMERGENCY_LEVEL, "EMERGENCY")
logging.addLevelName(FATAL_LEVEL, "FATAL")
logging.addLevelName(CRITICAL_LEVEL, "CRITICAL")
logging.addLevelName(ALERT_LEVEL, "ALERT")
logging.addLevelName(NOTICE_LEVEL, "NOTICE")


def _add_custom_level_method(level_name: str, level_value: int):
    """
    Add a custom logging method to the Logger class.
    
    Args:
        level_name (str): Name of the custom log level
        level_value (int): Numeric value for the log level
    """
    def log_method(self, message, *args, **kwargs):
        if self.isEnabledFor(level_value):
            self._log(level_value, message, args, **kwargs)
    
    setattr(logging.Logger, level_name.lower(), log_method)


# Add custom logging methods
_add_custom_level_method("EMERGENCY", EMERGENCY_LEVEL)
_add_custom_level_method("FATAL", FATAL_LEVEL)
_add_custom_level_method("CRITICAL", CRITICAL_LEVEL)
_add_custom_level_method("ALERT", ALERT_LEVEL)
_add_custom_level_method("NOTICE", NOTICE_LEVEL)

if RICH_AVAILABLE:

    class CustomRichFormatter(logging.Formatter):
        """Custom Rich formatter with syntax highlighting support."""

        LEVEL_STYLES = {
            logging.DEBUG: "bold #FFAA00",
            logging.INFO: "bold #00ffff",
            logging.WARNING: "black on #ffff00",
            logging.ERROR: "#ffffff on red",
            logging.CRITICAL: "bright_white on #550000",
            FATAL_LEVEL: "bright_white on #0055FF",
            EMERGENCY_LEVEL: "bright_white on #aa00ff",
            ALERT_LEVEL: "bright_white on #005500",
            NOTICE_LEVEL: "black on #00ffff",
            CRITICAL_LEVEL: "bright_white on #550000",
        }
        
        def __init__(self, lexer=None, show_background=True):
            """
            Initialize CustomRichFormatter with syntax highlighting options.
            
            Args:
                lexer (str, optional): Syntax highlighter lexer name (e.g., 'python', 'javascript')
                show_background (bool): Whether to show background colors for log levels
            """
            super().__init__()
            # Define styles for different log levels
            self.lexer = lexer
            if not show_background:
                self.LEVEL_STYLES.update({
                    logging.WARNING: "#ffff00",
                    logging.ERROR: "red",
                    logging.CRITICAL: "bold #550000",
                    FATAL_LEVEL: "#0055FF",
                    EMERGENCY_LEVEL: "#aa00ff",
                    ALERT_LEVEL: "#005500",
                    NOTICE_LEVEL: "#00ffff",
                    CRITICAL_LEVEL: "#550000",
                    logging.INFO: "bold #00FF00",
                })
        
        def format(self, record):
            """
            Format log record with Rich styling and optional syntax highlighting.
            
            Args:
                record (logging.LogRecord): The log record to format
                
            Returns:
                rich.text.Text: The formatted log message with Rich styling
            """
            lexer = getattr(record, "lexer", self.lexer or None)
            level_style = self.LEVEL_STYLES.get(record.levelno, "")
            prefix = f"{record.levelname} - ({record.filename}:{record.lineno}) "
            prefix_text = Text(prefix, style=level_style)
            # print(f"lexer: {lexer}, level: {record.levelno}, style: {level_style}")
            if lexer:
                try:
                    # Apply syntax highlighting
                    syntax = Syntax(
                        record.getMessage(), 
                        lexer, 
                        theme="fruity", 
                        line_numbers=False
                    )
                    text_obj = syntax.highlight(record.getMessage())
                    if text_obj.plain.endswith("\n"):
                        text_obj = text_obj[:-1]
                    prefix_text.append(text_obj)
                    return prefix_text
                except Exception:
                    # Fallback if syntax highlighting fails
                    pass
            
            # Handle Text objects
            if isinstance(record.msg, Text):
                return record.msg
                
            # Default formatting
            log_fmt = f"{record.levelname} - {record.getMessage()} ({record.filename}:{record.lineno})"
            return Text(log_fmt, style=level_style)
else:
    # Fallback: CustomRichFormatter = CustomFormatter
    CustomRichFormatter = CustomFormatter

def _check_logging_disabled():
    """Check environment variables to see if logging should be disabled."""
    NO_LOGGING = os.getenv('NO_LOGGING', '0').lower() in ['1', 'true', 'yes']
    LOGGING_DISABLED = os.getenv('LOGGING', '1').lower() in ['0', 'false', 'no']

    if NO_LOGGING or LOGGING_DISABLED:
        # Set a very high level to effectively disable all logging
        root_logger = logging.getLogger()
        root_logger.setLevel(logging.CRITICAL + 99999)
        # Remove all existing handlers to prevent any output
        root_logger.handlers = []
        return True
    return False

def setup_logging_custom(
    level: Union[str, int] = 'DEBUG',
    show_background=True,
    format_template=None,
    show_time=True,
    show_name=True,
    show_pid=True,
    show_level=True,
    show_path=True,
    exceptions=[]
):
    """
    Setup basic logging with custom formatter (ANSI colors).
    
    Args:
        level (Union[str, int]): Logging level (e.g., 'DEBUG', logging.DEBUG)
        show_background (bool): Whether to show background colors for log levels
        format_template (str, optional): Custom format template string
        show_time (bool): Whether to show timestamp in logs
        show_name (bool): Whether to show logger name in logs
        show_pid (bool): Whether to show process ID in logs
        show_level (bool): Whether to show log level in logs
        show_path (bool): Whether to show file path and line number in logs
        exceptions (list): List of logger names to set to CRITICAL level
        
    Returns:
        logging.Logger: Configured logger instance
    """
    if _check_logging_disabled():
        return logging.getLogger()

    if isinstance(level, str):
        logging.basicConfig(level=getattr(logging, level))
    else:
        logging.basicConfig(level=level)

    if exceptions:
        for i in exceptions:
            if isinstance(i, str): 
                logging.getLogger(str(i)).setLevel('CRITICAL')

    logger = logging.getLogger()

    # Update handlers with custom formatter
    for handler in logger.handlers:
        handler.setFormatter(CustomFormatter(
                show_background,
                format_template,
                show_time,
                show_name,
                show_pid,
                show_level,
                show_path,
            )
        )
    
    return logger

def setup_logging(
    name: Optional[str] = None,
    lexer: Optional[str] = None,
    logtofile: bool = False,
    logfile: Optional[str] = None, 
    show_locals: bool = False, 
    level: Union[str, int] = 'DEBUG',
    show_level: bool = False,
    show_time: bool = True,
    omit_repeated_times: bool = True,
    show_path: bool = True,
    enable_link_path: bool = True,
    highlighter=None,
    markup: bool = False,
    rich_tracebacks: bool = False,
    tracebacks_width: Optional[int] = None,
    tracebacks_extra_lines: int = 3,
    tracebacks_theme: Optional[str] = None,
    tracebacks_word_wrap: bool = True,
    tracebacks_show_locals: bool = False,
    tracebacks_suppress: Iterable[Union[str, ModuleType]] = (),
    locals_max_length: int = 10,
    locals_max_string: int = 80,
    log_time_format: Union[str, FormatTimeCallable] = "[%x %X]",
    keywords: Optional[List[str]] = None,
    show_background=True,
    exceptions=[]
) -> logging.Logger:
    """
    Setup enhanced logging with Rich formatting.
    
    Args:
        lexer (str, optional): Syntax highlighter for code (e.g., 'python', 'javascript')
        logfile (str, optional): Path to log file (auto-generated if None)
        show_locals (bool): Show local variables in tracebacks
        level (Union[str, int]): Logging level (e.g., 'DEBUG', logging.DEBUG)
        show_level (bool): Whether to show log level in console output
        show_time (bool): Whether to show timestamp in console output
        omit_repeated_times (bool): Whether to omit repeated timestamps
        show_path (bool): Whether to show file path in console output
        enable_link_path (bool): Whether to enable clickable file paths
        highlighter: Rich highlighter instance
        markup (bool): Whether to enable Rich markup in log messages
        rich_tracebacks (bool): Whether to use Rich formatted tracebacks
        tracebacks_width (int, optional): Width of traceback display
        tracebacks_extra_lines (int): Extra lines to show in tracebacks
        tracebacks_theme (str, optional): Theme for traceback syntax highlighting
        tracebacks_word_wrap (bool): Whether to wrap long lines in tracebacks
        tracebacks_show_locals (bool): Whether to show local variables in tracebacks, same as `show_locals`
        tracebacks_suppress (Iterable): Modules to suppress in tracebacks
        locals_max_length (int): Maximum length of local variable representations
        locals_max_string (int): Maximum length of string representations
        log_time_format (Union[str, FormatTimeCallable]): Time format for logs
        keywords (List[str], optional): Keywords to highlight in logs
        show_background (bool): Whether to show background colors for log levels
        exceptions (list): List of logger names to set to CRITICAL level
        
    Returns:
        logging.Logger: Configured logger instance
    """
    if _check_logging_disabled():
        return logging.getLogger()
    
    # print(f"Setting up logging with level: {level}, show_locals: {show_locals}, logfile: {logfile}, lexer: {lexer}")
    # Auto-generate logfile name if not provided
    # if logfile is None:
    #     try:
    #         main_file = inspect.stack()[-1].filename
    #         base = os.path.splitext(os.path.basename(main_file))[0]
    #         logfile = f"{base}.log"
    #     except (IndexError, AttributeError):
    #         logfile = "app.log"

    if exceptions:
        for i in exceptions:
            if isinstance(i, str): 
                logging.getLogger(str(i)).setLevel('CRITICAL')
    
    if isinstance(level, str):
        logging.basicConfig(level=getattr(logging, level))
    else:
        logging.basicConfig(level=level)
    # console.log(f"logtofile = {logtofile}, logfile = {logfile}")
    if logfile is None:
        try:
            main_file = inspect.stack()[-1].filename
            # Check if it contains an invalid character or not a python file
            if not main_file or main_file.startswith('<') or not main_file.endswith(('.py', '.pyc')):
                logfile = "app.log"
            else:
                base = os.path.splitext(os.path.basename(main_file))[0]
                logfile = f"{base}.log"
        except Exception as e:
            logfile = "app.log"
    # console.log(f"logtofile = {logtofile}, logfile = {logfile}")
    # Setup Rich handler for console
    rich_handler = RichHandler(
        show_time=show_time,
        omit_repeated_times=omit_repeated_times,
        show_level=show_level,
        show_path=show_path,
        enable_link_path=enable_link_path,
        highlighter=highlighter,
        markup=markup,
        rich_tracebacks=rich_tracebacks,
        tracebacks_width=tracebacks_width or shutil.get_terminal_size()[0],
        tracebacks_extra_lines=tracebacks_extra_lines,
        tracebacks_theme=tracebacks_theme or 'fruity',
        tracebacks_word_wrap=tracebacks_word_wrap,
        tracebacks_show_locals=show_locals or tracebacks_show_locals,
        tracebacks_suppress=tracebacks_suppress,
        locals_max_length=locals_max_length,
        locals_max_string=locals_max_string,
        log_time_format=log_time_format,
        keywords=keywords,
    )
    rich_handler._log_render.emojis = False
    rich_handler.setFormatter(CustomRichFormatter(lexer, show_background))

    if logtofile:
        # Setup file handler
        file_handler = logging.FileHandler(logfile, encoding="utf-8")
        file_handler.setFormatter(logging.Formatter(
            "%(asctime)s - %(name)s - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)"
        ))

    # Configure root logger
    logger = logging.getLogger(name)
    logger.setLevel(level)
    logger.handlers.clear()
    logger.addHandler(rich_handler)
    if logtofile: logger.addHandler(file_handler)
    
    # Patch RichHandler to handle Text objects properly
    def custom_render_message(self, record, message):
        use_markup = getattr(record, "markup", self.markup)
        if isinstance(message, Text):
            return message
        else:
            return Text.from_markup(message) if use_markup else Text(message)
    
    # Apply patch to RichHandler
    for handler in logger.handlers:
        if isinstance(handler, RichHandler):
            handler.render_message = custom_render_message.__get__(handler)
    
    return logger


def get_def() -> str:
    """
    Get current function/class definition name for logging context.
    
    Returns:
        str: Formatted string with function/class context information
    """
    name = ''
    
    try:
        # Try to get function name from stack
        frame = inspect.stack()[1]
        name = str(frame.function)
    except (IndexError, AttributeError) as e:
        logging.exception("Error getting name from stack[1]: %s", e)
    
    # Fallback to stack[2] if needed
    if not name:
        try:
            frame = inspect.stack()[2]
            name = str(frame.function)
        except (IndexError, AttributeError) as e:
            logging.exception("Error getting name from stack[2]: %s", e)
    
    # Handle module-level calls
    if not name or name == "<module>":
        # Try to get class name
        try:
            frame = inspect.stack()[1]
            self_obj = frame.frame.f_locals.get('self')
            if self_obj:
                class_name = self_obj.__class__.__name__
                if class_name != "NoneType":
                    name = f"[#00ffff]({class_name}) --> "
        except Exception as e:
            logging.exception("Error getting class from stack[1]: %s", e)
        
        # Look for calling function in stack
        if not name or name == "<module>":
            try:
                for frame_info in inspect.stack()[3:]:
                    if isinstance(frame_info.lineno, int) and frame_info.function != '<module>':
                        name = f"[#ff5500]{frame_info.function}\\[[white on red]{frame_info.lineno}][/] --> "
                        break
            except Exception as e:
                logging.exception("Error scanning stack: %s", e)
    
    # Ultimate fallback to filename
    if not name or name == "<module>":
        try:
            filename = os.path.basename(inspect.stack()[0].filename)
            name = filename
        except Exception as e:
            logging.exception("Error getting filename: %s", e)
            name = "unknown"
    
    return name or "unknown"

def test():
    """Test function to verify logger setup with different configurations."""
    logger = setup_logging_custom()
    try:
        console.print("[italic]Test function to verify logger setup (CustomFormatter).[/]\n")
    except:
        print("Test function to verify logger setup (CustomFormatter).\n")
    logger.emergency("This is an emergency message")
    logger.critical("This is a critical message")
    logger.alert("This is an alert message")
    logger.fatal("This is a fatal message")
    logger.error("This is an error message")
    logger.warning("This is a warning message")
    logger.notice("This is a notice message")
    logger.debug("This is a debug message")
    logger.info("This is an info message")
    print("=" * shutil.get_terminal_size()[0])
    
    logger = setup_logging_custom(show_background=False)
    try:
        console.print("[italic]Test function to verify logger setup (CustomFormatter), No Background Color.[/]\n")
    except:
        print("Test function to verify logger setup (CustomFormatter), No Background Color.\n")
    logger.emergency("This is an emergency message")
    logger.critical("This is a critical message")
    logger.alert("This is an alert message")
    logger.fatal("This is a fatal message")
    logger.error("This is an error message")
    logger.warning("This is a warning message")
    logger.notice("This is a notice message")
    logger.debug("This is a debug message")
    logger.info("This is an info message")
    print("=" * shutil.get_terminal_size()[0])
    
    logger = setup_logging(logtofile=True)
    
    try:
        console.print("[italic]Test function to verify logger setup (CustomRichFormatter).[/]\n")
    except:
        print("Test function to verify logger setup (CustomRichFormatter).\n")
    logger.emergency("This is an emergency message")
    logger.critical("This is a critical message")
    logger.alert("This is an alert message")
    logger.fatal("This is a fatal message")
    logger.error("This is an error message")
    logger.warning("This is a warning message")
    logger.notice("This is a notice message")
    logger.debug("This is a debug message")
    logger.info("This is an info message")
    print("=" * shutil.get_terminal_size()[0])
    
    logger = setup_logging(show_background=False)
    
    try:
        console.print("[italic]Test function to verify logger setup (CustomRichFormatter), No Background Color.[/]\n")
    except:
        print("Test function to verify logger setup (CustomRichFormatter), No Background Color.\n")
    logger.emergency("This is an emergency message")
    logger.critical("This is a critical message")
    logger.alert("This is an alert message")
    logger.fatal("This is a fatal message")
    logger.error("This is an error message")
    logger.warning("This is a warning message")
    logger.notice("This is a notice message")
    logger.debug("This is a debug message")
    logger.info("This is an info message")
    
def run_test():
    # Run test if executed directly
    test()
    # Example usage of get_def
    print(f"get_def(): {get_def()}")
    try:
        from . example_usage import main, ExampleClass
    except Exception as e:
        from example_usage import main, ExampleClass

    main()
    
    # Test class-based logging
    obj = ExampleClass()
    result = obj.example_method()
    print(f"Result: {result}")

if __name__ == "__main__":
    run_test()