import os
import openai
import streamlit as st
import PyPDF2

# Set up OpenAI API key
openai.api_key = os.getenv("OPENAI_KEY")

# Streamlit app
def main():
    st.title("EB2-NIW Petition Draft Generator")
    st.write("Upload a research paper PDF and get a drafted petition addressing the three prongs for EB2-NIW.")

    # File upload
    uploaded_file = st.file_uploader("Upload PDF", type="pdf")

    if uploaded_file is not None:
        # Extract text from PDF
        pdf_text = extract_text_from_pdf(uploaded_file)

        # Button to extract and show text from PDF
        if st.button("Extract Text"):
            if pdf_text:
                st.subheader("Extracted Text:")
                st.write(pdf_text)
            else:
                st.error("Failed to extract text from the PDF.")

        # Button to generate petition draft
        if st.button("Generate Petition Draft"):
            if pdf_text.strip() == "":
                st.error("No text extracted from the PDF. Please check the file and try again.")
            else:
                # Generate the petition draft based on the extracted text
                petition_draft = generate_petition_draft(pdf_text)
                st.success("Petition Draft:")
                st.write(petition_draft)


# Function to extract text from a PDF
def extract_text_from_pdf(pdf_file):
    pdf_reader = PyPDF2.PdfReader(pdf_file)
    text = ""
    for page in pdf_reader.pages:
        text += page.extract_text()
    return text

# Function to generate EB2-NIW petition draft using GPT-3.5
def generate_petition_draft(research_text):
    # Prompt for GPT-3.5
    prompt = (
        "A researcher has asked you to draft an EB2-NIW petition. "
        "Based on the following research paper text, draft a 2-page petition that proves the following three prongs:\n\n"
        "Prong 1: The applicant's proposed endeavor has substantial merit and national importance.\n"
        "Prong 2: The applicant is well-positioned to advance the proposed endeavor.\n"
        "Prong 3: On balance, it would be beneficial to the United States to waive the requirements of the PERM labor certification.\n\n"
        "Research paper text:\n\n" + research_text
    )

    # Generate response from GPT-3.5
    response = openai.ChatCompletion.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "You are an experienced immigration lawyer."},
            {"role": "user", "content": prompt}
        ],
        max_tokens=10000,
        temperature=0
    )

    petition_draft = response.choices[0].message['content'].strip()
    return petition_draft

if __name__ == "__main__":
    main()