# -*- coding: utf-8 -*-
"""Logging utilities."""
import json
import os
import sys
from typing import Optional, Literal, Any

from loguru import logger

from agentscope.studio._client import _studio_client
from agentscope.web.gradio.utils import (
    generate_image_from_name,
    send_msg,
    get_reset_msg,
    thread_local_data,
)

LOG_LEVEL = Literal[
    "TRACE",
    "DEBUG",
    "INFO",
    "SUCCESS",
    "WARNING",
    "ERROR",
    "CRITICAL",
]

LEVEL_SAVE_LOG = "SAVE_LOG"
LEVEL_SAVE_MSG = "SAVE_MSG"
LEVEL_DISPLAY_MSG = "DISPLAY_MSG"

_SPEAKER_COLORS = [
    ("\033[90m", "\033[0m"),
    ("\033[91m", "\033[0m"),
    ("\033[92m", "\033[0m"),
    ("\033[93m", "\033[0m"),
    ("\033[94m", "\033[0m"),
    ("\033[95m", "\033[0m"),
    ("\033[96m", "\033[0m"),
    ("\033[97m", "\033[0m"),
]

_SPEAKER_TO_COLORS = {}


def _get_speaker_color(speaker: str) -> tuple[str, str]:
    """Get the color markers for a speaker. If the speaker is new, assign a
    new color. Otherwise, return the color that was assigned to the speaker.

    Args:
        speaker (`str`):
            The speaker to be assigned a color.

    Returns:
        `tuple[str, str]`: A color marker tuple, e.g. ("<blue>", "</blue>").
    """
    global _SPEAKER_COLORS, _SPEAKER_TO_COLORS
    if speaker in _SPEAKER_TO_COLORS:
        return _SPEAKER_TO_COLORS[speaker]
    else:
        markers = _SPEAKER_COLORS[
            (len(_SPEAKER_TO_COLORS) + 1)
            % len(
                _SPEAKER_COLORS,
            )
        ]
        # Record the color for this speaker
        _SPEAKER_TO_COLORS[speaker] = markers
        return markers


# add chat function for logger
def _chat(
    message: dict,
    *args: Any,
    disable_studio: bool = False,
    **kwargs: Any,
) -> None:
    """
    Log a chat message with the format of"<speaker>: <content>". If the
    running instance is registered in the studio, the message will be sent
    and display in the studio.

    Args:
        message (`dict`):
            The message to be logged as "<name/role>: <content>", which must
            be an object of Msg class.
    """
    # Push message to studio if it is active
    if _studio_client.active:
        _studio_client.push_message(message)

    # Save message into chat file, add default to ignore not serializable
    # objects
    logger.log(
        LEVEL_SAVE_MSG,
        json.dumps(message, ensure_ascii=False, default=lambda _: None),
        *args,
        **kwargs,
    )

    # Print message in terminal with specific format
    if isinstance(message, dict):
        contain_name_or_role = "name" in message or "role" in message
        contain_content = "content" in message
        contain_url = message.get("url", None) is not None

        # print content if contain name or role and contain content
        if contain_name_or_role:
            speaker = message.get("name", None) or message.get("role", None)
            (m1, m2) = _get_speaker_color(speaker)

            print_str = []
            print_str_without_markers = []
            if contain_content:
                print_str.append(
                    f"{m1}\033[1m{speaker}\033[0m{m2}: {message['content']}",
                )
                print_str_without_markers.append(
                    f"{speaker}: {message['content']}",
                )

            if contain_url:
                print_str.append(
                    f"{m1}\033[1m{speaker}\033[0m{m2}: {message['url']}",
                )
                print_str_without_markers.append(
                    f"{speaker}: {message['url']}",
                )

            if len(print_str) > 0:
                print_str = "\n".join(print_str)
                print_str_without_markers = "\n".join(
                    print_str_without_markers,
                )

                logger.log(LEVEL_DISPLAY_MSG, print_str, *args, **kwargs)
                logger.log(
                    LEVEL_SAVE_LOG,
                    print_str_without_markers,
                    *args,
                    **kwargs,
                )

                if hasattr(thread_local_data, "uid") and not disable_studio:
                    log_studio(message, thread_local_data.uid, **kwargs)
                return

    logger.log(LEVEL_DISPLAY_MSG, message, *args, **kwargs)
    logger.log(LEVEL_SAVE_LOG, message, *args, **kwargs)


def log_studio(message: dict, uid: str, **kwargs: Any) -> None:
    """Send chat message to studio.

    Args:
        message (`dict`):
            The message to be logged. It should have "name"(or "role") and
            "content" keys, and the message will be logged as "<name/role>:
            <content>".
        uid (`str`):
            The local value 'uid' of the thread.
    """
    if uid:
        get_reset_msg(uid=uid)
        name = message.get("name", "default") or message.get("role", "default")
        avatar = kwargs.get("avatar", None) or generate_image_from_name(
            message["name"],
        )

        msg = message["content"]
        flushing = True
        if "url" in message and message["url"]:
            flushing = False
            if isinstance(message["url"], str):
                message["url"] = [message["url"]]
            for i in range(len(message["url"])):
                msg += "\n" + f"""<img src="{message['url'][i]}"/>"""
        if "audio_path" in message and message["audio_path"]:
            flushing = False
            if isinstance(message["audio_path"], str):
                message["audio_path"] = [message["audio_path"]]
            for i in range(len(message["audio_path"])):
                msg += (
                    "\n"
                    + f"""<audio src="{message['audio_path'][i]}"
                controls/></audio>"""
                )
        if "video_path" in message and message["video_path"]:
            flushing = False
            if isinstance(message["video_path"], str):
                message["video_path"] = [message["video_path"]]
            for i in range(len(message["video_path"])):
                msg += (
                    "\n"
                    + f"""<video src="{message['video_path'][i]}"
                controls/></video>"""
                )

        send_msg(
            msg,
            role=name,
            uid=uid,
            flushing=flushing,
            avatar=avatar,
        )


def _level_format(record: dict) -> str:
    """Format the log record."""
    # Display the chat message
    if record["level"].name in [LEVEL_DISPLAY_MSG, LEVEL_SAVE_LOG]:
        return "{message}\n"
    else:
        return (
            "<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | <level>{"
            "level: <8}</level> | <cyan>{name}</cyan>:<cyan>{"
            "function}</cyan>:<cyan>{line}</cyan> - <level>{"
            "message}</level>\n"
        )


def setup_logger(
    path_log: Optional[str] = None,
    level: LOG_LEVEL = "INFO",
) -> None:
    r"""Setup `loguru.logger` and redirect stderr to logging.

    Args:
        path_log (`str`, defaults to `""`):
            The directory of log files.
        level (`str`, defaults to `"INFO"`):
            The logging level, which is one of the following: `"TRACE"`,
            `"DEBUG"`, `"INFO"`, `"SUCCESS"`, `"WARNING"`, `"ERROR"`,
            `"CRITICAL"`.
    """
    # avoid reinit in subprocess
    if not hasattr(logger, "chat"):
        # add chat function for logger
        logger.level(LEVEL_SAVE_LOG, no=51)
        logger.level(LEVEL_DISPLAY_MSG, no=52)

        # save chat message into file
        logger.level(LEVEL_SAVE_MSG, no=53)
        logger.chat = _chat

        # set logging level
        logger.remove()
        # standard output for all logging except chat
        logger.add(
            sys.stdout,
            filter=lambda record: record["level"].name
            not in [LEVEL_SAVE_LOG, LEVEL_SAVE_MSG],
            format=_level_format,
            enqueue=True,
            level=level,
        )

    if path_log is not None:
        os.makedirs(path_log, exist_ok=True)
        path_log_file = os.path.join(path_log, "logging.log")
        path_chat_file = os.path.join(
            path_log,
            "logging.chat",
        )

        # save all logging into file
        logger.add(
            path_log_file,
            filter=lambda record: record["level"].name
            not in [LEVEL_SAVE_MSG, LEVEL_DISPLAY_MSG],
            format=_level_format,
            enqueue=True,
            level=level,
        )

        logger.add(
            path_chat_file,
            format="{message}",
            enqueue=True,
            level=LEVEL_SAVE_MSG,
        )
