import re
import gradio as gr
from transformers import pipeline

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

# 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


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

    return twolists(source_lines, generated_lines)


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

    print("__poem")
    print(poem)

    # pass the full poem into the generator
    result = generator("".join(poem), max_length=30, num_return_sequences=3)
    # result = generator(poem, max_length=30, num_return_sequences=3)

    print("____result")
    print(result[0]["generated_text"])

    response = result[0]["generated_text"]

    # TODO: remove each source_line from the poem
    response = response.replace("".join(poem), "")

    print("____response")
    print(response)

    generated_lines.append(response)

    # html = [f"<p>{line}</p>" for line in build_poem(source_lines, generated_lines)]
    html = [f"{line}<br/>" for line in build_poem(source_lines, generated_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