import streamlit as st
import os
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
from langchain.prompts import PromptTemplate, ChatPromptTemplate
from langchain_community.document_loaders import UnstructuredPDFLoader
import tempfile

# Set up Hugging Face API token
hk = os.getenv('hf')
os.environ['HUGGINGFACEHUB_API_TOKEN'] = hk
os.environ['HF_TOKEN'] = hk

# Access the LLM
llm_skeleton = HuggingFaceEndpoint(
    repo_id='meta-llama/Llama-3.2-3B-Instruct',
    provider='novita',
    temperature=0.7,
    max_new_tokens=150,
    task='conversational'
)
llm = ChatHuggingFace(
    llm=llm_skeleton,
    repo_id='meta-llama/Llama-3.2-3B-Instruct',
    provider='novita',
    temperature=0.7,
    max_new_tokens=150,
    task='conversational'
)

# App Layout
st.title("📄 Resume & Job Description Extractor")

# Upload Resume PDF
resume_file = st.file_uploader("Upload Resume (PDF)", type=["pdf"])

# Upload Job Description PDF (optional)
jd_file = st.file_uploader("Upload Job Description (PDF)", type=["pdf"])

# Or Input Job Description Text
jd_text = st.text_area("Or paste Job Description text here")

if st.button("Extract Data"):
    if resume_file:
        # Save the uploaded resume to a temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_resume:
            tmp_resume.write(resume_file.read())
            resume_path = tmp_resume.name
        
        # Load resume with UnstructuredPDFLoader
        loader = UnstructuredPDFLoader(resume_path)
        resume_text = loader.load()[0].page_content
        
        # Extract resume data using LLM
        resume_prompt = f"""
        Extract the following from the resume:
        1. Name
        2. Education
        3. Experience
        4. Skills
        5. Project Names and Results
        
        Resume:
        {resume_text}
        """
        resume_data = llm.invoke(resume_prompt)
        st.subheader("Extracted Resume Data")
        st.write(resume_data)

        # Clean up temp file
        os.unlink(resume_path)

    if jd_file or jd_text:
        if jd_file:
            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_jd:
                tmp_jd.write(jd_file.read())
                jd_path = tmp_jd.name

            loader = UnstructuredPDFLoader(jd_path)
            jd_text_extracted = loader.load()[0].page_content
            os.unlink(jd_path)  # Clean up temp file
        else:
            jd_text_extracted = jd_text

        # Extract JD data using LLM
        jd_prompt = f"""
        Extract the following from the job description:
        1. Job ID
        2. Company Name
        3. Role
        4. Experience Required
        5. Skills Required
        6. Education Required
        7. Location
        
        Job Description:
        {jd_text_extracted}
        """
        jd_data = llm.invoke(jd_prompt)
        st.subheader("Extracted Job Description Data")
        st.write(jd_data)