import os
import requests
import gradio as gr
from langchain.llms.base import LLM
from langchain import LLMChain, PromptTemplate
from langchain.memory import ConversationBufferMemory
from groq import Groq

# ========== API Keys ==========
GROQ_API_KEY = "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"
ELEVENLABS_API_KEY = "sk_1253d5341f2936003bc788544d3206731872b30fa6b6dbfc"
ELEVENLABS_VOICE_ID = "21m00Tcm4TlvDq8ikWAM"

os.environ["GROQ_API_KEY"] = GROQ_API_KEY

# ========== Groq LLM via LangChain ==========

client = Groq(api_key=GROQ_API_KEY)

class GroqLLM(LLM):
    def _call(self, prompt: str, stop=None):
        completion = client.chat.completions.create(
            messages=[{"role": "user", "content": prompt}],
            model="llama-3.1-8b-instant",
        )
        return completion.choices[0].message.content

    @property
    def _llm_type(self) -> str:
        return "groq"

# ========== LangChain Chain Setup ==========

template = """You are a helpful assistant.
{chat_history}
User: {user_message}
Assistant:"""

prompt = PromptTemplate(
    input_variables=["chat_history", "user_message"],
    template=template
)

memory = ConversationBufferMemory(memory_key="chat_history")

llm_chain = LLMChain(
    llm=GroqLLM(),
    prompt=prompt,
    memory=memory
)

# ========== ElevenLabs TTS Function ==========

def generate_audio_elevenlabs(text, filename="response_audio.mp3"):
    url = f"https://api.elevenlabs.io/v1/text-to-speech/{ELEVENLABS_VOICE_ID}"
    headers = {
        "xi-api-key": ELEVENLABS_API_KEY,
        "Content-Type": "application/json"
    }
    data = {
        "text": text,
        "model_id": "eleven_monolingual_v1"
    }

    response = requests.post(url, headers=headers, json=data)

    if response.status_code == 200:
        with open(filename, "wb") as f:
            f.write(response.content)
        return filename
    else:
        print("Error generating audio:", response.status_code, response.text)
        return None

# ========== Chatbot Handler ==========

def respond(message, chat_history):
    response = llm_chain.predict(user_message=message)
    audio_file = generate_audio_elevenlabs(response)

    chat_history.append((message, response))
    return chat_history, audio_file

# ========== Gradio UI ==========

with gr.Blocks() as demo:
    chatbot = gr.Chatbot(label="Groq + ElevenLabs Chatbot")
    msg = gr.Textbox(label="Ask something...")
    audio_output = gr.Audio(label="Audio Reply", autoplay=True, type="filepath")

    state = gr.State([])

    msg.submit(respond, [msg, state], [chatbot, audio_output])
    msg.submit(lambda: "", None, msg)  # Clear textbox after sending

# ========== Launch App ==========
demo.launch()