import os
import yaml
from fuzzywuzzy import fuzz
from huggingface_hub import InferenceClient
import gradio as gr

# 🔐 Secure token from Hugging Face Secrets
API_KEY = os.getenv("HF_TOKEN")
MODEL_NAME = "meta-llama/Llama-2-7b-chat-hf"
client = InferenceClient(model=MODEL_NAME, token=API_KEY)

# 📦 Load and validate YAML dataset
def load_yaml(path="data.yaml"):
    if not os.path.exists(path):
        print(f"❌ YAML file not found: {path}")
        return []
    try:
        with open(path, 'r', encoding='utf-8') as f:
            data = yaml.safe_load(f)
            print(f"✅ Loaded {len(data)} items from YAML")
            return data
    except Exception as e:
        print("❌ YAML loading error:", e)
        return []

qa_data = load_yaml()

# 🔍 Fuzzy match logic across qa_pairs and descriptions
def find_best_match(user_question, threshold=70):
    best_score = 0
    best_item = None
    best_answer = None

    for item in qa_data:
        # Match in qa_pairs
        for qa in item.get("qa_pairs", []):
            for q_field in ["question_bn", "question_en"]:
                question = qa.get(q_field, "")
                score = fuzz.partial_ratio(user_question.lower(), question.lower())
                if score > best_score and score >= threshold:
                    best_score = score
                    best_item = item
                    best_answer = qa.get("answer_bn") if "question_bn" in q_field else qa.get("answer_en")

        # Fallback to description_bn / description_en
        for field in ["description_bn", "description_en"]:
            if field in item:
                score = fuzz.partial_ratio(user_question.lower(), item[field].lower())
                if score > best_score and score >= threshold:
                    best_score = score
                    best_item = item
                    best_answer = item[field]

    return best_item, best_answer

# 🧠 Prompt builder with emotion and category
def build_prompt(item, user_question, answer_text):
    name_en = item.get("name_en", "")
    name_bn = item.get("name_bn", "")
    emotion = item.get("emotion", "neutral")
    category = item.get("category", "general")

    prompt = f"""
You are a thoughtful, bilingual assistant. Answer clearly and briefly, like a kind human.

Context:
🔹 Name (EN): {name_en}
🔹 Name (BN): {name_bn}
🔹 Emotion: {emotion}
🔹 Category: {category}
🔹 Matched Answer: {answer_text}

User Question: {user_question}

Instructions:
- Answer in Bengali if question is in Bengali, otherwise in English.
- Keep it short, clear, and relevant.
- Reflect the emotion tag if possible.
"""
    return prompt

# 🤖 Ask model
def ask_model(user_question):
    item, answer_text = find_best_match(user_question)
    if not item:
        return "❌ কোনো মিল পাওয়া যায়নি YAML-এ"

    prompt = build_prompt(item, user_question, answer_text)
    try:
        response = client.text_generation(prompt=prompt, max_new_tokens=150)
        return response.strip()
    except Exception as e:
        return f"❌ Model error: {e}"

# 🔊 Voice preview logic
def get_voice_path(item):
    voice_file = item.get("voice", "")
    voice_path = os.path.join("media", "voice", voice_file)
    return voice_path if os.path.exists(voice_path) else None

# 🎛️ Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🧠 Emotion-aware Bengali QA Chatbot (YAML Powered)")

    with gr.Row():
        user_input = gr.Textbox(label="প্রশ্ন লিখুন (বাংলা বা English)", scale=3)
        submit_btn = gr.Button("উত্তর দিন", scale=1)

    answer_output = gr.Textbox(label="উত্তর", lines=3)
    voice_output = gr.Audio(label="🔊 কণ্ঠস্বর (যদি থাকে)", interactive=False)

    def full_response(user_question):
        item, answer_text = find_best_match(user_question)
        if not item:
            return "❌ কোনো মিল পাওয়া যায়নি YAML-এ", None
        prompt = build_prompt(item, user_question, answer_text)
        response = client.text_generation(prompt=prompt, max_new_tokens=150).strip()
        voice_path = get_voice_path(item)
        return response, voice_path

    submit_btn.click(fn=full_response, inputs=user_input, outputs=[answer_output, voice_output])

demo.launch()