import datetime
import inspect

class TerminalColors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

class Logger:
    LEVELS = {"TRACE": 0, "INFO": 1, "DEBUG": 2, "WARN": 3, "ERROR": 4}
    COLORS = {"TRACE": TerminalColors.HEADER, "INFO": TerminalColors.OKCYAN, "DEBUG": TerminalColors.OKGREEN,
              "WARN": TerminalColors.WARNING, "ERROR": TerminalColors.FAIL}

    def __init__(self, level="INFO", logfile=None, log_format="{timestamp} {level} {message}", reset_logfile=False):
        self.level = level
        self.logfile = logfile
        self.log_format = log_format

        if reset_logfile and logfile:
            with open(logfile, 'w') as file:
                file.truncate(0)  # This will clear the content of the file

    def _get_caller_info(self):
        frame = inspect.stack()[3]
        filename = frame.filename.split('/')[-1]  # Extracts the last part after the final '/'
        line_number = frame.lineno
        return filename, line_number

    def _log_to_file(self, message):
        if self.logfile:
            with open(self.logfile, 'a') as file:
                file.write(message + '\n')

    def _print_log(self, level_name, *args):
        if Logger.LEVELS[level_name] >= Logger.LEVELS[self.level]:
            timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            message = " ".join(map(str, args))

            if level_name in ["TRACE", "DEBUG"]:
                filename, line_number = self._get_caller_info()
                message = f"({filename}:{line_number}) {message}"

            log_message = self.log_format.format(timestamp=timestamp, level=level_name, message=message)
            print(f"{Logger.COLORS[level_name]}{log_message}{TerminalColors.ENDC}")
            self._log_to_file(log_message)

    def trace(self, *args):
        self._print_log("TRACE", *args)

    def info(self, *args):
        self._print_log("INFO", *args)

    def warn(self, *args):
        self._print_log("WARN", *args)
    
    def debug(self, *args):
        self._print_log("DEBUG", *args)

    def error(self, *args):
        self._print_log("ERROR", *args)