import streamlit as st
from transformers import AutoTokenizer, AutoModel
import torch
import spacy
import numpy as np

st.set_page_config(
    page_title="ATS Resume Optimizer",
    page_icon="📄",
    layout="wide",
    initial_sidebar_state="expanded"
)

TEST_DATA = {
    "Data Scientist (Senior)": {
        "resume": "Led a data science team for 5 years. Developed scalable ETL pipelines using Python, SQL, and TensorFlow. Successfully reduced model latency by 20%.",
        "job_desc": "Seeking a Senior Data Scientist who will spearhead complex data pipelines. Must be proficient in Python, SQL, and have experience with TensorFlow. Strong project management skills required."
    },
    "Frontend Developer": {
        "resume": "Built responsive web applications using React, JavaScript, and CSS. Implemented UI designs and optimized performance. Experienced in Git workflows and REST APIs.",
        "job_desc": "Looking for a Frontend Developer skilled in React, JavaScript, CSS, and web performance optimization. Knowledge of REST APIs and Git required."
    },
    "Project Manager": {
        "resume": "Managed multiple projects across software development teams. Skilled in Agile methodologies, stakeholder management, and risk assessment. Led teams to deliver projects on time.",
        "job_desc": "Hiring a Project Manager with experience in Agile frameworks, team coordination, and risk management. Must ensure timely delivery and quality control."
    },
    "Marketing Specialist": {
        "resume": "Executed digital marketing campaigns across social media and email. Analyzed performance metrics and improved ROI by 25%. Skilled in SEO, SEM, and content creation.",
        "job_desc": "Seeking a Marketing Specialist to manage digital campaigns, analyze KPIs, and optimize marketing ROI. SEO, SEM, and content creation skills are essential."
    },
    "Machine Learning Engineer": {
        "resume": "Developed ML models for predictive analytics using Python, PyTorch, and scikit-learn. Implemented end-to-end pipelines and deployed models in production.",
        "job_desc": "Looking for a Machine Learning Engineer experienced in Python, PyTorch, scikit-learn, and model deployment. Must be able to create production-ready ML pipelines."
    }
}

with st.sidebar:
    st.title("📄 ATS Resume Optimizer")
    st.markdown("Select a test dataset to quickly populate resume and job description fields:")
    
    test_choice = st.selectbox("Choose Test Data", ["Custom"] + list(TEST_DATA.keys()))
    
    if test_choice != "Custom":
        resume_input_default = TEST_DATA[test_choice]["resume"]
        job_desc_input_default = TEST_DATA[test_choice]["job_desc"]
    else:
        resume_input_default = ""
        job_desc_input_default = ""

    st.markdown("---")
    st.caption("Built with Streamlit + MiniLM + spaCy")

@st.cache_resource
def load_models():
    """Loads lightweight MiniLM and spaCy models compatible with Hugging Face Spaces."""
    
    model_name = "sentence-transformers/all-MiniLM-L6-v2"

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name, trust_remote_code=True)

    try:
        nlp = spacy.load("en_core_web_sm")
    except OSError:
        st.error("spaCy model 'en_core_web_sm' not found.")
        raise
    return tokenizer, model, nlp

tokenizer, model, nlp = load_models()

def embed(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True)
    with torch.no_grad():
        output = model(**inputs)
    embedding = output.last_hidden_state.mean(dim=1)
    return embedding

def calculate_ats_match(resume_text: str, job_desc_text: str) -> tuple[float, list]:
    
    emb_resume = embed(resume_text)
    emb_job = embed(job_desc_text)

    cosine_score = torch.nn.functional.cosine_similarity(emb_resume, emb_job).cpu().item()
    ats_score = round(float(cosine_score) * 100, 2)
    
    job_doc = nlp(job_desc_text.lower())
    resume_doc = nlp(resume_text.lower())

    job_keywords = {token.text for token in job_doc 
                    if token.pos_ in ["NOUN", "PROPN", "VERB"] and not token.is_stop and len(token.text) > 2}

    resume_keywords = {token.text for token in resume_doc 
                       if token.pos_ in ["NOUN", "PROPN", "VERB"] and not token.is_stop and len(token.text) > 2}

    missing_keywords = list(job_keywords - resume_keywords)
    return ats_score, missing_keywords

st.markdown("## 🚀 Maximize Your ATS Compatibility")
st.markdown("Paste your resume and job description to get a semantic match score and optimization suggestions.")

col1, col2 = st.columns(2)

with col1:
    resume_input = st.text_area(
        "📄 Resume Text", 
        height=350, 
        placeholder="Paste your resume here...",
        value=resume_input_default
    )

with col2:
    job_desc_input = st.text_area(
        "💼 Job Description", 
        height=350, 
        placeholder="Paste the job description here...",
        value=job_desc_input_default
    )

st.markdown("---")
if st.button("🔍 Analyze ATS Match Score"):
    if resume_input.strip() and job_desc_input.strip():
        with st.spinner("Calculating semantic similarity and extracting keywords..."):
            score, missing = calculate_ats_match(resume_input, job_desc_input)
        
        st.metric(label="ATS Match Score", value=f"{score}%")
        
        if score > 75:
            feedback = "🎯 HIGH MATCH. Ready to submit!"
            color = "green"
        elif score > 50:
            feedback = "⚠️ MEDIUM MATCH. Review keyword suggestions."
            color = "orange"
        else:
            feedback = "❌ LOW MATCH. Major revision required."
            color = "red"
        
        st.markdown(f"### Optimization Feedback: <span style='color:{color}'>{feedback}</span>", unsafe_allow_html=True)
        
        st.markdown("### 📝 Missing High-Value Keywords")
        if missing:
            st.info(f"Consider adding these concepts: **{', '.join(missing[:15])}**")
        else:
            st.success("Your resume covers all high-value concepts. Great job! 🎉")
    else:
        st.warning("Please paste text into both the Resume and Job Description boxes.")