import gradio as gr
import logging
import PyPDF2

from config import SHARE_GRADIO_WITH_PUBLIC_URL
from chains import qa_chain, summarization_chain

logger = logging.getLogger(__name__)

# Translation dictionary
TRANSLATIONS = {
    "en": {
        "title": "# 📚 Study Buddy: AI Learning Assistant",
        "subtitle": "## 🤖 A smart, user-friendly chatbot for students!",
        "summary_subtitle": "## 📄 Upload Notes for Summarization",
        "chat_input_label": "Type your question here:",
        "chat_placeholder": "e.g., Explain Newton's laws",
        "chat_button_label": "Get Answer",
        "summary_button_label": "Summarize Notes",
        "upload_file_label": "Upload .txt or .pdf file",
        "summary_output_label": "Summary",
        "language_label": "Language / Langue",
        "ai_response_label": "AI Response"
    },
    "fr": {
        "title": "# 📚 Study Buddy: Assistant d'apprentissage IA",
        "subtitle": "## 🤖 Un chatbot intelligent et convivial pour les étudiants!",
        "summary_subtitle": "## 📄 Subir notas para resumir",
        "chat_input_label": "Tapez votre question ici:",
        "chat_placeholder": "ex : Expliquez les lois de Newton",
        "chat_button_label": "Obtenir une réponse",
        "summary_button_label": "Résumer les notes",
        "upload_file_label": "Téléchargez un fichier .txt ou .pdf",
        "summary_output_label": "Résumé",
        "language_label": "Langue / Language",
        "ai_response_label": "Réponse de l'IA"
    }
}

# Function to process user queries
def chatbot_response(user_input, lang):
    try:
        response_output = qa_chain.invoke({"question": user_input})
        response = response_output.content
        logger.info("chatbot_response completed")
        print("> chatbot_response completed")
        return response
    except Exception as e:
        msg = f"Error : {e}"
        logger.exception(msg)
        print(msg)
        return TRANSLATIONS[lang].get("error_message", "Sorry, an error occurred while processing your request.")

# Function to summarize notes
def summarize_pdf(pdf, lang):
    try:
        with open(pdf, "rb") as file:
            reader = PyPDF2.PdfReader(file)
            page = reader.pages[0]  # Get the first page
            text = page.extract_text()
            print(text)
        summary = summarize_text(text, lang)
        logger.info("summarize_pdf completed")
        print("> summarize_pdf completed")
        return summary
    except Exception as e:
        msg = f"Error : {e}"
        logger.exception(msg)
        print(msg)
        return TRANSLATIONS[lang].get("error_message", "Sorry, an error occurred while summarizing your notes.")
    
# Function to summarize notes
def summarize_text(text, lang):
    try:
        summary_output = summarization_chain.invoke({"document_text": text})
        print(summary_output)
        summary = summary_output.content
        logger.info("summarize_text completed")
        print("> summarize_text completed")
        return summary
    except Exception as e:
        msg = f"Error : {e}"
        logger.exception(msg)
        print(msg)
        return TRANSLATIONS[lang].get("error_message", "Sorry, an error occurred while summarizing your notes.")

# Function to update UI labels dynamically
def update_language(lang):
    return (
        TRANSLATIONS[lang]["title"],
        TRANSLATIONS[lang]["subtitle"],
        TRANSLATIONS[lang]["chat_input_label"],
        TRANSLATIONS[lang]["chat_placeholder"],
        TRANSLATIONS[lang]["chat_button_label"],
        TRANSLATIONS[lang]["upload_file_label"],
        TRANSLATIONS[lang]["summary_button_label"],
        TRANSLATIONS[lang]["summary_output_label"],
        TRANSLATIONS[lang]["ai_response_label"]
    )

# Gradio UI
def create_interface():
    with gr.Blocks(css="body { font-family: sans-serif; background-color: #f9f9f9; }") as study_buddy:
        
        # Default to English
        lang = "en"

        title = gr.Markdown(f"{TRANSLATIONS[lang]['title']}")
        
        with gr.Row():
            with gr.Column():
                gr.Markdown("", height=4) 
                
                language = gr.Radio(
                    choices=["en", "fr"],
                    value=lang,
                    label=TRANSLATIONS[lang]["language_label"]
                )
                
                gr.Markdown("", height=4)
                               
                subtitle = gr.Markdown(f"{TRANSLATIONS[lang]['subtitle']}")
                
                chat_input = gr.Textbox(
                    # label=TRANSLATIONS[lang]["chat_input_label"]
                    label= "Type your question here: / Tapez votre question ici:",
                    lines=4,
                    placeholder=TRANSLATIONS[lang]["chat_placeholder"]
                )

            with gr.Column():
                gr.Markdown("", height=4)
                summary_subtitle = gr.Markdown(f"{TRANSLATIONS[lang]['summary_subtitle']}")
                file_input = gr.File(label="File / Fichier",file_types=[".pdf", ".txt"])
                file_label = gr.Markdown(TRANSLATIONS[lang]["upload_file_label"])  # Separate label
                
        with gr.Row():
            with gr.Column():
                
                chat_button = gr.Button(TRANSLATIONS[lang]["chat_button_label"], variant="primary")
                chat_output = gr.Textbox(
                    # label=TRANSLATIONS[lang]["ai_response_label"], 
                    label = "AI Response / Réponse de l'IA",
                    lines=5, interactive=True
                )

                # Bind chatbot response function
                chat_button.click(
                    chatbot_response, 
                    inputs=[chat_input, language], 
                    outputs=chat_output
                )

            with gr.Column():
            
                summary_button = gr.Button(TRANSLATIONS[lang]["summary_button_label"], variant="primary")
                summary_output = gr.Textbox(
                    # label=TRANSLATIONS[lang]["summary_output_label"],
                    label="Summary / Résumé",
                    lines=5, 
                    interactive=True
                )

                # Bind summarization function
                summary_button.click(
                    summarize_pdf, 
                    inputs=[file_input, language], 
                    outputs=summary_output
                )

        # Update labels dynamically when the language changes
        def update_labels(lang):
            return (
                TRANSLATIONS[lang]["title"],
                TRANSLATIONS[lang]["subtitle"],
                TRANSLATIONS[lang]["summary_subtitle"],
                TRANSLATIONS[lang]["chat_input_label"],
                TRANSLATIONS[lang]["chat_placeholder"],
                TRANSLATIONS[lang]["chat_button_label"],
                TRANSLATIONS[lang]["summary_button_label"],
                TRANSLATIONS[lang]["summary_output_label"],
                TRANSLATIONS[lang]["ai_response_label"],
                TRANSLATIONS[lang]["upload_file_label"]
            )

        language.change(
            update_labels, 
            inputs=[language], 
            outputs=[
                title, subtitle, summary_subtitle, 
                chat_input, chat_input, 
                chat_button, 
                summary_button, summary_output, chat_output, file_label  # Update file label separately
            ]
        )

    return study_buddy

if __name__ == "__main__":
    study_buddy = create_interface()
    #study_buddy.launch(share=SHARE_GRADIO_WITH_PUBLIC_URL)
    study_buddy.launch()