import os
import streamlit as st
from PyPDF2 import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import FAISS
from langchain.chains.question_answering import load_qa_chain
from langchain.llms import OpenAI
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
from io import BytesIO
from deep_translator import GoogleTranslator

# Set your OpenAI API key here
os.environ["OPENAI_API_KEY"] = os.getenv("OpenAIKey")

st.title("LegalAIDe")

if 'translated_answer' not in st.session_state:
    st.session_state.translated_answer = ""

uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")

if uploaded_file is not None:
    pdf_reader = PdfReader(uploaded_file)
    
    raw_text = ''
    for page in pdf_reader.pages:
        text = page.extract_text()
        if text:
            raw_text += text
    
    text_splitter = RecursiveCharacterTextSplitter(
        separators=["\n\n", "\n", ".", "!", "?"],
        chunk_size=1000,
        chunk_overlap=200,
        length_function=len
    )

    texts = text_splitter.split_text(raw_text)
    embeddings = OpenAIEmbeddings()
    document_search = FAISS.from_texts(texts, embeddings)
    
    query = st.text_input("Enter your question:")

    if st.button("Get Answer"):
        docs = document_search.similarity_search(query)
        chain = load_qa_chain(OpenAI(), chain_type="stuff")

        # Add context about your role
        context = "You are a lawyer and need assistance with legal questions."

        # Use a more explicit prompt for better context
        prompt = f"Context: {context}\n\nQuestion: {query}\n\nPlease provide a detailed answer based on the given documents."

        answer = chain.run(input_documents=docs, question=prompt)
        st.write("Answer:", answer)

        st.session_state.answer = answer

    # Language selection and translation
    translation_language = st.text_input("Enter translation language:")  # Corrected to use st.text_input

if st.button("Translate Answer"):
    if translation_language:  # Check if a language was entered
        try:
            translated_answer = GoogleTranslator(source='auto', target=translation_language).translate(st.session_state.answer)
            st.session_state.translated_answer = translated_answer
        except Exception as e:
            st.error(f"An error occurred: {e}")
    else:
        st.error("Please enter a valid language.")

# Display translated answer if available
if st.session_state.get('translated_answer'):
    st.write("Translated Answer:", st.session_state.translated_answer)

    if st.button("Generate PDF"):
        # Generate PDF
        pdf_buffer = BytesIO()
        c = canvas.Canvas(pdf_buffer, pagesize=letter)
        c.drawString(100, 750, "Question:")
        c.drawString(100, 735, query)
        c.drawString(100, 715, "Answer:")
        text_lines = st.session_state.answer.split('\n')
        y = 700
        for line in text_lines:
            if y < 50:
                c.showPage()
                y = 750
            c.drawString(100, y, line)
            y -= 15

        if st.session_state.translated_answer:
            c.drawString(100, y, "Translated Answer:")
            y -= 15
            text_lines = st.session_state.translated_answer.split('\n')
            for line in text_lines:
                if y < 50:
                    c.showPage()
                    y = 750
                c.drawString(100, y, line)
                y -= 15

        c.save()

        pdf_buffer.seek(0)
        st.download_button(
            label="Download Answer as PDF",
            data=pdf_buffer,
            file_name="answer.pdf",
            mime="application/pdf"
        )