import streamlit as st
import PyPDF2
import os
from langchain.llms import HuggingFaceHub
from langchain.prompts import PromptTemplate

# Set up API token
os.environ["HUGGINGFACEHUB_API_TOKEN"] = st.secrets["HF_TOKEN"]

# Initialize LLM
llm = HuggingFaceHub(
    repo_id="mistralai/Mistral-7B-Instruct-v0.3", 
    model_kwargs={"temperature": 0.5}
)

# Function to extract text from 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() + "\n"
    return text

# Streamlit UI
st.title("Cold Email & Cover Letter Generator for Professors")

# User inputs
position_details = st.text_area("Enter details about the research position:", 
                                "e.g., Professor's name, research focus, university, lab details, etc.")
resume_file = st.file_uploader("Upload your CV/Resume (PDF)", type=["pdf"])

if st.button("Generate Cold Email & Cover Letter"):
    if position_details and resume_file:
        # Extract text from the uploaded PDF
        resume_text = extract_text_from_pdf(resume_file)

        # Define prompt for cold email
        email_prompt = PromptTemplate.from_template(
            """
            Based on the following details of a research position and a student's CV, generate a professional cold email:
            
            Research Position Details:
            {position}
            
            Student's CV:
            {resume}
            
            The email should be formal, concise, and engaging, expressing interest in the position while highlighting relevant skills.
            """
        )
        
        email_content = llm(email_prompt.format(position=position_details, resume=resume_text))
        
        # Define prompt for cover letter
        cover_prompt = PromptTemplate.from_template(
            """
            Generate a professional cover letter based on the following research position details and a student's CV:
            
            Research Position Details:
            {position}
            
            Student's CV:
            {resume}
            
            The cover letter should be well-structured, highlighting the student's background, relevant skills, research interests, and motivation for applying.
            """
        )
        
        cover_content = llm(cover_prompt.format(position=position_details, resume=resume_text))
        
        # Display results
        st.subheader("Generated Cold Email")
        st.write(email_content)
        
        st.subheader("Generated Cover Letter")
        st.write(cover_content)
    else:
        st.warning("Please provide both the position details and upload a CV.")