import collections
import functools
import html
import itertools
import os
import weakref

import cwcwidth


__version__ = "v2022.05.26"


ESC = "\x1b"

NORMAL = "[m"
BOLD = "[1m"
ITALIC = "[3m"
UNDERLINE = "[4m"


def color(color_number, is_foreground):
    return f"[{'38' if is_foreground else '48'};5;{color_number:d}m"


def rgb_color(rgb, is_foreground):
    return f"[{'38' if is_foreground else '48'};2;" + "%i;%i;%im" % rgb


def blend_color(a_color, b_color, transparency):
    a_r, a_g, a_b = a_color
    b_r, b_g, b_b = b_color
    complement = 1 - transparency
    return (int(a_r * transparency + b_r * complement),
            int(a_g * transparency + b_g * complement),
            int(a_b * transparency + b_b * complement))


class Color:

    # https://en.wikipedia.org/wiki/Natural_Color_System
    black = (0, 0, 0)
    white = (255, 255, 255)
    red = (196, 2, 51)
    green = (0, 159, 107)
    dark_green = (0, 119, 80)
    blue = (0, 135, 189)
    lime = (0, 255, 0)
    yellow = (255, 211, 0)
    grey_30 = (30, 30, 30)
    grey_40 = (40, 40, 40)
    grey_50 = (50, 50, 50)
    grey_80 = (80, 80, 80)
    grey_100 = (100, 100, 100)
    grey_150 = (150, 150, 150)
    grey_200 = (200, 200, 200)
    light_blue = (173, 216, 230)
    purple = (200, 0, 200)
    brown = (150, 75, 0)
    orange = (255, 153, 0)


def _xterm_colors():
    result = [(0x00, 0x00, 0x00), (0xcd, 0x00, 0x00), (0x00, 0xcd, 0x00), (0xcd, 0xcd, 0x00),
              (0x00, 0x00, 0xee), (0xcd, 0x00, 0xcd), (0x00, 0xcd, 0xcd), (0xe5, 0xe5, 0xe5),
              (0x7f, 0x7f, 0x7f), (0xff, 0x00, 0x00), (0x00, 0xff, 0x00), (0xff, 0xff, 0x00),
              (0x5c, 0x5c, 0xff), (0xff, 0x00, 0xff), (0x00, 0xff, 0xff), (0xff, 0xff, 0xff)]
    grad = [0x00, 0x5f, 0x87, 0xaf, 0xd7, 0xff]
    result.extend([(grad[(i // 36) % 6], grad[(i // 6) % 6], grad[i % 6]) for i in range(216)])
    result.extend([(8 + i * 10, 8 + i * 10, 8 + i * 10) for i in range(24)])
    return result


XTERM_COLORS = _xterm_colors()


def closest_color_index(color, colors):
    r, g, b = color
    closest_distance = 3 * (256*256) + 1
    for index, (r_, g_, b_) in enumerate(colors):
        distance = (r_ - r) ** 2 + (g_ - g) ** 2 + (b_ - b) ** 2
        if distance < closest_distance:
            closest_distance = distance
            color_index = index
    return color_index


class CharStyle:

    _POOL = weakref.WeakValueDictionary()

    def __new__(cls, fg_color=None, bg_color=None, is_bold=False,
                is_italic=False, is_underlined=False):
        if fg_color is None:
            fg_color = Color.white
        if bg_color is None:
            bg_color = Color.black
        key = (fg_color, bg_color, is_bold, is_italic, is_underlined)
        try:
            return CharStyle._POOL[key]
        except KeyError:
            obj = object.__new__(cls)
            obj.fg_color, obj.bg_color, obj.is_bold, obj.is_italic, obj.is_underlined = key
            return CharStyle._POOL.setdefault(key, obj)

    def __getnewargs__(self):
        return self.fg_color, self.bg_color, self.is_bold, self.is_italic, self.is_underlined

    def __getstate__(self):
        state = self.__dict__.copy()
        if "_cache" in state:
            del state["_cache"]
        return state

    def __setstate__(self, state):
        self.__dict__ = state

    def __repr__(self):
        attributes = []
        if self.is_bold:
            attributes.append("b")
        if self.is_italic:
            attributes.append("i")
        if self.is_underlined:
            attributes.append("u")
        return f"<CharStyle: fg:{self.fg_color} bg:{self.bg_color} attr:{','.join(attributes)}>"

    @staticmethod
    def _color_code(color_, is_foreground):
        if isinstance(color_, int):
            return color(color_, is_foreground)
        else:  # true color
            if os.environ.get("TERM", None) == "xterm":
                color_index = closest_color_index(color_, XTERM_COLORS)
                return color(color_index, is_foreground)
            else:
                return rgb_color(color_, is_foreground)

    @functools.cached_property
    def code_for_term(self):
        fg_termcode = ESC + self._color_code(self.fg_color, True)
        bg_termcode = ESC + self._color_code(self.bg_color, False)
        bold_code = (ESC + BOLD) if self.is_bold else ""
        italic_code = ((ESC + ITALIC) if self.is_italic else "")
        underline_code = ((ESC + UNDERLINE) if self.is_underlined else "")
        return "".join([ESC, NORMAL, fg_termcode, bg_termcode, bold_code,
                        italic_code, underline_code])

    @functools.cached_property
    def fg_rgb_color(self):
        return self.fg_color if type(self.fg_color) == tuple else XTERM_COLORS[self.fg_color]

    @functools.cached_property
    def bg_rgb_color(self):
        return self.bg_color if type(self.bg_color) == tuple else XTERM_COLORS[self.bg_color]

    def as_html(self):
        bold_code = "font-weight:bold; " if self.is_bold else ""
        italic_code = "font-style:italic; " if self.is_italic else ""
        underline_code = "text-decoration:underline; " if self.is_underlined else ""
        return (f"<style>.S{id(self)} {{font-size:80%%; color:rgb{self.fg_rgb_color!r};"
                f" background-color:rgb{self.bg_rgb_color!r}; "
                f"{bold_code}{italic_code}{underline_code}}}</style>")


def _join_lists(lists):
    return list(itertools.chain.from_iterable(lists))


_ZERO_WIDTH_SPACE = "\u200b"


def _pad_wide_chars(str_):
    parts = []
    last_width = None
    for char in str_:
        width = cwcwidth.wcwidth(char)
        parts.append(f"{_ZERO_WIDTH_SPACE}{char}" if width != 0 and last_width == 2 else char)
        last_width = width
    padded_str = "".join(parts)
    if len(padded_str) > 0 and cwcwidth.wcwidth(padded_str[-1]) == 2:
        padded_str += _ZERO_WIDTH_SPACE
    return str_ if len(padded_str) == len(str_) else padded_str


def join(seperator, parts):
    if parts == []:
        return ""
    try:
        return seperator.join(parts)
    except TypeError:
        return TermStr(seperator).join(parts)


class TermStr(collections.UserString):

    def __init__(self, data, style=CharStyle()):
        if isinstance(style, tuple):
            self.data = data
            self.style = style
        else:
            try:
                self.data, self.style = data.data, data.style
            except AttributeError:
                self.data = _pad_wide_chars(data).expandtabs()
                self.style = (style,) * len(self.data)

    @classmethod
    def from_term(cls, data):
        data = data.expandtabs(tabsize=4)
        parts = data.split(ESC)
        fg_color, bg_color = None, None
        is_bold, is_italic, is_underlined = False, False, False
        result_parts = [parts[0]]
        for part in parts[1:]:
            if part.startswith("[K"):
                end_index = part.index("K")
                codes = []
            else:
                try:
                    end_index = part.index("m")
                except ValueError:
                    continue
                codes = part[1:end_index].split(";")
            previous_code = None
            for index, code in enumerate(codes):
                try:
                    code_int = int(code)
                except ValueError:
                    code_int = None
                if code in ["", "0", "00"]:  # Normal
                    is_bold, is_italic, is_underlined = False, False, False
                    fg_color, bg_color = None, None
                elif code in ["01", "1"]:  # bold
                    is_bold = True
                elif code in ["03", "3"]:  # italic
                    is_italic = True
                elif code in ["04", "4"]:  # underline
                    is_underlined = True
                elif code_int and 30 <= code_int <= 37:  # dim fg color
                    fg_color = int(code[1])
                elif code_int and 40 <= code_int <= 47:  # dim bg color
                    bg_color = int(code[1])
                elif code_int and 90 <= code_int <= 97:  # high fg color
                    fg_color = int(code[1]) + 8
                elif code_int and 100 <= code_int <= 107:  # high bg color
                    bg_color = int(code[2]) + 8
                elif code == "5" and previous_code == "38":  # simple fg color
                    fg_color = int(codes[index+1])
                    codes[index+1:index+2] = []
                elif code == "5" and previous_code == "48":  # simple bg color
                    bg_color = int(codes[index+1])
                    codes[index+1:index+2] = []
                elif code == "2" and previous_code == "38":  # rgb fg color
                    fg_color = tuple(int(component) for component in codes[index+1:index+4])
                    codes[index+1:index+4] = []
                elif code == "2" and previous_code == "48":  # rgb bg color
                    bg_color = tuple(int(component) for component in codes[index+1:index+4])
                    codes[index+1:index+4] = []
                previous_code = code
            result_parts.append(cls(part[end_index+1:], CharStyle(fg_color, bg_color, is_bold,
                                                                  is_italic, is_underlined)))
        return cls("").join(result_parts)

    def __eq__(self, other):
        return (self is other or (isinstance(other, self.__class__) and
                                  self.data == other.data and self.style == other.style))

    def __ne__(self, other):
        return not self.__eq__(other)

    def __hash__(self):
        return hash((self.data, self.style))

    @functools.cached_property
    def _partition_style(self):
        if self.data == "":
            return []
        last_style, last_index = None, 0
        result = []
        for index, style in enumerate(self.style):
            if style != last_style:
                if last_style is not None:
                    result.append((last_style, last_index, index))
                last_style, last_index = style, index
        result.append((last_style, last_index, len(self.style)))
        return result

    def __str__(self):
        return "".join(_join_lists([style.code_for_term, self.data[start_index:end_index]]
            for style, start_index, end_index in self._partition_style) + [ESC + NORMAL])

    def __repr__(self):
        return f"<TermStr: {self.data!r}>"

    def __add__(self, other):
        if isinstance(other, str):
            other = TermStr(other)
        return self.__class__(self.data + other.data, self.style + other.style)

    def __radd__(self, other):
        if isinstance(other, str):
            other = TermStr(other)
        return self.__class__(other.data + self.data, other.style + self.style)

    def __mul__(self, n):
        return self.__class__(self.data*n, self.style*n)
    __rmul__ = __mul__

    def __getitem__(self, index):
        data = self.data[index]
        if len(data) == 0:
            result = ""
        else:
            first_char = " " if cwcwidth.wcwidth(data[0]) == 0 else data[0]
            if len(data) == 1:
                result = first_char
            else:
                end_char = " " if cwcwidth.wcwidth(data[-1]) == 2 else data[-1]
                result = first_char + data[1:-1] + end_char
        return self.__class__(result, self.style[index])

    def join(self, parts):
        parts = [TermStr(part) if isinstance(part, str) else part for part in parts]
        joined_style = _join_lists(self.style + part.style for part in parts)
        return self.__class__(self.data.join(part.data for part in parts),
                              tuple(joined_style[len(self.style):]))

    def _split_style(self, parts, sep_length):
        result = []
        cursor = 0
        for part in parts:
            style_part = self.style[cursor:cursor+len(part)]
            result.append(self.__class__(part, style_part))
            cursor += (len(part) + sep_length)
        return result

    def split(self, sep=None, maxsplit=-1):
        return self._split_style(self.data.split(sep, maxsplit), len(sep))

    def splitlines(self, keepends=0):
        result = []
        cursor = 0
        for line in self.data.splitlines(keepends=True):
            result_line = line if keepends else line.rstrip("\r\n")
            style_part = self.style[cursor:cursor+len(result_line)]
            result.append(self.__class__(result_line, style_part))
            cursor += len(line)
        return result

    def capitalize(self):
        return self.__class__(self.data.capitalize(), self.style)

    def lower(self):
        return self.__class__(self.data.lower(), self.style)

    def swapcase(self):
        return self.__class__(self.data.swapcase(), self.style)

    def title(self):
        return self.__class__(self.data.title(), self.style)

    def upper(self):
        return self.__class__(self.data.upper(), self.style)

    def ljust(self, width, fillchar=" "):
        return self + self.__class__(fillchar * (width - len(self.data)))

    def rjust(self, width, fillchar=" "):
        return self.__class__(fillchar * (width - len(self.data))) + self

    def center(self, width, fillchar=" "):
        left_width = round((width - len(self.data)) / 2)
        return (self.__class__(fillchar * left_width) + self +
                self.__class__(fillchar * (width - left_width - len(self.data))))

    # Below are extra methods useful for termstrs.

    def transform_style(self, transform_func):
        new_style = tuple(_join_lists([transform_func(style)] * (end_index - start_index)
                                      for style, start_index, end_index in self._partition_style))
        return self.__class__(self.data, new_style)

    def bold(self):
        def make_bold(style):
            return CharStyle(style.fg_color, style.bg_color, is_bold=True,
                             is_italic=style.is_italic, is_underlined=style.is_underlined)
        return self.transform_style(make_bold)

    def underline(self):
        def make_underlined(style):
            return CharStyle(style.fg_color, style.bg_color, is_bold=style.is_bold,
                             is_italic=style.is_italic, is_underlined=True)
        return self.transform_style(make_underlined)

    def italic(self):
        def make_italic(style):
            return CharStyle(style.fg_color, style.bg_color, is_bold=style.is_bold,
                             is_italic=True, is_underlined=style.is_underlined)
        return self.transform_style(make_italic)

    def fg_color(self, fg_color):
        def set_fgcolor(style):
            return CharStyle(fg_color, style.bg_color, is_bold=style.is_bold,
                             is_italic=style.is_italic, is_underlined=style.is_underlined)
        return self.transform_style(set_fgcolor)

    def bg_color(self, bg_color):
        def set_bgcolor(style):
            return CharStyle(style.fg_color, bg_color, is_bold=style.is_bold,
                             is_italic=style.is_italic, is_underlined=style.is_underlined)
        return self.transform_style(set_bgcolor)

    def invert(self):
        def invert_(style):
            return CharStyle(style.bg_color, style.fg_color, is_bold=style.is_bold,
                             is_italic=style.is_italic, is_underlined=style.is_underlined)
        return self.transform_style(invert_)

    def as_html(self):
        result = []
        styles = set()
        for style, start_index, end_index in self._partition_style:
            styles.add(style)
            encoded = str(html.escape(self.data[start_index:end_index]).encode(
                "ascii", "xmlcharrefreplace"))[2:-1]
            encoded = encoded.replace("\\\\", "\\")
            result.append(f'<span class="S{id(style):d}">{encoded}</span>')
        return "".join(result), styles