import streamlit as st
import tempfile
import os
import random
import json
from pathlib import Path
from PyPDF2 import PdfReader
from openai import OpenAI

# Import our CrewAI resume optimization module
from crewai_resume_optimization import optimize_resume, generate_interview_questions, extract_job_description

# ---------------------------
# Helper Function: Extract text from PDF (for resume)
# ---------------------------
def extract_pdf_text(uploaded_file):
    # Ensure file size is less than 10MB
    uploaded_file.seek(0, os.SEEK_END)
    if uploaded_file.tell() > 10 * 1024 * 1024:
        st.error("File exceeds 10MB limit.")
        return ""
    uploaded_file.seek(0)
    reader = PdfReader(uploaded_file)
    text = ""
    for page in reader.pages:
        page_text = page.extract_text()
        if page_text:
            text += page_text + "\n"
    return text

# ---------------------------
# OpenAI Helper for Keyword Extraction
# ---------------------------
def extract_keywords_from_text(text):
    prompt = f"Extract the most important keywords from the following job input, separated by commas:\n\n{text}"
    messages = [
        {"role": "system", "content": "You are a job analysis assistant."},
        {"role": "user", "content": prompt}
    ]
    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    completion = client.chat.completions.create(model="gpt-4o-mini", messages=messages)
    return completion.choices[0].message.content.strip()

# ---------------------------
# Sidebar: File Upload & Input Mode
# ---------------------------
st.sidebar.title("Job & Resume Setup")

# Resume upload (PDF only)
uploaded_resume = st.sidebar.file_uploader("Upload your resume (PDF, max 10MB)", type="pdf")

# Job input mode selection: "Job Description" vs "Job URL"
mode = st.sidebar.radio("Select Job Input Mode", ("Job Description", "Job URL"))
job_input = ""
if mode == "Job Description":
    job_input = st.sidebar.text_area("Paste the job description here:")
elif mode == "Job URL":
    job_input = st.sidebar.text_input("Enter the job URL:")

# Button to extract keywords from job input.
if st.sidebar.button("Extract Keywords"):
    if job_input:
        # For simplicity, we use the same function regardless of mode.
        st.session_state.job_keywords = extract_keywords_from_text(job_input)
        st.sidebar.success("Keywords extracted!")
    else:
        st.sidebar.error("Please provide job input.")

# Button to trigger optimization.
if st.sidebar.button("Optimize Resume"):
    st.session_state.optimize_trigger = True

# ---------------------------
# Session State Initialization
# ---------------------------
if "resume_text" not in st.session_state:
    st.session_state.resume_text = None
if "job_keywords" not in st.session_state:
    st.session_state.job_keywords = None
if "optimized_resume" not in st.session_state:
    st.session_state.optimized_resume = None
if "interview_questions" not in st.session_state:
    st.session_state.interview_questions = None

# ---------------------------
# Process Resume Upload
# ---------------------------
if uploaded_resume is not None:
    st.session_state.resume_text = extract_pdf_text(uploaded_resume)
    if st.session_state.resume_text:
        st.sidebar.success("Resume uploaded and processed successfully!")
    else:
        st.sidebar.error("Failed to extract text from the resume.")

# ---------------------------
# Main Area: Display Results
# ---------------------------
st.title("Optimized Resume & Interview Preparation")

if st.session_state.resume_text is None:
    st.info("Please upload your resume from the sidebar to begin.")
else:
    # Display job keywords if available.
    if st.session_state.job_keywords:
        st.subheader("Extracted Job Keywords")
        st.write(st.session_state.job_keywords)
    else:
        st.info("Please provide job input and extract keywords from the sidebar.")
    
    # Trigger optimization if requested.
    if st.session_state.get("optimize_trigger", False):
        if st.session_state.resume_text and st.session_state.job_keywords:
            with st.spinner("Optimizing resume..."):
                st.session_state.optimized_resume = optimize_resume(st.session_state.resume_text, st.session_state.job_keywords)
            with st.spinner("Generating interview questions..."):
                st.session_state.interview_questions = generate_interview_questions(st.session_state.optimized_resume, st.session_state.job_keywords)
            st.success("Resume optimized and interview questions generated!")
            st.session_state.optimize_trigger = False
        else:
            st.error("Please ensure you have uploaded your resume and extracted job keywords.")
    
    # Display the optimized resume and interview questions if available.
    if st.session_state.optimized_resume:
        st.header("Optimized Resume")
        st.markdown(st.session_state.optimized_resume)
        st.download_button("Download Optimized Resume (md)",
                           st.session_state.optimized_resume,
                           file_name="tailored_resume.md",
                           mime="text/markdown")
    if st.session_state.interview_questions:
        st.header("Interview Questions")
        st.markdown(st.session_state.interview_questions)
        st.download_button("Download Interview Questions (md)",
                           st.session_state.interview_questions,
                           file_name="interview_materials.md",
                           mime="text/markdown")
        
st.info("You can open the downloaded markdown (.md) files with Notepad or any text editor.")