import os
import gradio as gr
import requests
from sentence_transformers import SentenceTransformer, util
import torch
import json
import urllib.parse
import soundfile as sf

# Fetch Hugging Face API Token securely from environment variables
HF_API_TOKEN = os.getenv("HF_API_TOKEN")  # This fetches the token securely

WHISPER_API_URL = "https://api-inference.huggingface.co/models/openai/whisper-large"
LLAMA_API_URL = "https://api-inference.huggingface.co/models/meta-llama/Llama-2-7b-chat-hf"

# Load SentenceTransformer model for retrieval
retriever_model = SentenceTransformer("distiluse-base-multilingual-cased-v2")

# Load dataset
with open("qa_dataset.json", "r", encoding="utf-8") as f:
    qa_data = json.load(f)

# Function to transcribe audio using Whisper
def transcribe_audio(audio_file):
    headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
    with open(audio_file, "rb") as f:
        response = requests.post(WHISPER_API_URL, headers=headers, data=f)
    return response.json()["text"]

# Function to generate TTS audio URL (Google Translate API for Tamil Voice)
def get_tts_audio_url(text, lang="ta"):
    # URL encode the text to ensure special characters are handled
    safe_text = urllib.parse.quote(text)
    return f"https://translate.google.com/translate_tts?ie=UTF-8&q={safe_text}&tl={lang}&client=tw-ob"

# Function to retrieve a relevant response from the Q&A dataset using SentenceTransformer
def get_bot_response(query):
    query_embedding = retriever_model.encode(query, convert_to_tensor=True)
    qa_embeddings = retriever_model.encode([qa["question"] for qa in qa_data], convert_to_tensor=True)

    scores = util.pytorch_cos_sim(query_embedding, qa_embeddings)
    best_idx = torch.argmax(scores)

    top_qa = qa_data[best_idx]
    prompt = f"User asked: {query}\nRelevant FAQ: {top_qa['question']}\nAnswer: {top_qa['answer']}\nNow generate a helpful and fluent Tamil response to the user query."

    # Use LLaMA for generating the refined response
    headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
    payload = {
        "inputs": prompt,
        "parameters": {"temperature": 0.7, "max_new_tokens": 150, "return_full_text": False},
    }
    response = requests.post(LLAMA_API_URL, headers=headers, json=payload)
    result = response.json()

    if isinstance(result, list) and "generated_text" in result[0]:
        return result[0]["generated_text"]
    else:
        return "மன்னிக்கவும், நான் இந்த கேள்விக்கு பதில் தர முடியவில்லை."

# Gradio interface function
def chatbot(audio, message, system_message, max_tokens, temperature, top_p):
    if audio is not None:
        sample_rate, audio_data = audio  # ✅ Correct order
        sf.write("temp.wav", audio_data, sample_rate)  # Save audio
        try:
            transcript = transcribe_audio("temp.wav")
            message = transcript  # Use transcribed text
        except Exception as e:
            return f"Audio transcription failed: {str(e)}", None

    try:
        response = get_bot_response(message)
        audio_url = get_tts_audio_url(response)
        return response, audio_url
    except Exception as e:
        return f"Error in generating response: {str(e)}", None


# Define Gradio interface
demo = gr.Interface(
    fn=chatbot,
    inputs=[
        gr.Audio(type="numpy", label="Speak to the Bot"),  # Adjusted for microphone input
        gr.Textbox(value="How can I help you?", label="Text Input (optional)"),
        gr.Textbox(value="You are a friendly chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
    ],
    outputs=[gr.Textbox(label="Response"), gr.Audio(label="Bot's Voice Response (Tamil)")],
    live=True,
)

if __name__ == "__main__":
    demo.launch()