import streamlit as st
import os
from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    Settings,
)
from llama_index.core import PromptTemplate
from llama_index.llms.gemini import Gemini
from llama_index.embeddings.gemini import GeminiEmbedding
import logging
import google.generativeai as genai
from dotenv import load_dotenv
from pathlib import Path

load_dotenv()

# Set logging level
logging.basicConfig(level=logging.INFO)


# Configure Gemini Pro
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))

model_gemini_pro_vision = "gemini-pro-vision"

# Configure Gemini models
Settings.llm = Gemini(model=model_gemini_pro_vision,
                        api_key=os.getenv("GOOGLE_API_KEY"))
Settings.embed_model = GeminiEmbedding(
    model_name="models/embedding-001",
                        api_key=os.getenv("GOOGLE_API_KEY")
    )

def load_and_index_pdf(pdf_path):
    """Loads and index the pdf.
    
    Args : 
    pdf_path (str) : The path to the pdf file
    
    Returns : 
    index (llama_index.core.VectorStoreIndex): The vector index
    """
    try:
        logging.info(f"Loading PDF document from: {pdf_path}")
        documents = SimpleDirectoryReader(input_files=[pdf_path]).load_data()
        if documents:
            logging.info("Creating vector store index")
            index = VectorStoreIndex.from_documents(documents)
            return index
        else:
            logging.warning("No documents found in the PDF")
            return None
    except Exception as e:
        logging.error(f"Error loading and indexing PDF: {e}")
        return None
    
def translate_text(french_text, index):
    """Translates french text to Yipunu.
    
    Args : 
    french_text (str): The french text to translate.
    index (llama_index.core.VectorStoreIndex): The vector index.
    
    Returns:
     (str): The yipunu translation or an error message.
    """

    try:
        logging.info(f"Initiating translation of: {french_text}")

        template = (
            "Tu es un excellent traducteur du français vers le yipunu. Tu traduis le texte sans donner d'explication. "
            "Texte: {french_text} "
            "Traduction:"
        )

        prompt_template = PromptTemplate(template)

        query_engine = index.as_query_engine(
            text_qa_template=prompt_template
        )
        response = query_engine.query(french_text)
        logging.info(f"Translation Result: {response.response}")
        return response.response
    except Exception as e:
        logging.error(f"Error during translation: {e}")
        return f"Error during translation: {str(e)}"
        

def main():
    """Main function for streamlit app."""
    
    st.title("French to Yipunu Translation App")

    # PDF File Upload
    uploaded_file = st.file_uploader("Upload a PDF file containing the Punu grammar:", type="pdf")

    if uploaded_file is not None:
            # Save file to a temporary location
        temp_file_path = Path("temp_file.pdf")
        with open(temp_file_path, "wb") as f:
            f.write(uploaded_file.read())

        index = load_and_index_pdf(str(temp_file_path))
        if index:
            french_text = st.text_area("Enter French Text:", "Ni vosi yipunu")
            if st.button("Translate"):
                 translation = translate_text(french_text, index)
                 st.success(f"Yipunu Translation: {translation}")

        # Clean up temp files
        os.remove(temp_file_path)
    else:
        st.info("Please upload a pdf containing the punu grammar.")


if __name__ == "__main__":
    main()