# src/telegram_bot.py
import os
import telebot
from langchain.chains import RetrievalQA
from langchain_community.llms import HuggingFaceEndpoint
from retriever import get_retriever

BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
HF_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")

bot = telebot.TeleBot(BOT_TOKEN)

retriever = get_retriever()
llm = HuggingFaceEndpoint(repo_id="tiiuae/falcon-7b-instruct", huggingfacehub_api_token=HF_TOKEN)
qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever)

user_lang_pref = {}

@bot.message_handler(commands=['start'])
def send_welcome(message):
    bot.reply_to(message, "Welcome to KnowYourRight Bot! Reply with '1' for English or '2' for Pidgin.")
    user_lang_pref[message.chat.id] = None

@bot.message_handler(func=lambda m: user_lang_pref.get(m.chat.id) is None)
def set_language(message):
    if message.text.strip() == "1":
        user_lang_pref[message.chat.id] = "english"
        bot.reply_to(message, "Okay! Ask me your question in English.")
    elif message.text.strip() == "2":
        user_lang_pref[message.chat.id] = "pidgin"
        bot.reply_to(message, "Okay! Ask me your question in Pidgin.")
    else:
        bot.reply_to(message, "Please type 1 for English or 2 for Pidgin.")

@bot.message_handler(func=lambda m: user_lang_pref.get(m.chat.id) is not None)
def handle_query(message):
    lang_choice = user_lang_pref[message.chat.id]
    query = message.text
    if lang_choice == "pidgin":
        query = f"Answer in Nigerian Pidgin: {query}"
    result = qa_chain.run(query)
    bot.reply_to(message, result)

if __name__ == "__main__":
    bot.polling()