import streamlit as st
import google.generativeai as genai
from docx import Document
import PyPDF2

# 🔹 Streamlit UI
st.set_page_config(page_title="ATS Resume Optimizer", layout="wide")
st.title("📄 ATS-Compliant Resume Optimizer")

# 🔹 Configure Google Gemini API
GOOGLE_API_KEY = st.secrets["GEMINI_API_KEY"]
genai.configure(api_key=GOOGLE_API_KEY)
model = genai.GenerativeModel("gemini-1.5-pro")

# 🔹 Resume Upload
st.subheader("Upload Your Resume")
uploaded_file = st.file_uploader("Upload Resume (PDF or DOCX)", type=["pdf", "docx"])

# 🔹 Job Description Input
st.subheader("Paste Job Description")
job_description = st.text_area("Enter Job Description", height=150)

def extract_text_from_resume(uploaded_file):
    """Extract raw text from resume (PDF/DOCX)"""
    if not uploaded_file:
        return ""
    try:
        if uploaded_file.type == "application/pdf":
            pdf_reader = PyPDF2.PdfReader(uploaded_file)
            return "\n".join(page.extract_text() for page in pdf_reader.pages if page.extract_text())
        elif uploaded_file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
            doc = Document(uploaded_file)
            return "\n".join(paragraph.text for paragraph in doc.paragraphs if paragraph.text)
    except Exception as e:
        st.error(f"❌ Error processing resume: {e}")
        return ""
    return ""

# 🔹 Process Resume & JD
def generate_optimized_resume(resume_text, job_description):
    """
    Uses LLM to enhance the resume for ATS optimization.
    - Extracts JD keywords & integrates them into the resume
    - Optimizes Professional Summary, Skills, Work Experience, and Projects
    - Ensures ATS compliance with structured formatting
    """

    prompt = f"""
    You are an AI Career Assistant specializing in ATS-optimized resumes.

    🔹 **TASK:** Modify the user's resume to align with the provided job description (JD).  
    🔹 **GOAL:** Ensure **keyword integration**, structured formatting, and ATS compliance while maintaining originality.

    **📄 Job Description (JD):**
    {job_description}

    **📌 Original Resume:**
    {resume_text}

    ---  
    **🔹 PROCESS:**  
    1️⃣ **Extract Keywords from JD** → Identify essential skills, technologies, and job-related terms.  
    2️⃣ **Enhance Key Resume Sections:**  
       - **Professional Summary:** Rewrite using extracted JD keywords.  
       - **Skills Section:** Categorize into 6 key areas; list as comma-separated values.  
       - **Work Experience:** Naturally integrate JD keywords & **quantify achievements** (e.g., "Improved efficiency by 30%").  
       - **Projects:** Highlight relevant projects & **technologies/tools used**.  
    3️⃣ **Ensure Structured, ATS-Friendly Formatting** → Use clean, readable formatting without tables/graphics.  
    4️⃣ **Adaptability Statement (if needed)** → Add: _"Although my resume wasn’t selected earlier, I am confident in learning & delivering results."_  

    ---  
    **🚀 OUTPUT FORMAT (Plain Text):**  
    Return the **fully optimized resume** in a clean, readable format, ensuring ATS compliance.

    """
    
    # Call Gemini API
    model = genai.GenerativeModel("gemini-1.5-pro")  
    response = model.generate_content(prompt)
    
    return response.text if response else "❌ Error generating resume."


# 🔹 Enhance Resume Button
if st.button("🚀 Enhance Resume"):
    if uploaded_file and job_description.strip():
        with st.spinner("Optimizing your resume..."):
            resume_text = extract_text_from_resume(uploaded_file)
            enhanced_resume = generate_optimized_resume(resume_text, job_description)
            
            # ✅ Display Enhanced Resume
            st.subheader("📄 Optimized Resume")
            st.text_area("Your ATS-Optimized Resume:", enhanced_resume, height=400)
            
            # ✅ Download Option
            st.download_button(
                label="📥 Download Optimized Resume",
                data=enhanced_resume,
                file_name="optimized_resume.txt",
                mime="text/plain"
            )
    else:
        st.warning("⚠️ Please upload a resume and provide a job description.")