import datetime
import threading
from langchain_core.prompts import ChatPromptTemplate
import time
from langchain_groq import ChatGroq
import gradio as gr
import os
from gtts import gTTS
import io
import pygame
groq_api_key = os.environ.get("GROQ_API_KEY")

llm = ChatGroq(
    model="mixtral-8x7b-32768",  
    temperature=0.7,
    api_key=groq_api_key,
    verbose = False
)

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "You are an intelligent AI assistant. Answer the questions as accurately as possible."),
        ("human", "{input}")
    ]
)

chain = prompt | llm

convo_history = []
audio_playing = False

def stop_audio():
    global audio_playing
    if audio_playing:
        pygame.mixer.music.stop()
        audio_playing = False

def play_tts(text):
    global audio_playing
    audio_playing = True

    def speak_async():
        global audio_playing
        pygame.mixer.init()
        fp = io.BytesIO()
        gTTS(text).write_to_fp(fp)
        fp.seek(0)
        pygame.mixer.music.load(fp)
        pygame.mixer.music.play()
        while pygame.mixer.music.get_busy():
            pygame.time.Clock().tick(10)
        audio_playing = False
        pygame.mixer.quit()

    thread = threading.Thread(target=speak_async)
    thread.start()

def ChatBot(text):

    global convo_history, audio_playing

    stop_audio()

    convo_history.append(f"Human: {text}")

    full_context = "\n".join(convo_history)
    response = chain.invoke({"input": full_context}).content

    convo_history.append(f"AI: {response}")

    print("\n".join(convo_history))

    play_tts(response)

    return "\n".join(convo_history)


demo = gr.Interface(
    fn=ChatBot,
    inputs="text",
    outputs="text",
    title="AudioBOT",
    description="Ask any question to the AI. It remembers the context of the conversation and speaks its response.",
)

demo.launch()