import json
import streamlit as st
import requests

headers = {"Authorization": f"Bearer api_LbZppGQTIlpuKxWWbyNLvgPXLxXCbKYiMr"}
API_URL = "https://api-inference.huggingface.co/models/vblagoje/bart_eli5"
API_URL_TTS = "https://api-inference.huggingface.co/models/espnet/kan-bayashi_ljspeech_vits"


def query_eli_model(payload):
    data = json.dumps(payload)
    response = requests.request("POST", API_URL, headers=headers, data=data)
    return json.loads(response.content.decode("utf-8"))

def query_audio_tts(payload):
    data = json.dumps(payload)
    response = requests.request("POST", API_URL_TTS, headers=headers, data=data)
    return response.content

st.set_page_config(
    page_title="AI assistant",
    initial_sidebar_state="expanded"
)

st.markdown(""" 
    <style> 
        .row-widget.stTextInput > div:first-of-type {
            background: #fff;
            display: flex;
            border: 1px solid #dfe1e5;
            box-shadow: none;
            border-radius: 24px;
            height: 50px;
            width: auto;
            padding: 10px;
            margin: 10px auto 30px;
        }
        .row-widget.stTextInput > div:first-of-type:hover,
        .row-widget.stTextInput > div:first-of-type:focus {
            box-shadow: 1px 1px 2px 1px rgba(0, 0, 0, 0.2);
        }
        .row-widget.stTextInput .st-bq {
            background-color: #fff;
        }
        .row-widget.stButton > button {
            border-radius: 24px;
            background-color: #B6C9B1;
            color: #fff;
            border: none;
            padding: 6px 20px;
            float: right;
            background-image: none;
        }
        .row-widget.stButton > button:hover {
            box-shadow: 1px 1px 2px 1px rgba(0, 0, 0, 0.2);
        }
        .row-widget.stButton > button:focus {
            border: none;
            color: #fff;
        }
    </style> """, unsafe_allow_html=True)

st.title('AI Assistant')

question = st.text_input('Enter a question')

if question:
    with st.spinner("Generating an answer..."):

        model_input = f'question: {question}'

        if model_input:
            data = query_eli_model({
                "inputs": model_input, 
                "parameters": {
                    "min_length": 30,
                    "max_length": 200,
                    "do_sample": False, 
                    "early_stopping": True,
                    "num_beams": 8,
                    "temperature": 1.0,
                    "top_k": None,
                    "top_p": None,
                    "no_repeat_ngram_size": 3,
                    "num_return_sequences": 1
                }
            })
        

            if data:
                generated_answer = data[0]['generated_text']

                st.markdown(
                    " ".join([
                        "<div style='margin-bottom: 50px;'>",
                        "<h2 class='font-title text-bold'>The answer:</h2>",
                        '<div style="padding: 30px;background-color: #B6C9B1; border-radius: 10px;">',
                        f'<p>{generated_answer}</p>',
                        "</div>",
                        "</div>"
                    ]),
                    unsafe_allow_html=True
                )

                if generated_answer:
                    with st.spinner("Generating an audio..."):
                        audio_file = query_audio_tts({
                            "inputs": generated_answer,
                            "parameters": {
                                "vocoder_tag": "str_or_none(none)",
                                "threshold": 0.5,
                                "minlenratio": 0.0,
                                "maxlenratio": 10.0,
                                "use_att_constraint": False,
                                "backward_window": 1,
                                "forward_window": 3,
                                "speed_control_alpha": 1.0,
                                "noise_scale": 0.333,
                                "noise_scale_dur": 0.333
                            }
                        })

                        with open("out.flac", "wb") as f:
                            f.write(audio_file)
                            
                            st.audio("out.flac")