#!/usr/bin/env python3

"""Simple script which adds matching rainbow colors to data read on stdin."""

import sys
from typing import Literal
from enum import auto, Enum
from dataclasses import dataclass, field

CSI = '\033['


def SGR(*xs):
    """
    Build a CSI escape sequence.

    https://en.wikipedia.org/wiki/ANSI_escape_code#CSI_(Control_Sequence_Introducer)_sequences
    """
    return CSI + ';'.join(str(x) for x in xs) + 'm'


class Color(Enum):
    """Known CSI colors."""

    # BLACK = 30
    RED = 31
    GREEN = auto()
    YELLOW = auto()
    BLUE = auto()
    MAGENTA = auto()
    CYAN = auto()


@dataclass
class ColorGenerator:
    """
    Handles colors by "depth".

    :param colors:
        List of colors to use.
    :param _idx:
        Current index.
    """

    colors: list[Color] = field(default_factory=lambda: list(Color))
    _idx: int = 0

    def __call__(self, c: int):
        """
        Update the current color, and return it.

        :param c:
            Any positive or negative number changes the interval value
            by 1 or -1. 0 simply returns the current value.
        """
        if c > 0:
            ret = self.colors[self._idx % len(self.colors)]
            self._idx += 1
        elif c < 0:
            self._idx -= 1
            ret = self.colors[self._idx % len(self.colors)]
        else:
            ret = self.colors[self._idx % len(self.colors)]
        return ret


def color(color: Color, c: str):
    """Write a highlighted string."""
    sys.stdout.write(SGR(1, color.value))
    sys.stdout.write(c)
    sys.stdout.write(SGR())


def __main() -> None:
    in_string: Literal[False] | Literal['"'] | Literal["'"] = False

    paren = 0
    brace = 0
    brack = 0

    col = ColorGenerator()

    while c := sys.stdin.read(1):
        match c:
            case '(':
                color(col(1), c)
                paren += 1
            case ')':
                color(col(-1), c)
                paren -= 1
            case '[':
                color(col(1), c)
                brack += 1
            case ']':
                color(col(-1), c)
                brack -= 1
            case '{':
                color(col(1), c)
                brace += 1
            case '}':
                color(col(-1), c)
                brace -= 1
            case "'":
                if in_string == "'":
                    in_string = False
                else:
                    in_string = "'"
                sys.stdout.write(c)
            case '"':
                if in_string == '"':
                    in_string = False
                else:
                    in_string = '"'
                sys.stdout.write(c)
            case '\\':
                sys.stdout.write(sys.stdin.read(1))
            case c:
                sys.stdout.write(c)


if __name__ == '__main__':
    __main()