import re
import random
import math

text = """
What the fuck did you just fucking say about me, you little bitch? I'll have you know I graduated top of my class in the Navy Seals, and I've been involved in numerous secret raids on Al-Quaeda, and I have over 300 confirmed kills. I am trained in gorilla warfare and I'm the top sniper in the entire US armed forces. You are nothing to me but just another target. I will wipe you the fuck out with precision the likes of which has never been seen before on this Earth, mark my fucking words. You think you can get away with saying that shit to me over the Internet? Think again, fucker. As we speak I am contacting my secret network of spies across the USA and your IP is being traced right now so you better prepare for the storm, maggot. The storm that wipes out the pathetic little thing you call your life. You're fucking dead, kid. I can be anywhere, anytime, and I can kill you in over seven hundred ways, and that's just with my bare hands. Not only am I extensively trained in unarmed combat, but I have access to the entire arsenal of the United States Marine Corps and I will use it to its full extent to wipe your miserable ass off the face of the continent, you little shit. If only you could have known what unholy retribution your little "clever" comment was about to bring down upon you, maybe you would have held your fucking tongue. But you couldn't, you didn't, and now you're paying the price, you goddamn idiot. I will shit fury all over you and you will drown in it. You're fucking dead, kiddo.
""".strip()

highlight = ["fuck", "shit", "bitch", "dead", "kill", "kid", "combat"]

foregrounds = list(range(30, 38))
backgrounds = list(range(40, 48))

def break_paragraph_greedy(tokens, width):
    lines = []
    length = 0
    line = []
    for token in tokens:
        if length + len(token) + len(line) > width:
            lines.append(line)
            line = [token]
            length = len(token)
        else:
            line.append(token)
            length += len(token)
    if len(line) > 0:
        lines.append(line)
    return lines

# algorithm recursion is simple:
# knuth_plass(tokens, width, i) = the score of the paragraph starting at token i
# kp[i] stores the best break index for the paragraph starting at token[i]
def knuth_plass_core(kp, tokens, width, max_width, i):
    assert(len(kp) == len(tokens))
    if kp[i] != None:
        return kp[i]
    else:
        # try all breaks from i to n
        this_line_n_tokens = 1 # place at least the first token
        this_line_length = len(tokens[i]) # place at least the first token
        best_score = math.inf
        best_break = None
        is_last_line = True
        #score = lambda: (width - this_line_length) ** 2
        # break is between tokens[i] and tokens[j]
        for j in range(i + 1, len(tokens)):
            # this_line_n_tokens is the number of spaces the line is about to have
            if this_line_length + len(tokens[j]) + this_line_n_tokens > max_width: # exceeds our hard cutoff width
                is_last_line = False
                break
            if j == len(tokens) - 1:
                score = 0
            else:
                score = (width - (this_line_length + len(tokens[j]) + this_line_n_tokens))**4 + knuth_plass_core(kp, tokens, width, max_width, j + 1)
            if score < best_score:
                best_score = score
                best_break = j + 1
            this_line_n_tokens += 1
            this_line_length += len(tokens[j])
        # special case for last line, don't penalize for not taking up the full line
        if is_last_line:
            best_break = None
            best_score = 0
        #assert(best_break != None)
        kp[i] = best_break
        return best_score

def knuth_plass(tokens, width, max_width):
    kp = [None for _ in tokens]
    knuth_plass_core(kp, tokens, width, max_width, 0)
    lines = []
    last_index = 0
    while last_index != None:
        lines.append(tokens[last_index:kp[last_index]])
        last_index = kp[last_index]
    return lines

def should_highlight(str: str):
    for p in highlight:
        if p in str:
            return True
    return False

def color(token: str):
    subtokens = re.findall(r"[\w'-]+|[.,!?;]", token)
    for i, subtoken in enumerate(subtokens):
        if subtoken in ".,!?;":
            pass
        else:
            if should_highlight(subtoken) or random.randint(1, 10) <= 2:
                f = random.choice(foregrounds)
                b = random.choice(backgrounds)
                subtokens[i] = f"\033[{f};{b}m{subtoken}\033[0m"
            else:
                f = random.choice(foregrounds)
                subtokens[i] = f"\033[{f}m{subtoken}\033[0m"
    return "".join(subtokens)

def main():
    tokens = text.split(" ")

    #lines = break_paragraph_greedy(tokens, 40)
    lines = knuth_plass(tokens, 40, 42)

    text_index = 0
    character_count = 0
    print("```ansi")
    for line in lines:
        colored_line = []
        for token in line:
            colored_line.append(color(token))
            text_index += 1
        colored_line = " ".join(colored_line)
        l = len(colored_line.encode("utf-8"))
        if character_count + l < 4000 - len("```ansi\n```"):
            character_count += l + 1
            print(colored_line)
        else:
            print("```")
            print("```ansi")
            print(colored_line)
            character_count = l + 1
    print("```")

main()