import gradio as gr
from llama_cpp import Llama
import urllib

# Initialize the Llama model
llm = Llama(model_path="ggml-model-Q4_K_M.gguf", chat_format="llama-3")


def completion_to_song(c):
    lines = c.split("\n")
    kept_lines = []
    notes = False
    for l in lines:

        # Record if we've hit music
        if "|" in l:
            notes = True

        # Stop if we then go back to the start of another song
        if "X" in l and notes:
            break
        if "T" in l and notes:
            break

        # Stop on an empty line
        if len(l.strip()) < 2 and notes:
            break

        # Otherwise keep the line
        kept_lines.append(l)

    return "\n".join(kept_lines)


def generate_html(input):
    song = completion_to_song(input)
    url_song = urllib.parse.quote(song, safe="~@#$&()*!+=:;,?/'")
    html_song = song.replace("\n", "<br>")
    html_text = (
        '<p><a href="https://editor.drawthedots.com?t='
        + url_song
        + '" target="_blank"><b>EDIT LINK - click to open abcjs editor (allows download and playback)</b></a></p>'
    )
    return html_text


# Function to generate ABC notation
def generate_abc(input):
    messages = [
        {
            "role": "system",
            "content": "you are a helpful assistant that generates ABC notation music. Only use ABC notation music",
        },
        {"role": "user", "content": f"{input}"},
    ]
    response = llm.create_chat_completion(messages=messages, temperature=1, top_k=10)
    return response["choices"][0]["message"][
        "content"
    ]  # Assuming this returns valid ABC notation


# Gradio Blocks Interface
with gr.Blocks(
    head="""<script src="https://cdnjs.cloudflare.com/ajax/libs/abcjs/6.4.4/abcjs-basic-min.js" type="text/javascript"></script>"""
) as demo:
    gr.Markdown("# ABC Notation Generator with LLM")
    gr.Markdown(
        "Enter a prompt to generate music in ABC notation, and view the rendered sheet music."
    )

    # Input section
    with gr.Row():
        user_input = gr.Textbox(
            label="Enter your prompt",
            value="Design melodic orchestrations by following the indicated chord progression. 'A', 'D', 'E7', 'A', 'E/G#', 'A', 'Bm', 'A7/C#', 'D', 'E7', 'A', 'A', 'D', 'A', 'A', 'D', 'A', 'A', 'D', 'A', 'D', 'A/D#', 'E', 'A', 'D', 'A', 'A', 'D', 'A', 'E7'",
        )

    # Output section
    with gr.Row():
        abc_output = gr.Textbox(label="Output ABC notation")

    # Button to trigger generation
    with gr.Row():
        abc_html = gr.HTML(label="link")
    submit_btn = gr.Button("Generate Music")

    # Button click logic
    def on_submit(input_text):
        abc_notation = generate_abc(input_text)  # Generate ABC notation
        html_output = generate_html(abc_notation)
        return abc_notation, html_output

    submit_btn.click(on_submit, inputs=user_input, outputs=[abc_output, abc_html])

# Launch the Gradio app
demo.launch()