import streamlit as st
import os
from langchain_community.llms import HuggingFaceHub
from PyPDF2 import PdfReader
import docx2txt

# Set Hugging Face token (Get yours from https://huggingface.co/settings/tokens)
os.environ["HUGGINGFACEHUB_API_TOKEN"] = st.secrets["HUGGINGFACEHUB_API_TOKEN"]

# Initialize Mistral 7B model
repo_id = "mistralai/Mistral-7B-v0.1"
llm = HuggingFaceHub(
    repo_id=repo_id,
    model_kwargs={"temperature": 0.5, "max_length": 4096}
)

def extract_text(uploaded_file):
    """Extract text from PDF or DOCX files"""
    if uploaded_file.type == "application/pdf":
        pdf_reader = PdfReader(uploaded_file)
        return "\n".join([page.extract_text() for page in pdf_reader.pages])
    elif uploaded_file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
        return docx2txt.process(uploaded_file)
    return ""

def analyze_resume(resume_text, job_desc):
    """Analyze resume using Mistral 7B"""
    prompt = f"""
    Analyze this resume against the job description below. Provide:
    1. Top 5 missing keywords from job description
    2. Three suggested improvements
    3. ATS compatibility score (0-100)
    
    Resume:
    {resume_text[:3000]}
    
    Job Description:
    {job_desc[:2000]}
    
    Format output using markdown with clear sections.
    """
    return llm.invoke(prompt)

def check_ats_compliance(text):
    """Basic ATS formatting checks"""
    forbidden_elements = ["table", "image", "column", "graphic"]
    return all(element not in text.lower() for element in forbidden_elements)

def generate_ats_version(text):
    """Convert text to ATS-friendly format"""
    return text.replace("\n", " ").strip()

# Streamlit UI
st.set_page_config(page_title="ATS Resume Optimizer", layout="wide")

st.title("📄 Free ATS Resume Optimizer")
st.markdown("Upload your resume and job description to get AI-powered optimization tips!")

col1, col2 = st.columns(2)

with col1:
    uploaded_file = st.file_uploader("Upload Resume (PDF/DOCX)", type=["pdf", "docx"])
    job_description = st.text_area("Paste Job Description", height=200)

with col2:
    if uploaded_file and job_description:
        if st.button("Analyze and Optimize"):
            with st.spinner("Processing with Mistral 7B..."):
                resume_text = extract_text(uploaded_file)
                
                # Perform analysis
                analysis = analyze_resume(resume_text, job_description)
                ats_compliant = check_ats_compliance(resume_text)
                optimized_text = generate_ats_version(resume_text)
                
                # Display results
                st.subheader("Analysis Results")
                st.markdown(analysis)
                
                st.subheader("ATS Compliance Check")
                st.metric("Format Status", "Pass ✅" if ats_compliant else "Fail ❌")
                
                # Download optimized version
                st.download_button(
                    label="Download ATS-Optimized Resume",
                    data=optimized_text,
                    file_name="optimized_resume.txt",
                    mime="text/plain"
                )
    else:
        st.info("👈 Upload your resume and paste a job description to get started")

st.markdown("---")
st.caption("Note: This tool uses Mistral 7B AI model for analysis. Always review suggestions before submitting applications.")