import streamlit as st
from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.schema.runnable import RunnableLambda
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from PyPDF2 import PdfReader
import os

# Streamlit UI setup
st.set_page_config(page_title="AI Job Interview Prep Tool")
st.title("AI-Powered Job Interview Preparation")

# It's good practice to load these from a .env file, but for demonstration purposes, we are setting them directly
os.environ["OPENAI_API_KEY"] = os.environ.get("s1")

# User Inputs
company = st.text_input("Company Name")
profile = st.text_input("Profile Name")
jd = st.text_area("Job Description")
resume_file = st.file_uploader("Upload Resume (Optional)", type=["pdf"])

# Extract text from PDF if resume is uploaded
resume_text = ""
if resume_file is not None:
    pdf = PdfReader(resume_file)
    for page in pdf.pages:
        resume_text += page.extract_text()

# Submit button
submit = st.button("Generate Interview Questions")

if submit:
    if not company or not profile or not jd:
        st.warning("Please fill in all mandatory fields.")
    else:
        # Set up OpenAI LLM
        llm = ChatOpenAI(temperature=0.7, model_name="gpt-4")

        # Prompt Template
        base_prompt = PromptTemplate(
            input_variables=["company", "profile", "jd", "resume"],
            template="""
You are an expert HR and Technical Interviewer.
Given the following details:
Company: {company}
Profile: {profile}
Job Description: {jd}
Resume (if provided): {resume}

Generate a structured set of interview questions:
- Round 1: Technical Interview (7-10 questions including 2 hands-on tasks)
- Round 2: Technical + Case Study (7-10 questions including at least 1 scenario or case-study based question)
- Round 3: Business and HR Round (7-10 questions focused on communication, decision making, and company alignment)

Present the questions grouped by round in clear bullet format.
            """
        )

        chain = LLMChain(llm=llm, prompt=base_prompt)
        response = chain.run({
            "company": company,
            "profile": profile,
            "jd": jd,
            "resume": resume_text or "Not Provided"
        })

        st.subheader("Suggested Interview Questions")
        st.write(response)