import re
import gradio as gr
from transformers import pipeline

# Set up the generatove model transformer pipeline
generator = pipeline("text-generation", model="tmp_trainer")

# A sequence of lines both those typed in and the line so far
# when save is clicked the txt file is downloaded
source_lines = []
generated_lines = []


def twolists(list1, list2):
    newlist = []
    a1 = len(list1)
    a2 = len(list2)

    for i in range(max(a1, a2)):
        if i < a1:
            newlist.append(list1[i])
        if i < a2:
            newlist.append(list2[i])

    return newlist


def build_poem(source_lines, generated_lines):
    # This function structures the space already, set boundarys to the possible.
    # return a list with the two lists interspersed
    return twolists(source_lines, generated_lines)


def add_to_lines(new_line):
    source_lines.append(new_line)
    # NOTE: pass the whole array of lines to the generator
    # There is a compounding of the effect, a spiral of interaction.
    poem = build_poem(source_lines, generated_lines)

    # pass the last two lines of the poem
    prompt = "\n".join(poem[-2:])
    result = generator(prompt, max_length=50, num_return_sequences=3)
    response = result[0]["generated_text"]

    # Get only the new generated text
    response = response.replace(prompt, "")

    generated_lines.append(response)

    lines = []
    for line in build_poem(source_lines, generated_lines):
        if (len(lines) % 2) == 1:
            line = f"<p style='text-align: right;'>{line}</p>"
        else:
            line = f"<p>{line}</p>"
        lines.append(line)

    html = "".join(lines)

    return html


demo = gr.Interface(
    fn=add_to_lines,
    inputs=gr.Textbox(
        lines=1,
        value="The drought had lasted now for ten million years, and the reign of the terrible lizards had long since ended.",
    ),
    outputs="html",
    allow_flagging="never",
)


if __name__ == "__main__":
    demo.launch()


def downloadtext():
    # somehow print the values from the list
    pass