import time

import gradio as gr
from sentence_transformers import SentenceTransformer

import httpx
import json

from utils import get_tags_for_prompts, get_mubert_tags_embeddings, get_pat

minilm = SentenceTransformer('all-MiniLM-L6-v2')
mubert_tags_embeddings = get_mubert_tags_embeddings(minilm)


def get_track_by_tags(tags, pat, duration, maxit=20, loop=False):
    if loop:
        mode = "loop"
    else:
        mode = "track"
    r = httpx.post('https://api-b2b.mubert.com/v2/RecordTrackTTM',
                   json={
                       "method": "RecordTrackTTM",
                       "params": {
                           "pat": pat,
                           "duration": duration,
                           "tags": tags,
                           "mode": mode
                       }
                   })

    rdata = json.loads(r.text)
    assert rdata['status'] == 1, rdata['error']['text']
    trackurl = rdata['data']['tasks'][0]['download_link']

    print('Generating track ', end='')
    for i in range(maxit):
        r = httpx.get(trackurl)
        if r.status_code == 200:
            return trackurl
        time.sleep(1)


def generate_track_by_prompt(email, prompt, duration, loop=False):
    try:
        pat = get_pat(email)
        _, tags = get_tags_for_prompts(minilm, mubert_tags_embeddings, [prompt, ])[0]
        return get_track_by_tags(tags, pat, int(duration), loop=loop), "Success", ",".join(tags)
    except Exception as e:
        return None, str(e), ""


block = gr.Blocks()

with block:
    gr.HTML(
        """
            <div style="text-align: center; max-width: 700px; margin: 0 auto;">
              <div
                style="
                  display: inline-flex;
                  align-items: center;
                  gap: 0.8rem;
                  font-size: 1.75rem;
                "
              >
                <h1 style="font-weight: 900; margin-bottom: 7px;">
                  Mubert
                </h1>
              </div>
              <p style="margin-bottom: 10px; font-size: 94%">
                All music is generated by Mubert API – <a href="https://mubert.com" style="text-decoration: underline;" target="_blank">www.mubert.com</a>
              </p>
            </div>
        """
    )
    with gr.Group():
        with gr.Box():
            email = gr.Textbox(label="email")
            prompt = gr.Textbox(label="prompt")
            duration = gr.Slider(label="duration (seconds)", value=30)
            is_loop = gr.Checkbox(label="Generate loop")
            out = gr.Audio()
            result_msg = gr.Text(label="Result message")
            tags = gr.Text(label="Tags")
            btn = gr.Button("Submit").style(full_width=True)
        
        btn.click(fn=generate_track_by_prompt, inputs=[email, prompt, duration, is_loop], outputs=[out, result_msg, tags])
        gr.HTML('''
        <div class="footer" style="text-align: center; max-width: 700px; margin: 0 auto;">
                    <p>Demo by <a href="https://huggingface.co/Mubert" style="text-decoration: underline;" target="_blank">Mubert</a>
                    </p>
        </div>
        ''')

block.launch()