import gradio as gr
from huggingface_hub import InferenceClient


client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")


def generate_text(messages):
    print("generate_text")
    print(messages)
    generated = ""
    for token in client.chat_completion(messages, max_tokens=50,stream=True):
        content = (token.choices[0].delta.content)
        generated += content
        #print(content)
        #print(''.join(list(content)))
        yield generated
        #print(token.choices[0].delta)
    
    #return generated+"." #no stram version

def call_generate_text(message, history):
    #if len(message) == 0:
    #    messages.append({"role": "system", "content": "you response around 10 words"})
   
    print(message)
    print(history)

    user_message = [{"role":"user","content":message}]
    messages = history + user_message
    try:
        
        assistant_message={"role":"assistant","content":""}
        text_generator = generate_text(messages)

        for text_chunk in text_generator:
            print(f"chunk={text_chunk}")
            assistant_message["content"] = text_chunk
            updated_history = messages + [assistant_message]
            yield "", updated_history

    except RuntimeError  as e:
        print(f"An unexpected error occurred: {e}")
        yield  "",history

head = '''
<script src="https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/ort.webgpu.min.js" ></script>
<script type="module">
        import { matcha_tts,env ,start_thread_play_tts} from "https://akjava.github.io/Matcha-TTS-Japanese/js-esm/matcha_tts_onnx_en.js";
        window.MatchaTTSEn = matcha_tts
        window.start_threads = async function(){
            //sadly onload unstable
            //console.log("start_threads")
            await start_thread_play_tts();
            await start_multi_line_tts();
        }
</script>

<script>
let is_start_threads = false
let last_chatbot_size = 0
let tts_text_index = 0
let tts_texts = []

const interval = 10

async function start_multi_line_tts() {
  //console.log("start_multi_line_tts")
  //console.log(tts_texts.length)
  if (tts_texts.length > tts_text_index){
  const tts_text = tts_texts[tts_text_index]
  tts_text_index += 1
  console.log(tts_text)
  if (tts_text!=""){
    //console.log("statar await load")
    await window.MatchaTTSEn(tts_text,"/file=models/ljspeech_sim.onnx")
    //console.log("end await ")
  }

  
  }
  setTimeout(start_multi_line_tts, interval);
}


function reset_tts_text(){
    if (is_start_threads == false){
        window.start_threads()
        is_start_threads = true
    }
    console.log("new messages")
    tts_text_index = 0
    tts_texts = []
}
function replaceSpecialChars(text) {
    const pattern = /[^a-zA-Z0-9,.!?-_']/g;
    return text.replace(pattern, ' ');
}


function update_tts_texts(text){
    //console.log(text)
    const replaced_text = replaceSpecialChars(text)
    const new_texts = []
    const splited = replaced_text.split(/[.!?]+\s/);
    for (let i = 0; i < splited.length; i++) {
    const value = splited[i].trim();
    
    if (i === splited.length - 1) {
        if (value.endsWith(".") || value.endsWith("?") || value.endsWith("!")){
            new_texts.push(value);
        }
        //console.log("Last element:", value);
    } else {
        // その他の要素に対する処理
        new_texts.push(value);
    }
    }
    tts_texts=new_texts
    
}

function update_chatbot(chatbot){
    //console.log(chatbot)
    if (chatbot.length!=last_chatbot_size){
        last_chatbot_size = chatbot.length
        reset_tts_text()
    }
    text = (chatbot[chatbot.length -1])["content"]
    update_tts_texts(text)
    
}

window.replaceSpecialChars = replaceSpecialChars

window.update_chatbot = update_chatbot
window.update_tts_texts = update_tts_texts
window.reset_tts_text = reset_tts_text  

</script>
'''

with gr.Blocks(title="LLM with TTS",head=head) as demo:
    gr.Markdown("## A LLM is unstable:The inference client used in this demo exhibits inconsistent performance. While it can provide responses in milliseconds, it sometimes becomes unresponsive and times out.")
    gr.Markdown("## TTS talke a long loading time:Please be patient, the first response may have a delay of up to over 20 seconds while loading.")
    gr.Markdown("**Mistral-7B-Instruct-v0.3/LJSpeech**.LLM and TTS models will change without notice.")
    
    js = """
    function(chatbot){
    window.update_chatbot(chatbot)
    //text = (chatbot[chatbot.length -1])["content"]
    //tts_text = window.replaceSpecialChars(text)
    //console.log(tts_text)
    //window.MatchaTTSEn(tts_text,"/file=models/ljspeech_sim.onnx")
    }
    """
    chatbot = gr.Chatbot(type="messages")
    chatbot.change(None,[chatbot],[],js=js)
    msg = gr.Textbox()
    with gr.Row():
        clear = gr.ClearButton([msg, chatbot])
        submit = gr.Button("Submit",variant="primary").click(call_generate_text, inputs=[msg, chatbot], outputs=[msg,chatbot])

    gr.HTML("""
    <br>
    <div id="footer">
    <b>Spaces</b><br>
     <a href="https://huggingface.co/spaces/Akjava/matcha-tts_vctk-onnx" style="font-size: 9px" target="link">Match-TTS VCTK-ONNX</a> | 
     <a href="https://huggingface.co/spaces/Akjava/matcha-tts-onnx-benchmarks" style="font-size: 9px" target="link">Match-TTS ONNX-Benchmark</a> | 
     <a href="https://huggingface.co/spaces/Akjava/AIChat-matcha-tts-onnx-en" style="font-size: 9px" target="link">AIChat-Matcha-TTS ONNX English</a> | 
     
      <br><br>
    <b>Credits</b><br>
    <a href="https://github.com/akjava/Matcha-TTS-Japanese" style="font-size: 9px" target="link">Matcha-TTS-Japanese</a> | 
    <a href = "http://www.udialogue.org/download/cstr-vctk-corpus.html" style="font-size: 9px"  target="link">CSTR VCTK Corpus</a> |
    <a href = "https://github.com/cmusphinx/cmudict" style="font-size: 9px"  target="link">CMUDict</a> |
    <a href = "https://huggingface.co/docs/transformers.js/index" style="font-size: 9px"  target="link">Transformer.js</a> |
    <a href = "https://huggingface.co/cisco-ai/mini-bart-g2p" style="font-size: 9px"  target="link">mini-bart-g2p</a> |
    <a href = "https://onnxruntime.ai/docs/get-started/with-javascript/web.html" style="font-size: 9px"  target="link">ONNXRuntime-Web</a> |
    <a href = "https://github.com/akjava/English-To-IPA-Collections" style="font-size: 9px"  target="link">English-To-IPA-Collections</a> |
    <a href ="https://huggingface.co/papers/2309.03199" style="font-size: 9px"  target="link">Matcha-TTS Paper</a>
    </div>
    """)
    
    msg.submit(call_generate_text, [msg, chatbot], [msg, chatbot])

import os
dir ="/home/user/app/"
dir = "C:\\Users\\owner\\Documents\\pythons\\huggingface\\mistral-7b-v0.3-matcha-tts-en"
demo.launch(allowed_paths=[os.path.join(dir,"models","ljspeech_sim.onnx")])