import py
import _formatting

class colorstr(str):
    singlecolor = '\x00'

    def __init__(self, obj):
        if isinstance(obj, colorstr) and self.singlecolor == '\x00':
            self.attrs = obj.attrs
        else:
            self.attrs = self.singlecolor * len(self)

    def __add__(self, other):
        s = colorstr(str.__add__(self, other))
        s.attrs = self.attrs + colorstr(other).attrs
        return s

    def __radd__(self, other):
        s = colorstr(str.__add__(other, self))
        s.attrs = colorstr(other).attrs + self.attrs
        return s

    def __mul__(self, count):
        s = colorstr(str.__mul__(self, count))
        s.attrs = self.attrs * count
        return s

    __rmul__ = __mul__

    def __getitem__(self, index):
        s = colorstr(str.__getitem__(self, index))
        s.attrs = self.attrs[index]
        return s

    def __getslice__(self, start, end):
        s = colorstr(str.__getslice__(self, start, end))
        s.attrs = self.attrs.__getslice__(start, end)
        return s

##    def __mod__(self, values):
##        import _formatting
##        if isinstance(values, tuple):
##            return _formatting.format(self, values, None)
##        else:
##            if hasattr(values, 'keys'):
##                return _formatting.format(self, (values,), values)
##            else:
##                return _formatting.format(self, (values,), None)

    def capitalize(self):
        s = colorstr(str.capitalize(self))
        s.attrs = self.attrs
        return s

    def center(self, width):
        d = width - len(self) 
        if d>0:
            offset = d//2
            centered = offset * ' ' + self + (d - offset) * ' '
        else:
            centered = self
        return centered

    def join(self, seq):
        if not isinstance(seq, (list, tuple)):
            seq = list(seq)
        text = str_join(self, seq)
        for x in seq:
            if not isinstance(x, str):
                return text
        text = colorstr(text)
        text.attrs = str_join(self.attrs, [colorstr(x).attrs for x in seq])
        return text

    def __str__(self):
        return self

    def __repr__(self):
        return '<%s %s>' % (self.__class__.__name__, str.__repr__(self))

    def uncolored(self):
        return uncolored(self)


class red(colorstr):    singlecolor = chr(31)
class green(colorstr):  singlecolor = chr(32)

uncolored = str.__str__


def ansitext(s):
    if isinstance(s, colorstr):
        result = []
        currentcolor = '\x00'
        for c, a in zip(s, s.attrs):
            if a != currentcolor:
                result.append('\x1b[%dm' % ord(a))
                currentcolor = a
            result.append(c)
        if currentcolor != '\x00':
            result.append('\x1b[0m')
        s = str_join('', result)
    return s


class ANSIFile:

    def __init__(self, output):
        self._output = output

    def write(self, s):
        self._output.write(ansitext(s))

    def writelines(self, lines):
        self._output.writelines(map(ansitext, lines))


def invoke():
    import sys
    if sys.stdout.isatty():
        sys.stdout = ANSIFile(sys.stdout)

# ____________________________________________________________
#
# Patching some methods of str...  :-(

def colorjoin(self, seq):
    if not isinstance(seq, (list, tuple)):
        seq = list(seq)
    for x in seq:
        if isinstance(x, colorstr):
            return colorstr(self).join(seq)
    return str_join(self, seq)

def colormod(self, values):
    if isinstance(values, tuple):
        return _formatting.format(self, values, None)
    else:
        if hasattr(values, 'keys'):
            return _formatting.format(self, (values,), values)
        else:
            return _formatting.format(self, (values,), None)

str_join = str.join
str_mod  = str.__mod__
py.magic.extractdict(str)['join'] = colorjoin
py.magic.extractdict(str)['__mod__'] = colormod   DOES NOT WORK

# ____________________________________________________________

if __name__ == '__main__':
    invoke()