import os

# ---------------- SAFE CACHE ----------------
os.environ["HF_HOME"] = "/tmp/hf"
os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf"
os.environ["TORCH_HOME"] = "/tmp/torch"

# ---------------- IMPORTS ----------------
import streamlit as st
import pdfplumber
import re
import pandas as pd
import spacy
import json

from keybert import KeyBERT
from sentence_transformers import SentenceTransformer, util
from streamlit_lottie import st_lottie

# =========================================================
# PAGE CONFIG
# =========================================================
st.set_page_config(
    page_title="AI Resume ATS Matcher",
    page_icon="📄",
    layout="wide"
)

# =========================================================
# CUSTOM CSS (LIGHT MODERN UI)
# =========================================================
st.markdown("""
<style>

/* Main Background */
.stApp {
    background-color: #f8fafc;
}

/* Header */
.main-title {
    font-size: 42px;
    font-weight: 800;
    color: #111827;
    text-align: center;
    margin-bottom: 5px;
}

.sub-title {
    font-size: 18px;
    color: #6b7280;
    text-align: center;
    margin-bottom: 30px;
}

/* Upload Box */
.upload-box {
    background: white;
    padding: 25px;
    border-radius: 18px;
    border: 1px solid #e5e7eb;
    box-shadow: 0px 2px 10px rgba(0,0,0,0.04);
    margin-bottom: 20px;
}

/* Candidate Card */
.candidate-card {
    background: white;
    padding: 20px;
    border-radius: 18px;
    border: 1px solid #e5e7eb;
    margin-bottom: 20px;
    box-shadow: 0px 2px 10px rgba(0,0,0,0.05);
}

/* Skill Tags */
.skill-tag {
    display: inline-block;
    background: #eff6ff;
    color: #2563eb;
    padding: 6px 12px;
    border-radius: 30px;
    margin: 4px;
    font-size: 13px;
    font-weight: 500;
}

/* Buttons */
.stButton > button {
    width: 100%;
    height: 55px;
    border-radius: 14px;
    background-color: #2563eb;
    color: white;
    font-size: 18px;
    font-weight: 700;
    border: none;
}

.stButton > button:hover {
    background-color: #1d4ed8;
    color: white;
}

/* Dataframe */
[data-testid="stDataFrame"] {
    border-radius: 15px;
    overflow: hidden;
}

/* Hide Streamlit Footer */
footer {
    visibility: hidden;
}

</style>
""", unsafe_allow_html=True)

# =========================================================
# LOAD MODELS
# =========================================================
@st.cache_resource
def load_models():

    model = SentenceTransformer("all-MiniLM-L6-v2")

    nlp = spacy.load("en_core_web_sm")

    kw_model = KeyBERT(model=model)

    return model, nlp, kw_model


model, nlp, kw_model = load_models()

# =========================================================
# LOAD ANIMATION
# =========================================================
def load_lottiefile(filepath):

    try:
        with open(filepath, "r") as f:
            return json.load(f)

    except:
        return None


lottie_animation = load_lottiefile("Animation.json")

# =========================================================
# HEADER
# =========================================================
if lottie_animation:
    st_lottie(lottie_animation, height=180)

st.markdown(
    '<div class="main-title">🤖 AI Resume ATS </div>',
    unsafe_allow_html=True
)

st.markdown(
    '<div class="sub-title">Upload multiple resumes and instantly rank the best candidates</div>',
    unsafe_allow_html=True
)

# =========================================================
# HELPERS
# =========================================================
def extract_text_from_pdf(pdf_file):

    text = ""

    with pdfplumber.open(pdf_file) as pdf:

        for page in pdf.pages:

            page_text = page.extract_text()

            if page_text:
                text += page_text + "\n"

    return text.strip()


def clean_text(text):

    return re.sub(r"\s+", " ", text).strip()


def extract_skills_auto(text, top_n=15):

    text = text[:3000]

    keywords = kw_model.extract_keywords(
        text,
        keyphrase_ngram_range=(1, 3),
        stop_words='english',
        top_n=top_n
    )

    return [kw[0] for kw in keywords]


def calculate_similarity(text1, text2):

    emb1 = model.encode(text1, convert_to_tensor=True)

    emb2 = model.encode(text2, convert_to_tensor=True)

    score = util.pytorch_cos_sim(
        emb1,
        emb2
    ).item()

    return round(score * 100, 2)

# =========================================================
# INPUT SECTION
# =========================================================
#st.markdown('<div class="upload-box">', unsafe_allow_html=True)

st.markdown("### 📄 Upload Resumes")

resume_files = st.file_uploader(
    "Upload Multiple PDF Resumes",
    type=["pdf"],
    accept_multiple_files=True
)

st.markdown("### 📝 Paste Job Description")

job_description = st.text_area(
    "Enter Job Description",
    height=220,
    placeholder="Paste the job description here..."
)

st.markdown('</div>', unsafe_allow_html=True)

# =========================================================
# MATCH BUTTON
# =========================================================
if st.button("🚀 Match Resumes"):

    if resume_files and job_description:

        with st.spinner("Analyzing resumes..."):

            jd_skills = extract_skills_auto(job_description)

            results = []

            progress_bar = st.progress(0)

            total_files = len(resume_files)

            # =========================================================
            # PROCESS RESUMES
            # =========================================================
            for idx, resume_file in enumerate(resume_files):

                try:

                    resume_text = clean_text(
                        extract_text_from_pdf(resume_file)
                    )

                    resume_skills = extract_skills_auto(
                        resume_text
                    )

                    score = calculate_similarity(
                        " ".join(resume_skills),
                        " ".join(jd_skills)
                    )

                    results.append({
                        "Resume": resume_file.name,
                        "ATS Score": score,
                        "Skills": resume_skills[:8]
                    })

                except Exception as e:

                    results.append({
                        "Resume": resume_file.name,
                        "ATS Score": 0,
                        "Skills": [f"Error: {str(e)}"]
                    })

                progress_bar.progress((idx + 1) / total_files)

            # =========================================================
            # SORT RESULTS
            # =========================================================
            results = sorted(
                results,
                key=lambda x: x["ATS Score"],
                reverse=True
            )

            # =========================================================
            # TOP STATS
            # =========================================================
            st.markdown("## 📊 ATS Analytics")

            col1, col2, col3 = st.columns(3)

            with col1:
                st.metric(
                    "Total Resumes",
                    len(results)
                )

            with col2:
                st.metric(
                    "Top ATS Score",
                    f"{results[0]['ATS Score']}%"
                )

            with col3:
                avg_score = round(
                    sum(r["ATS Score"] for r in results) / len(results),
                    2
                )

                st.metric(
                    "Average Score",
                    f"{avg_score}%"
                )

            st.divider()

            # =========================================================
            # TABLE
            # =========================================================
            st.markdown("## 🏆 Candidate Rankings")

            table_data = []

            for idx, item in enumerate(results, start=1):

                table_data.append({
                    "Rank": idx,
                    "Resume": item["Resume"],
                    "ATS Score": f"{item['ATS Score']}%"
                })

            df = pd.DataFrame(table_data)

            st.dataframe(
                df,
                use_container_width=True
            )

            # =========================================================
            # TOP CANDIDATES
            # =========================================================
            st.markdown("## ⭐ Best Candidates")

            top_candidates = results[:5]

            for idx, candidate in enumerate(top_candidates, start=1):

                score = candidate["ATS Score"]

                # Status
                if score >= 85:
                    status = "Excellent Match ✅"

                elif score >= 65:
                    status = "Good Match 👍"

                else:
                    status = "Average Match ⚠️"

                # CARD
                st.markdown(
                    '<div class="candidate-card">',
                    unsafe_allow_html=True
                )

                col1, col2 = st.columns([1, 5])

                # LEFT
                with col1:

                    st.metric(
                        label=f"🏅 Rank #{idx}",
                        value=f"{score}%"
                    )

                # RIGHT
                with col2:

                    st.markdown(
                        f"### 📄 {candidate['Resume']}"
                    )

                    st.progress(score / 100)

                    st.markdown(
                        f"**Status:** {status}"
                    )

                    st.markdown("#### 🛠 Matching Skills")

                    skill_html = ""

                    for skill in candidate["Skills"]:

                        skill_html += f"""
                        <span class="skill-tag">
                        {skill}
                        </span>
                        """

                    st.markdown(
                        skill_html,
                        unsafe_allow_html=True
                    )

                st.markdown(
                    '</div>',
                    unsafe_allow_html=True
                )

            # =========================================================
            # DOWNLOAD CSV
            # =========================================================
            csv = pd.DataFrame(results).to_csv(index=False).encode("utf-8")

            st.download_button(
                label="📥 Download ATS Report CSV",
                data=csv,
                file_name="ats_report.csv",
                mime="text/csv"
            )

    else:

        st.warning(
            "⚠️ Please upload resumes and enter a job description."
        )

# =========================================================
# FOOTER
# =========================================================
st.markdown("""
<br><br>
<hr>
<p style='text-align:center; color:gray;'>
Developed by <b>Ashvinkumar Bari</b> • AI Resume Screening Platform
</p>
""", unsafe_allow_html=True)