import re

from .const import COLOR_MAPPING

COLOR_PREFIX = "\033["  # ANSI escape code
COLOR_SUFFIX = "\033[0m"  # ANSI escape code
COLOR_DICT = {
    "system": "red",
    "human": "green",
    "ai": "yellow",
    "##": "blue",
    "###": "magenta",
    "[]": "cyan",
}


def add_color(s: str, color: str) -> str:
    if color in COLOR_MAPPING:
        color = COLOR_MAPPING[color]
    elif color not in COLOR_MAPPING.values():
        raise ValueError(f"Color {color} not supported.")
    return COLOR_PREFIX + color + s + COLOR_SUFFIX


def make_colorful(type: str, content: str, title: str = None) -> str:
    color = COLOR_DICT[type]
    if title is None:
        title = f"# {type} message".upper()
    title = add_color(title, color)
    lines = content.split("\n")
    new_lines = [title]
    for l in lines:
        if l.startswith("###"):
            l = add_color(l, COLOR_DICT["###"])
        elif l.startswith("##"):
            l = add_color(l, COLOR_DICT["##"])
        else:
            regex = r"\[(.*?)\]"
            matches = re.findall(regex, l)
            replaced = set()
            for m in matches:
                if m not in replaced:
                    l = l.replace(f"[{m}]", add_color(f"[{m}]", COLOR_DICT["[]"]))
                    replaced.add(m)
        new_lines.append(l)
    return "\n".join(new_lines)
