import tokenize
import token as token_type
import keyword
import io


def tokenize_with_whitespace(readline):
    """
    like tokenize.tokenize, except it tries to yield a token for every character in the file,
    including non-significant whitespace.
    yielded values are (string, token_type) tuples.
    Non-significant whitespace has a token_type of None.
    """
    prev_end = (0,0)
    for token in tokenize.tokenize(file.readline):
        #check if there was whitespace between the last token and this one
        if token.start != prev_end:
            #we only really care about horizontal whitespace. Vertical whitespace is usually present in the token stream as NEWLINE or similar.
            if token.start[0] - prev_end[0] == 0:
                dx = token.start[1] - prev_end[1]
                yield (" "*dx, None)
        yield (token.string, token.type)
        prev_end = token.end

def apply_style(readline, style):
    result = []
    for s, type in tokenize_with_whitespace(readline):
        if type in style:
            result.append(style[type](s))
        else:
            result.append(s)
    return "".join(result)

# bbcode_color = lambda s, color: f"[color={color}]{s}[/color]"
# bbcode_style = {
    # token_type.ENCODING: lambda s: "",
    # token_type.NAME: lambda s: bbcode_color(s, "blue") if keyword.iskeyword(s) else s,
    # token_type.STRING: lambda s: bbcode_color(s, "red")
# }

import colorama
colorama.init()
colorama_style = {
    token_type.ENCODING: lambda s: "",
    token_type.NAME: lambda s: colorama.Fore.CYAN + s + colorama.Fore.RESET if keyword.iskeyword(s) else s,
    token_type.STRING: lambda s: colorama.Fore.RED + s + colorama.Fore.RESET
}

with open("sample.py", "rb") as file:
    print(apply_style(file.readline, colorama_style))