import os
import shutil
from flask import Flask, render_template, request, jsonify
from llama_index.core import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, ChatPromptTemplate, Settings
from llama_index.llms.huggingface import HuggingFaceInferenceAPI
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from huggingface_hub import InferenceClient
from transformers import AutoTokenizer, AutoModel, XLMRobertaXLForMultipleChoice
from deep_translator import GoogleTranslator
import torch

# Ensure HF_TOKEN is set
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("HF_TOKEN environment variable not set.")

repo_id = "facebook/xlm-roberta-xl"
tokenizer = AutoTokenizer.from_pretrained(repo_id)
model = XLMRobertaXLForMultipleChoice.from_pretrained(repo_id)

PERSIST_DIR = "db"
PDF_DIRECTORY = 'data'

# Ensure directories exist
os.makedirs(PDF_DIRECTORY, exist_ok=True)
os.makedirs(PERSIST_DIR, exist_ok=True)
chat_history = []
current_chat_history = []

def data_ingestion_from_directory():
    # Clear previous data by removing the persist directory
    if os.path.exists(PERSIST_DIR):
        shutil.rmtree(PERSIST_DIR)  # Remove the persist directory and all its contents

    # Recreate the persist directory after removal
    os.makedirs(PERSIST_DIR, exist_ok=True)

    # Load new documents from the directory
    new_documents = SimpleDirectoryReader(PDF_DIRECTORY).load_data()

    # Create a new index with the new documents
    index = VectorStoreIndex.from_documents(new_documents)

    # Persist the new index
    index.storage_context.persist(persist_dir=PERSIST_DIR)

def handle_query(query):
    chat_text_qa_msgs = [
        (
            "user",
            """
            You are the Hotel voice chatbot and your name is hotel helper. Your goal is to provide accurate, professional, and helpful answers to user queries based on the hotel's data. Always ensure your responses are clear and concise. Give response within 10-15 words only. You need to give an answer in the same language used by the user.       
            {context_str}
            Question:
            {query_str}
            """
        )
    ]
    text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)

    storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
    index = load_index_from_storage(storage_context)
    context_str = ""
    for past_query, response in reversed(current_chat_history):
        if past_query.strip():
            context_str += f"User asked: '{past_query}'\nBot answered: '{response}'\n"

    query_engine = index.as_query_engine(text_qa_template=text_qa_template, context_str=context_str)
    print(query)
    answer = query_engine.query(query)

    if hasattr(answer, 'response'):
        response = answer.response
    elif isinstance(answer, dict) and 'response' in answer:
        response = answer['response']
    else:
        response = "Sorry, I couldn't find an answer."
    current_chat_history.append((query, response))
    return response

def evaluate_model(prompt, choice0, choice1):
    labels = torch.tensor(0).unsqueeze(0)  # choice0 is correct, batch size 1
    encoding = tokenizer([prompt, prompt], [choice0, choice1], return_tensors="pt", padding=True)
    outputs = model(**{k: v.unsqueeze(0) for k, v in encoding.items()}, labels=labels)  # batch size is 1

    # the linear classifier still needs to be trained
    loss = outputs.loss
    logits = outputs.logits
    return loss, logits

app = Flask(__name__)

# Data ingestion
data_ingestion_from_directory()

# Generate Response
def generate_response(query, language):
    try:
        # Call the handle_query function to get the response
        bot_response = handle_query(query)

        # Map of supported languages
        supported_languages = {
            "hindi": "hi",
            "bengali": "bn",
            "telugu": "te",
            "marathi": "mr",
            "tamil": "ta",
            "gujarati": "gu",
            "kannada": "kn",
            "malayalam": "ml",
            "punjabi": "pa",
            "odia": "or",
            "urdu": "ur",
            "assamese": "as",
            "sanskrit": "sa",
            "arabic": "ar",
            "australian": "en-AU",
            "bangla-india": "bn-IN",
            "chinese": "zh-CN",
            "dutch": "nl",
            "french": "fr",
            "filipino": "tl",
            "greek": "el",
            "indonesian": "id",
            "italian": "it",
            "japanese": "ja",
            "korean": "ko",
            "latin": "la",
            "nepali": "ne",
            "portuguese": "pt",
            "romanian": "ro",
            "russian": "ru",
            "spanish": "es",
            "swedish": "sv",
            "thai": "th",
            "ukrainian": "uk",
            "turkish": "tr"
        }

        # Initialize the translated text
        translated_text = bot_response

        # Translate only if the language is supported and not English
        try:
            if language in supported_languages:
                target_lang = supported_languages[language]
                translated_text = GoogleTranslator(source='en', target=target_lang).translate(bot_response)
                print(translated_text)
            else:
                print(f"Unsupported language: {language}")
        except Exception as e:
            # Handle translation errors
            print(f"Translation error: {e}")
            translated_text = "Sorry, I couldn't translate the response."

        # Append to chat history
        chat_history.append((query, translated_text))
        return translated_text
    except Exception as e:
        return f"Error fetching the response: {str(e)}"

# Route for the homepage
@app.route('/')
def index():
    return render_template('index.html')

# Route to handle chatbot messages
@app.route('/chat', methods=['POST'])
def chat():
    try:
        user_message = request.json.get("message")
        language = request.json.get("language")
        if not user_message:
            return jsonify({"response": "Please say something!"})

        bot_response = generate_response(user_message, language)
        return jsonify({"response": bot_response})
    except Exception as e:
        return jsonify({"response": f"An error occurred: {str(e)}"})

if __name__ == '__main__':
    app.run(debug=True)