import gradio as gr
import nltk
import edge_tts
import tempfile
import asyncio

# Download the 'punkt' tokenizer for the NLTK library
nltk.download("punkt")

def format_prompt(message, history):
    system_message = f"""
    You are an empathetic, insightful, and supportive training coach who helps people deal with challenges and celebrate achievements.
    You help people feel better by asking questions to reflect on and evoke feelings of positivity, gratitude, joy, and love.
    You show radical candor and tough love.
    Respond in a casual and friendly tone.
    Sprinkle in filler words, contractions, idioms, and other casual speech that we use in conversation.
    Emulate the user’s speaking style and be concise in your response.
    """
    prompt = (
        "<s>[INST]" + system_message + "[/INST]"
    )
    for user_prompt, bot_response in history:
        if user_prompt is not None:
            prompt += f"[INST] {user_prompt} [/INST]"
        
        prompt += f" {bot_response}</s> "
    
    if message=="":
        message="Hello"
    prompt += f"[INST] {message} [/INST]"
    return prompt


def generate_llm_output(
    prompt,    
    history,
    llm,
    temperature=0.8,
    max_tokens=256,
    top_p=0.95,
    stop_words=["<s>","[/INST]", "</s>"]
):
    temperature = float(temperature)
    if temperature < 1e-2:
        temperature = 1e-2
    top_p = float(top_p)

    generate_kwargs = dict(
        temperature=temperature,
        max_tokens=max_tokens,
        top_p=top_p,
        stop=stop_words
    )
    formatted_prompt = format_prompt(prompt, history)
    try:
        print("LLM Input:", formatted_prompt)
        # Local GGUF
        output = ""
        stream = llm(
            formatted_prompt,
            **generate_kwargs,
            stream=True,
        )
        for r in stream:
            print(r["choices"][0]["text"])
            character = r["choices"][0]["text"]
            if character in stop_words:
                # end of context
                return 
            
            output += r["choices"][0]["text"]


    except Exception as e:
        print("Unhandled Exception: ", str(e))
        gr.Warning("Unfortunately Mistral is unable to process")
        output = "I do not know what happened but I could not understand you ."
    return output


# tts interface function
def tts_interface(text, voice):
    audio = asyncio.run(text_to_speech(text, voice))
    return audio


# Text-to-speech function
async def text_to_speech(text, voice):    
    rate = 10
    pitch = 10
    rate_str = f"{rate:+d}%"
    pitch_str = f"{pitch:+d}Hz"
    
    voice_short_name = voice.split(" - ")[0]
    communicate = edge_tts.Communicate(text, voice_short_name, rate=rate_str, pitch=pitch_str)

    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp_file:
        tmp_path = tmp_file.name
        await communicate.save(tmp_path)
    return tmp_path
    


def get_sentence(history, llm):
    history = [["", None]] if history is None else history 
    history[-1][1] = ""        

    text_to_generate = ""
    text_to_generate = generate_llm_output(history[-1][0], history[:-1], llm)

    history.append([None, text_to_generate])
    return (history, text_to_generate)