from typing import Dict, Optional, Callable, Literal


__all__ = ["LoggerMixin"]
TypePrintColor = Literal['black', 'red', 'green', 'yellow', 'blue', 'cyan', 'magenta', 'white', 'reset']


try:
    from colorama import Fore, Style

    print_color_mapping = {
        "BLACK".lower(): Fore.BLACK,
        "RED".lower(): Fore.RED,
        "GREEN".lower(): Fore.GREEN,
        "YELLOW".lower(): Fore.YELLOW,
        "BLUE".lower(): Fore.BLUE,
        "CYAN".lower(): Fore.CYAN,
        "MAGENTA".lower(): Fore.MAGENTA,
        "WHITE".lower(): Fore.WHITE,
        "RESET".lower(): Fore.RESET,
    }


    class LoggerMixin:
        """"""
        logger: Optional[Callable]
        verbose: Optional[bool]

        def _logger(
                self,
                msg: str,
                color: Optional[TypePrintColor] = None
        ) -> None:
            """"""
            if self.logger: self.logger.info(msg)
            if self.verbose:
                if color:
                    print(print_color_mapping.get(color) + msg + Style.RESET_ALL, flush=True)
                else:
                    print(msg, flush=True)

        def _logger_base(self, msg: str, color: Optional[TypePrintColor] = None):
            """"""
            if self.logger: self.logger.info(msg)
            if self.verbose: print(print_color_mapping.get(color) + msg + Style.RESET_ALL, flush=True)

        def _logger_agent_start(
                self,
                msg: Optional[str] = None,
                name: Optional[str] = None,
                color: Optional[TypePrintColor] = "green"
        ):
            """"""
            if msg is None: msg = f"[{name}] Start ...\n"
            self._logger_base(msg=msg, color=color)

        def _logger_agent_end(
                self,
                msg: Optional[str] = None,
                name: Optional[str] = None,
                color: Optional[TypePrintColor] = "green"
        ):
            """"""
            if msg is None: msg = f"[{name}] End ...\n"
            self._logger_base(msg=msg, color=color)

        def _logger_agent_question(
                self,
                msg: Optional[str] = None,
                name: Optional[str] = None,
                content: Optional[str] = None,
                color: Optional[TypePrintColor] = "green"
        ):
            """"""
            if msg is None: msg = f"[{name}] User Question: {content}\n"
            self._logger_base(msg=msg, color=color)

        def _logger_messages_start(
                self,
                msg: Optional[str] = None,
                name: Optional[str] = None,
                color: Optional[TypePrintColor] = "red"
        ):
            """"""
            if msg is None: msg = f"{20 * '='} [{name}] Messages Start {20 * '='}\n"
            self._logger_base(msg=msg, color=color)

        def _logger_messages_end(
                self,
                msg: Optional[str] = None,
                name: Optional[str] = None,
                color: Optional[TypePrintColor] = "red"
        ):
            """"""
            if msg is None: msg = f"{20 * '='} [{name}] Messages End    {20 * '='}\n"
            self._logger_base(msg=msg, color=color)

        def _logger_messages(
                self,
                msg: Optional[str] = None,
                role: Optional[str] = None,
                content: Optional[str] = None,
                color: Optional[TypePrintColor] = "blue"
        ):
            """"""
            if msg is None: msg = f"{role}: [{content}]\n"
            self._logger_base(msg=msg, color=color)

        def _logger_agent_script(
                self,
                msg: Optional[str] = None,
                name: Optional[str] = None,
                script: Optional[str] = None,
                color: Optional[TypePrintColor] = "magenta"
        ):
            """"""
            if msg is None: msg = f"[{name}] Script: ```\n{script}\n```"
            self._logger_base(msg=msg, color=color)

        def _logger_agent_search(
                self,
                msg: Optional[str] = None,
                name: Optional[str] = None,
                content: Optional[str] = None,
                color: Optional[TypePrintColor] = "magenta"
        ):
            """"""
            if msg is None: msg = f"[{name}] Script: ```\n{content}\n```"
            self._logger_base(msg=msg, color=color)

        def _logger_agent_warning(self, msg: str, color: Optional[TypePrintColor] = "red"):
            """"""
            self._logger_base(msg=msg, color=color)

        def _logger_agent_final_answer(
                self,
                msg: Optional[str] = None,
                name: Optional[str] = None,
                content: Optional[str] = None,
                color: Optional[TypePrintColor] = "yellow"
        ):
            """"""
            if msg is None: msg = f"[{name}] Final Answer: {content}\n"
            self._logger_base(msg=msg, color=color)

        def _logger_dict(self, msg: Dict, color="green"):
            """"""
            for key, val in msg.items():
                if self.logger: self.logger.info(f"{key}: {val}")
                if self.verbose: print(print_color_mapping.get(color) + f"{key}: {val}" + Style.RESET_ALL, flush=True)

        def _logger_color(self, msg: str) -> None:
            """"""
            last_line = msg.split('\n')[-1]

            if self.logger: self.logger.info(msg)
            if self.verbose:
                if last_line.startswith("Answer"):
                    print(Fore.YELLOW + msg + Style.RESET_ALL)
                elif last_line.startswith("Action"):
                    print(Fore.GREEN + msg + Style.RESET_ALL)
                elif last_line.startswith("No Action"):
                    print(Fore.RED + msg + Style.RESET_ALL)
                elif last_line.startswith("Observation"):
                    print(Fore.BLUE + msg + Style.RESET_ALL)
                elif last_line.startswith("Running"):
                    print(Fore.CYAN + msg + Style.RESET_ALL)
                else:
                    print(msg)
                print()

except ModuleNotFoundError:
    print(ModuleNotFoundError("pip install colorama"))

    class LoggerMixin:
        """"""
        logger: Optional[Callable]
        verbose: Optional[bool]

        def _logger(
                self,
                msg: str,
                color: Optional[TypePrintColor] = None
        ) -> None:
            """"""
            if self.logger: self.logger.info(msg)
            if self.verbose:
                print(msg, flush=True)
