import streamlit as st
import pandas as pd
import numpy as np
import torch
import nltk
import os
import tempfile
import base64
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer, CrossEncoder
from nltk.tokenize import word_tokenize
import pdfplumber
import PyPDF2
from docx import Document
import csv
from datasets import load_dataset
import gc
from transformers import AutoModelForCausalLM, AutoTokenizer
import time
import faiss
import re

# Download NLTK resources
try:
    nltk.data.find('tokenizers/punkt')
except LookupError:
    nltk.download('punkt')

# Set page configuration
st.set_page_config(
    page_title="AI Resume Screener",
    page_icon="🎯",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Sidebar configuration
with st.sidebar:
    st.title("⚙️ Configuration")
    
    # Advanced options
    st.subheader("Advanced Options")
    top_k = st.selectbox("Number of results to display", [1,2,3,4,5], index=4)
    
    # LLM Settings
    st.subheader("LLM Settings")
    use_llm_explanations = st.checkbox("Generate AI Explanations", value=True)
    if use_llm_explanations:
        hf_token = st.text_input("Hugging Face Token (optional)", type="password", 
                                help="Enter your HF token for better rate limits")
    
    st.markdown("---")
    st.markdown("### 🤖 Advanced Pipeline")
    st.markdown("- **Stage 1**: FAISS Recall (Top 50)")
    st.markdown("- **Stage 2**: Cross-Encoder Re-ranking (Top 20)")
    st.markdown("- **Stage 3**: BM25 Keyword Matching")
    st.markdown("- **Stage 4**: LLM Intent Analysis")
    st.markdown("- **Final**: Combined Scoring (Top 5)")
    st.markdown("### 📊 Models Used")
    st.markdown("- **Embedding**: BAAI/bge-large-en-v1.5")
    st.markdown("- **Cross-Encoder**: ms-marco-MiniLM-L6-v2")
    st.markdown("- **LLM**: Qwen/Qwen3-1.7B")
    st.markdown("### 📈 Scoring Formula")
    st.markdown("**Final Score = Cross-Encoder (0-1) + BM25 (0.1-0.2) + Intent (0-0.3)**")

# Initialize session state
if 'embedding_model' not in st.session_state:
    st.session_state.embedding_model = None
if 'cross_encoder' not in st.session_state:
    st.session_state.cross_encoder = None
if 'results' not in st.session_state:
    st.session_state.results = []
if 'resume_texts' not in st.session_state:
    st.session_state.resume_texts = []
if 'file_names' not in st.session_state:
    st.session_state.file_names = []
if 'current_job_description' not in st.session_state:
    st.session_state.current_job_description = ""
if 'qwen3_1_7b_tokenizer' not in st.session_state:
    print("[Init] Loading Qwen3-1.7B Tokenizer...")
    st.session_state.qwen3_1_7b_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-1.7B")
    print("[Init] Qwen3-1.7B Tokenizer Loaded.")
if 'qwen3_1_7b_model' not in st.session_state:
    print("[Init] Loading Qwen3-1.7B Model...")
    st.session_state.qwen3_1_7b_model = AutoModelForCausalLM.from_pretrained(
        "Qwen/Qwen3-1.7B", torch_dtype="auto", device_map="auto"
    )
    print("[Init] Qwen3-1.7B Model Loaded.")

@st.cache_resource
def load_embedding_model():
    """Load and cache the BGE embedding model"""
    print("[Cache] Attempting to load Embedding Model (BAAI/bge-large-en-v1.5)...")
    try:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"[Cache] Using device: {device} for embedding model")
        with st.spinner("🔄 Loading BAAI/bge-large-en-v1.5 model..."):
            model = SentenceTransformer('BAAI/bge-large-en-v1.5', device=device)
            st.success("✅ Embedding model loaded successfully!")
            print("[Cache] Embedding Model (BAAI/bge-large-en-v1.5) LOADED.")
            return model
    except Exception as e:
        st.error(f"❌ Error loading embedding model: {str(e)}")
        return None

@st.cache_resource
def load_cross_encoder():
    """Load and cache the Cross-Encoder model"""
    print("[Cache] Attempting to load Cross-Encoder Model (ms-marco-MiniLM-L6-v2)...")
    try:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"[Cache] Using device: {device} for cross-encoder model")
        with st.spinner("🔄 Loading Cross-Encoder ms-marco-MiniLM-L6-v2..."):
            from sentence_transformers import CrossEncoder
            model = CrossEncoder('cross-encoder/ms-marco-MiniLM-L6-v2', device=device)
            st.success("✅ Cross-Encoder model loaded successfully!")
            print("[Cache] Cross-Encoder Model (ms-marco-MiniLM-L6-v2) LOADED.")
            return model
    except Exception as e:
        st.error(f"❌ Error loading Cross-Encoder model: {str(e)}")
        return None

def generate_qwen3_response(prompt, tokenizer, model, max_new_tokens=200):
    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=True
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=max_new_tokens
    )
    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
    response = tokenizer.decode(output_ids, skip_special_tokens=True).strip("\n")
    return response

class ResumeScreener:
    def __init__(self):
        print("[ResumeScreener] Initializing...")
        st.text("Initializing Screener: Loading embedding model...")
        self.embedding_model = load_embedding_model()
        st.text("Initializing Screener: Loading cross-encoder model...")
        self.cross_encoder = load_cross_encoder()
        print("[ResumeScreener] Initialized.")
        st.text("Screener Ready.")
    
    def extract_text_from_file(self, file_path, file_type):
        """Extract text from various file types"""
        try:
            if file_type == "pdf":
                with open(file_path, 'rb') as file:
                    with pdfplumber.open(file) as pdf:
                        text = ""
                        for page in pdf.pages:
                            text += page.extract_text() or ""
                        
                        if not text.strip():
                            # Fallback to PyPDF2
                            file.seek(0)
                            reader = PyPDF2.PdfReader(file)
                            text = ""
                            for page in reader.pages:
                                text += page.extract_text() or ""
                        return text
                        
            elif file_type == "docx":
                doc = Document(file_path)
                return " ".join([paragraph.text for paragraph in doc.paragraphs])
                
            elif file_type == "txt":
                with open(file_path, 'r', encoding='utf-8') as file:
                    return file.read()
                    
            elif file_type == "csv":
                with open(file_path, 'r', encoding='utf-8') as file:
                    csv_reader = csv.reader(file)
                    return " ".join([" ".join(row) for row in csv_reader])
                    
        except Exception as e:
            st.error(f"Error extracting text from {file_path}: {str(e)}")
            return ""
    
    def get_embedding(self, text):
        """Generate embedding for text using BGE model"""
        if self.embedding_model is None:
            st.error("No embedding model loaded!")
            return np.zeros(1024)  # BGE-large dimension
            
        try:
            # BGE models recommend adding instruction for retrieval
            # For queries (job description)
            if len(text) < 500:  # Assuming shorter texts are queries
                text = "Represent this sentence for searching relevant passages: " + text
            
            # Truncate text to avoid memory issues
            text = text[:8192] if text else ""
            
            # Generate embedding
            embedding = self.embedding_model.encode(text, 
                                                  convert_to_numpy=True,
                                                  normalize_embeddings=True)
            return embedding
                    
        except Exception as e:
            st.error(f"Error generating embedding: {str(e)}")
            return np.zeros(1024)  # BGE-large dimension
    
    def calculate_bm25_scores(self, resume_texts, job_description):
        """Calculate BM25 scores for keyword matching"""
        try:
            job_tokens = word_tokenize(job_description.lower())
            corpus = [word_tokenize(text.lower()) for text in resume_texts if text and text.strip()]
            
            if not corpus:
                return [0.0] * len(resume_texts)
                
            bm25 = BM25Okapi(corpus)
            scores = bm25.get_scores(job_tokens)
            return scores.tolist()
            
        except Exception as e:
            st.error(f"Error calculating BM25 scores: {str(e)}")
            return [0.0] * len(resume_texts)
    
    def advanced_pipeline_ranking(self, resume_texts, job_description):
        """Advanced pipeline: FAISS recall -> Cross-encoder -> BM25 -> LLM intent -> Final ranking"""
        print("[Pipeline] Advanced Pipeline Ranking started.")
        if not resume_texts:
            return []
        st.info("🔍 Stage 1: FAISS Recall - Finding top candidates...")
        print("[Pipeline] Calling faiss_recall.")
        top_50_indices = self.faiss_recall(resume_texts, job_description, top_k=50)
        print(f"[Pipeline] faiss_recall returned {len(top_50_indices)} indices.")

        st.info("🎯 Stage 2: Cross-Encoder Re-ranking - Selecting top candidates...")
        print("[Pipeline] Calling cross_encoder_rerank.")
        top_20_results = self.cross_encoder_rerank(resume_texts, job_description, top_50_indices, top_k=20)
        print(f"[Pipeline] cross_encoder_rerank returned {len(top_20_results)} results.")

        st.info("🔤 Stage 3: BM25 Keyword Matching...")
        print("[Pipeline] Calling add_bm25_scores.")
        top_20_with_bm25 = self.add_bm25_scores(resume_texts, job_description, top_20_results)
        print(f"[Pipeline] add_bm25_scores processed.")

        st.info("🤖 Stage 4: LLM Intent Analysis (Qwen3-1.7B)...")
        print("[Pipeline] Calling add_intent_scores.")
        top_20_with_intent = self.add_intent_scores(resume_texts, job_description, top_20_with_bm25)
        print(f"[Pipeline] add_intent_scores processed.")

        st.info("🏆 Stage 5: Final Combined Ranking...")
        print("[Pipeline] Calling calculate_final_scores.")
        final_results = self.calculate_final_scores(top_20_with_intent)
        print(f"[Pipeline] calculate_final_scores returned {len(final_results)} results.")
        print("[Pipeline] Advanced Pipeline Ranking finished.")
        return final_results[:5]  # Return top 5
    
    def faiss_recall(self, resume_texts, job_description, top_k=50):
        """Stage 1: Use FAISS for initial recall to find top 50 resumes"""
        print("[faiss_recall] Method started.")
        st.text("FAISS Recall: Embedding job description...")
        job_embedding = self.get_embedding(job_description)
        print("[faiss_recall] Job description embedded.")
        st.text(f"FAISS Recall: Embedding {len(resume_texts)} resumes...")
        resume_embeddings = []
        progress_bar = st.progress(0)
        
        for i, text in enumerate(resume_texts):
            if text:
                embedding = self.embedding_model.encode(text[:8192], 
                                                      convert_to_numpy=True,
                                                      normalize_embeddings=True)
                resume_embeddings.append(embedding)
            else:
                resume_embeddings.append(np.zeros(1024))
            progress_bar.progress((i + 1) / len(resume_texts))
            if i % 10 == 0: # Print progress every 10 resumes
                print(f"[faiss_recall] Embedded resume {i+1}/{len(resume_texts)}")
        
        progress_bar.empty()
        print("[faiss_recall] All resumes embedded.")
        st.text("FAISS Recall: Building FAISS index...")
        resume_embeddings = np.array(resume_embeddings).astype('float32')
        dimension = resume_embeddings.shape[1]
        index = faiss.IndexFlatIP(dimension)  # Inner product for cosine similarity
        index.add(resume_embeddings)
        print("[faiss_recall] FAISS index built.")
        st.text("FAISS Recall: Searching index...")
        job_embedding = job_embedding.reshape(1, -1).astype('float32')
        scores, indices = index.search(job_embedding, min(top_k, len(resume_texts)))
        print("[faiss_recall] FAISS search complete.")
        return indices[0].tolist()
    
    def cross_encoder_rerank(self, resume_texts, job_description, top_50_indices, top_k=20):
        """Stage 2: Use Cross-Encoder to re-rank top 50 and select top 20"""
        print("[cross_encoder_rerank] Method started.")
        try:
            if not self.cross_encoder:
                st.error("Cross-encoder not loaded!")
                return [(idx, 0.0) for idx in top_50_indices[:top_k]]
            
            # Prepare pairs for cross-encoder
            pairs = []
            valid_indices = []
            
            for idx in top_50_indices:
                if idx < len(resume_texts) and resume_texts[idx]:
                    # Truncate texts for cross-encoder
                    job_snippet = job_description[:512]
                    resume_snippet = resume_texts[idx][:512]
                    pairs.append([job_snippet, resume_snippet])
                    valid_indices.append(idx)
            
            if not pairs:
                return [(idx, 0.0) for idx in top_50_indices[:top_k]]
            
            st.text(f"Cross-Encoder: Preparing {len(pairs)} pairs for re-ranking...")
            print(f"[cross_encoder_rerank] Prepared {len(pairs)} pairs.")
            # Get cross-encoder scores
            progress_bar = st.progress(0)
            scores = []
            
            # Process in batches to avoid memory issues
            batch_size = 8
            for i in range(0, len(pairs), batch_size):
                batch = pairs[i:i+batch_size]
                batch_scores = self.cross_encoder.predict(batch)
                scores.extend(batch_scores)
                progress_bar.progress(min(1.0, (i + batch_size) / len(pairs)))
                print(f"[cross_encoder_rerank] Processed batch {i//batch_size + 1}")
            
            progress_bar.empty()
            print("[cross_encoder_rerank] All pairs scored.")
            st.text("Cross-Encoder: Re-ranking complete.")
            
            # Combine indices with scores and sort
            indexed_scores = list(zip(valid_indices, scores))
            indexed_scores.sort(key=lambda x: x[1], reverse=True)
            
            return indexed_scores[:top_k]
            
        except Exception as e:
            st.error(f"Error in cross-encoder re-ranking: {str(e)}")
            return [(idx, 0.0) for idx in top_50_indices[:top_k]]
    
    def add_bm25_scores(self, resume_texts, job_description, top_20_results):
        """Stage 3: Add BM25 scores to top 20 resumes"""
        print("[add_bm25_scores] Method started.")
        st.text("BM25: Calculating keyword scores...")
        try:
            # Get texts for top 20
            top_20_texts = [resume_texts[idx] for idx, _ in top_20_results]
            
            # Calculate BM25 scores
            bm25_scores = self.calculate_bm25_scores(top_20_texts, job_description)
            
            # Normalize BM25 scores to 0.1-0.2 range
            if bm25_scores and max(bm25_scores) > 0:
                max_bm25 = max(bm25_scores)
                min_bm25 = min(bm25_scores)
                if max_bm25 > min_bm25:
                    normalized_bm25 = [
                        0.1 + 0.1 * (score - min_bm25) / (max_bm25 - min_bm25) 
                        for score in bm25_scores
                    ]
                else:
                    normalized_bm25 = [0.15] * len(bm25_scores)
            else:
                normalized_bm25 = [0.15] * len(top_20_results)
            
            # Combine with existing results
            results_with_bm25 = []
            for i, (idx, cross_score) in enumerate(top_20_results):
                bm25_score = normalized_bm25[i] if i < len(normalized_bm25) else 0.15
                results_with_bm25.append((idx, cross_score, bm25_score))
            
            print("[add_bm25_scores] BM25 scores calculated and normalized.")
            st.text("BM25: Keyword scores added.")
            return results_with_bm25
            
        except Exception as e:
            st.error(f"Error adding BM25 scores: {str(e)}")
            return [(idx, cross_score, 0.15) for idx, cross_score in top_20_results]
    
    def add_intent_scores(self, resume_texts, job_description, top_20_with_bm25):
        """Stage 4: Add LLM intent analysis scores"""
        print("[add_intent_scores] Method started.")
        st.text(f"LLM Intent: Analyzing intent for {len(top_20_with_bm25)} candidates (Qwen3-1.7B)...")
        results_with_intent = []
        progress_bar = st.progress(0)
        
        for i, (idx, cross_score, bm25_score) in enumerate(top_20_with_bm25):
            intent_score = self.analyze_intent(resume_texts[idx], job_description)
            results_with_intent.append((idx, cross_score, bm25_score, intent_score))
            progress_bar.progress((i + 1) / len(top_20_with_bm25))
            print(f"[add_intent_scores] Intent analyzed for candidate {i+1}")
        
        progress_bar.empty()
        print("[add_intent_scores] All intents analyzed.")
        st.text("LLM Intent: Analysis complete.")
        return results_with_intent
    
    def analyze_intent(self, resume_text, job_description):
        """Analyze candidate's intent using Qwen3-1.7B LLM with thinking enabled."""
        print(f"[analyze_intent] Analyzing intent for one resume (Qwen3-1.7B)...")
        st.text("LLM Intent: Analyzing intent (Qwen3-1.7B)...")
        try:
            resume_snippet = resume_text[:15000]
            job_snippet = job_description[:5000]
            
            prompt = f"""You are given a job description and a candidate's resume.\nAnalyze the candidate's resume in detail against the job description to determine if they are genuinely seeking this specific job, or if their profile is a more general fit or perhaps a mismatch.\nProvide a step-by-step thought process for your decision.\nFinally, clearly answer: \"Is the candidate likely seeking THIS SPECIFIC job? Respond with 'Yes', 'Maybe', or 'No' and give a brief justification based on your thought process.\"\n\nJob Description:\n{job_snippet}\n\nCandidate Resume:\n{resume_snippet}\n\nResponse format:\n<think>\n[Your detailed step-by-step thought process comparing resume to JD, noting specific alignments or mismatches that indicate intent. Be thorough.]\n</think>\nIntent: [Yes/Maybe/No]\nReason: [Brief justification based on your thought process]"""
            
            response_text = generate_qwen3_response(
                prompt,
                st.session_state.qwen3_1_7b_tokenizer,
                st.session_state.qwen3_1_7b_model,
                max_new_tokens=20000
            )
            print(f"[analyze_intent] Qwen3-1.7B full response (first 100 chars): {response_text[:100]}...")

            thinking_content = "No detailed thought process extracted."
            intent_decision_part = response_text

            think_start_tag = "<think>"
            think_end_tag = "</think>"
            
            start_index = response_text.find(think_start_tag)
            end_index = response_text.rfind(think_end_tag)

            if start_index != -1 and end_index != -1 and start_index < end_index:
                thinking_content = response_text[start_index + len(think_start_tag):end_index].strip()
                intent_decision_part = response_text[end_index + len(think_end_tag):].strip()
                print(f"[analyze_intent] Thinking content extracted (first 50 chars): {thinking_content[:50]}...")
            else:
                print("[analyze_intent] <think> block not found or malformed in response.")

            response_lower = intent_decision_part.lower()
            intent_score = 0.1
            if 'intent: yes' in response_lower or 'intent:yes' in response_lower:
                intent_score = 0.3
            elif 'intent: no' in response_lower or 'intent:no' in response_lower:
                intent_score = 0.0
            
            print(f"[analyze_intent] Parsed Intent: {intent_score}, Decision part: {intent_decision_part[:100]}...")
            return intent_score

        except Exception as e:
            st.warning(f"Error analyzing intent with Qwen3-1.7B: {str(e)}")
            print(f"[analyze_intent] EXCEPTION: {str(e)}")
            return 0.1
    
    def calculate_final_scores(self, results_with_all_scores):
        """Stage 5: Calculate final combined scores"""
        print("[calculate_final_scores] Method started.")
        st.text("Final Ranking: Calculating combined scores...")
        try:
            final_results = []
            
            for idx, cross_score, bm25_score, intent_score in results_with_all_scores:
                # Normalize cross-encoder score to 0-1 range
                normalized_cross = max(0, min(1, cross_score))
                
                # Final Score = Cross-Encoder (0-1) + BM25 (0.1-0.2) + Intent (0-0.3)
                final_score = normalized_cross + bm25_score + intent_score
                
                final_results.append({
                    'index': idx,
                    'cross_encoder_score': normalized_cross,
                    'bm25_score': bm25_score,
                    'intent_score': intent_score,
                    'final_score': final_score
                })
            
            # Sort by final score
            final_results.sort(key=lambda x: x['final_score'], reverse=True)
            
            print("[calculate_final_scores] Final scores calculated and sorted.")
            st.text("Final Ranking: Complete.")
            return final_results
            
        except Exception as e:
            st.error(f"Error calculating final scores: {str(e)}")
            return []
    
    def extract_skills(self, text, job_description):
        """Extract skills from resume based on job description"""
        if not text:
            return []
            
        # Common tech skills
        common_skills = [
            "python", "java", "javascript", "react", "angular", "vue", "node.js", 
            "express", "django", "flask", "spring", "sql", "nosql", "html", "css",
            "aws", "azure", "gcp", "docker", "kubernetes", "jenkins", "git", "github",
            "agile", "scrum", "jira", "ci/cd", "devops", "microservices", "rest", "api",
            "machine learning", "deep learning", "data science", "artificial intelligence", 
            "tensorflow", "pytorch", "keras", "scikit-learn", "pandas", "numpy", 
            "matplotlib", "seaborn", "jupyter", "r", "sas", "spss", "tableau", "powerbi",
            "excel", "mysql", "postgresql", "mongodb", "redis", "elasticsearch", 
            "kafka", "rabbitmq", "spark", "hadoop", "hive", "airflow", "linux", "unix"
        ]
        
        # Extract potential skills from job description
        job_words = set(word.lower() for word in word_tokenize(job_description) if len(word) > 2)
        
        # Find matching skills
        found_skills = []
        text_lower = text.lower()
        
        # Check common skills that appear in both resume and job description
        for skill in common_skills:
            if skill in text_lower and any(skill in job_word for job_word in job_words):
                found_skills.append(skill)
        
        # Check for skills mentioned in job description
        for word in job_words:
            if len(word) > 3 and word in text_lower and word not in found_skills:
                # Basic filter to avoid common words
                if word not in ['with', 'have', 'that', 'this', 'from', 'what', 'when', 'where']:
                    found_skills.append(word)
        
        return list(set(found_skills))[:15]  # Return top 15 unique skills

def create_download_link(df, filename="resume_screening_results.csv"):
    """Create download link for results"""
    csv = df.to_csv(index=False)
    b64 = base64.b64encode(csv.encode()).decode()
    return f'<a href="data:file/csv;base64,{b64}" download="{filename}" class="download-btn">📥 Download Results CSV</a>'

# Main App Interface
st.title("🎯 AI-Powered Resume Screener")
st.markdown("*Find the perfect candidates using BAAI/bge-large-en-v1.5 embeddings and Qwen3-1.7B for intent analysis*")
st.markdown("---")

# Initialize screener
screener = ResumeScreener()

# Job Description Input
st.header("📝 Step 1: Enter Job Description")
job_description = st.text_area(
    "Enter the complete job description or requirements:",
    height=150,
    placeholder="Paste the job description here, including required skills, experience, and qualifications..."
)

# Resume Input Options
st.header("📄 Step 2: Upload Resumes")

# Show loaded resumes indicator
if st.session_state.resume_texts:
    col1, col2 = st.columns([3, 1])
    with col1:
        st.info(f"📚 {len(st.session_state.resume_texts)} resumes loaded and ready for analysis")
    with col2:
        if st.button("🗑️ Clear Resumes", type="secondary", help="Clear all loaded resumes to start fresh"):
            st.session_state.resume_texts = []
            st.session_state.file_names = []
            st.session_state.results = []
            st.session_state.current_job_description = ""
            st.rerun()

input_method = st.radio(
    "Choose input method:",
    ["📁 Upload Files", "🗂️ Load from CSV Dataset", "🔗 Load from Hugging Face Dataset"]
)

if input_method == "📁 Upload Files":
    uploaded_files = st.file_uploader(
        "Upload resume files",
        type=["pdf", "docx", "txt"],
        accept_multiple_files=True,
        help="Supported formats: PDF, DOCX, TXT"
    )
    
    if uploaded_files:
        with st.spinner(f"🔄 Processing {len(uploaded_files)} files..."):
            resume_texts = []
            file_names = []
            
            for file in uploaded_files:
                file_type = file.name.split('.')[-1].lower()
                
                with tempfile.NamedTemporaryFile(delete=False, suffix=f'.{file_type}') as tmp_file:
                    tmp_file.write(file.getvalue())
                    tmp_path = tmp_file.name
                
                text = screener.extract_text_from_file(tmp_path, file_type)
                if text.strip():
                    resume_texts.append(text)
                    file_names.append(file.name)
                
                os.unlink(tmp_path)
            
            st.session_state.resume_texts = resume_texts
            st.session_state.file_names = file_names
            
        if resume_texts:
            st.success(f"✅ Successfully processed {len(resume_texts)} resumes")

elif input_method == "🗂️ Load from CSV Dataset":
    csv_file = st.file_uploader("Upload CSV file with resume data", type=["csv"])
    
    if csv_file:
        try:
            df = pd.read_csv(csv_file)
            st.write("**CSV Preview:**")
            st.dataframe(df.head())
            
            text_column = st.selectbox(
                "Select column containing resume text:",
                df.columns.tolist()
            )
            
            name_column = st.selectbox(
                "Select column for candidate names/IDs (optional):",
                ["Use Index"] + df.columns.tolist()
            )
            
            if st.button("🚀 Process CSV Data"):
                with st.spinner("🔄 Processing CSV data..."):
                    resume_texts = []
                    file_names = []
                    
                    for idx, row in df.iterrows():
                        text = str(row[text_column])
                        if text and text.strip() and text.lower() != 'nan':
                            resume_texts.append(text)
                            
                            if name_column == "Use Index":
                                file_names.append(f"Resume_{idx}")
                            else:
                                file_names.append(str(row[name_column]))
                    
                    st.session_state.resume_texts = resume_texts
                    st.session_state.file_names = file_names
                
                if resume_texts:
                    st.success(f"✅ Successfully loaded {len(resume_texts)} resumes from CSV")
                    
        except Exception as e:
            st.error(f"❌ Error processing CSV: {str(e)}")

elif input_method == "🔗 Load from Hugging Face Dataset":
    st.markdown("**Popular Resume Datasets:**")
    st.markdown("- `ahmedheakl/resume-atlas`")
    st.markdown("- `InferenceFly/Resume-Dataset`")
    
    col1, col2 = st.columns([2, 1])
    with col1:
        dataset_name = st.text_input(
            "Dataset name:",
            value="ahmedheakl/resume-atlas",
            help="Enter Hugging Face dataset name"
        )
    with col2:
        dataset_split = st.selectbox("Split:", ["train", "test", "validation"], index=0)
    
    if st.button("🔗 Load from Hugging Face"):
        try:
            with st.spinner(f"🔄 Loading {dataset_name}..."):
                dataset = load_dataset(dataset_name, split=dataset_split)
                
            st.success(f"✅ Loaded dataset with {len(dataset)} entries")
            st.write("**Dataset Preview:**")
            
            preview_df = pd.DataFrame(dataset[:5])
            st.dataframe(preview_df)
            
            text_column = st.selectbox(
                "Select column with resume text:",
                dataset.column_names,
                index=dataset.column_names.index('resume_text') if 'resume_text' in dataset.column_names else 0
            )
            
            category_column = None
            if 'category' in dataset.column_names:
                categories = list(set(dataset['category']))
                category_column = st.selectbox(
                    "Filter by category (optional):",
                    ["All"] + categories
                )
            
            max_samples = st.slider("Maximum samples to load:", 10, min(1000, len(dataset)), 100)
            
            if st.button("🚀 Process Dataset"):
                with st.spinner("🔄 Processing dataset..."):
                    resume_texts = []
                    file_names = []
                    
                    filtered_dataset = dataset
                    
                    if category_column and category_column != "All":
                        filtered_dataset = dataset.filter(lambda x: x['category'] == category_column)
                    
                    sample_indices = list(range(min(max_samples, len(filtered_dataset))))
                    
                    for idx in sample_indices:
                        item = filtered_dataset[idx]
                        text = str(item[text_column])
                        
                        if text and text.strip() and text.lower() != 'nan':
                            resume_texts.append(text)
                            
                            if 'id' in item:
                                file_names.append(f"Resume_{item['id']}")
                            else:
                                file_names.append(f"Resume_{idx}")
                    
                    st.session_state.resume_texts = resume_texts
                    st.session_state.file_names = file_names
                
                if resume_texts:
                    st.success(f"✅ Successfully loaded {len(resume_texts)} resumes")
                    
        except Exception as e:
            st.error(f"❌ Error loading dataset: {str(e)}")

# Processing and Results
st.header("🔍 Step 3: Analyze Resumes")

# First button: Find top K candidates (fast ranking)
col1, col2 = st.columns([1, 1])

with col1:
    if st.button("🚀 Advanced Pipeline Analysis", 
                 disabled=not (job_description and st.session_state.resume_texts),
                 type="primary",
                 help="Run the complete 5-stage advanced pipeline"):
        print("--- Advanced Pipeline Analysis Button Clicked ---")
        if len(st.session_state.resume_texts) == 0:
            st.error("❌ Please upload resumes first!")
        elif not job_description.strip():
            st.error("❌ Please enter a job description!")
        else:
            print("[UI Button] Pre-checks passed. Starting spinner and pipeline.")
            with st.spinner("🚀 Running Advanced Pipeline Analysis..."):
                st.text("Pipeline Initiated: Starting advanced analysis...")
                try:
                    # Run the advanced pipeline
                    pipeline_results = screener.advanced_pipeline_ranking(
                        st.session_state.resume_texts, job_description
                    )
                    
                    # Prepare results for display
                    results = []
                    
                    for rank, result_data in enumerate(pipeline_results, 1):
                        idx = result_data['index']
                        name = st.session_state.file_names[idx]
                        text = st.session_state.resume_texts[idx]
                        
                        # Extract skills
                        skills = screener.extract_skills(text, job_description)
                        
                        results.append({
                            'rank': rank,
                            'name': name,
                            'final_score': result_data['final_score'],
                            'cross_encoder_score': result_data['cross_encoder_score'],
                            'bm25_score': result_data['bm25_score'],
                            'intent_score': result_data['intent_score'],
                            'skills': skills,
                            'text': text,
                            'text_preview': text[:500] + "..." if len(text) > 500 else text
                        })
                    
                    # Store in session state
                    st.session_state.results = results
                    st.session_state.current_job_description = job_description
                    
                    st.success(f"🚀 Advanced pipeline complete! Found top {len(st.session_state.results)} candidates.")
                    st.text("Displaying Top Candidates...")
                    
                except Exception as e:
                    st.error(f"❌ Error during analysis: {str(e)}")

# Display Results
if st.session_state.results:
    st.header("🏆 Top Candidates")
    
    # Create tabs for different views
    tab1, tab2, tab3 = st.tabs(["📊 Summary", "📋 Detailed Analysis", "📈 Visualizations"])
    
    with tab1:
        # Create summary dataframe with new scoring system
        summary_data = []
        for result in st.session_state.results:
            # Map intent score to text
            intent_text = "Yes" if result['intent_score'] == 0.3 else "Maybe" if result['intent_score'] == 0.1 else "No"
            
            summary_data.append({
                "Rank": result['rank'],
                "Candidate": result['name'],
                "Final Score": f"{result['final_score']:.2f}",
                "Cross-Encoder": f"{result['cross_encoder_score']:.2f}",
                "BM25": f"{result['bm25_score']:.2f}",
                "Intent": f"{intent_text} ({result['intent_score']:.1f})",
                "Top Skills": ", ".join(result['skills'][:5])
            })
        
        summary_df = pd.DataFrame(summary_data)
        
        # Style the dataframe
        def color_scores(val):
            if isinstance(val, str) and any(char.isdigit() for char in val):
                try:
                    # Extract numeric value
                    numeric_val = float(''.join(c for c in val if c.isdigit() or c == '.'))
                    if 'Final Score' in val or numeric_val >= 1.0:
                        if numeric_val >= 1.2:
                            return 'background-color: #d4edda'
                        elif numeric_val >= 1.0:
                            return 'background-color: #fff3cd'
                        else:
                            return 'background-color: #f8d7da'
                    else:
                        if numeric_val >= 0.7:
                            return 'background-color: #d4edda'
                        elif numeric_val >= 0.5:
                            return 'background-color: #fff3cd'
                        else:
                            return 'background-color: #f8d7da'
                except:
                    pass
            return ''
        
        styled_df = summary_df.style.applymap(color_scores, subset=['Final Score', 'Cross-Encoder', 'BM25'])
        st.dataframe(styled_df, use_container_width=True)
        
        # Download link
        detailed_data = []
        for result in st.session_state.results:
            intent_text = "Yes" if result['intent_score'] == 0.3 else "Maybe" if result['intent_score'] == 0.1 else "No"
            
            detailed_data.append({
                "Rank": result['rank'],
                "Candidate": result['name'],
                "Final_Score": result['final_score'],
                "Cross_Encoder_Score": result['cross_encoder_score'],
                "BM25_Score": result['bm25_score'],
                "Intent_Score": result['intent_score'],
                "Intent_Analysis": intent_text,
                "Skills": "; ".join(result['skills']),
                "Resume_Preview": result['text_preview']
            })
        
        download_df = pd.DataFrame(detailed_data)
        st.markdown(create_download_link(download_df), unsafe_allow_html=True)
    
    with tab2:
        # Detailed results with new scoring breakdown
        for result in st.session_state.results:
            intent_text = "Yes" if result['intent_score'] == 0.3 else "Maybe" if result['intent_score'] == 0.1 else "No"
            
            with st.expander(f"#{result['rank']}: {result['name']} (Final Score: {result['final_score']:.2f})"):
                col1, col2 = st.columns([1, 2])
                
                with col1:
                    st.metric("🏆 Final Score", f"{result['final_score']:.2f}")
                    
                    st.write("**📊 Score Breakdown:**")
                    st.metric("🎯 Cross-Encoder", f"{result['cross_encoder_score']:.2f}", help="Semantic relevance (0-1)")
                    st.metric("🔤 BM25 Keywords", f"{result['bm25_score']:.2f}", help="Keyword matching (0.1-0.2)")
                    st.metric("🤖 Intent Analysis", f"{intent_text} ({result['intent_score']:.1f})", help="Job seeking likelihood (0-0.3)")
                    
                    st.write("**🎯 Matching Skills:**")
                    skills_per_column = 5
                    skill_cols = st.columns(2)
                    for idx, skill in enumerate(result['skills'][:10]):
                        with skill_cols[idx % 2]:
                            st.write(f"• {skill}")
                
                with col2:
                    st.write("**📄 Resume Preview:**")
                    st.text_area("", result['text_preview'], height=200, disabled=True, key=f"preview_{result['rank']}")
    
    with tab3:
        # Score visualization
        if len(st.session_state.results) > 1:
            # Bar chart
            st.subheader("Score Comparison")
            
            chart_data = pd.DataFrame({
                'Candidate': [r['name'][:20] + '...' if len(r['name']) > 20 else r['name'] 
                             for r in st.session_state.results],
                'Final Score': [r['final_score'] for r in st.session_state.results],
                'Cross-Encoder': [r['cross_encoder_score'] for r in st.session_state.results],
                'BM25': [r['bm25_score'] for r in st.session_state.results],
                'Intent': [r['intent_score'] for r in st.session_state.results]
            })
            
            st.bar_chart(chart_data.set_index('Candidate'))
            
            # Score distribution
            col1, col2 = st.columns(2)
            
            with col1:
                st.subheader("Score Distribution")
                score_ranges = {
                    'Excellent (≥1.2)': sum(1 for r in st.session_state.results if r['final_score'] >= 1.2),
                    'Good (1.0-1.2)': sum(1 for r in st.session_state.results if 1.0 <= r['final_score'] < 1.2),
                    'Fair (0.8-1.0)': sum(1 for r in st.session_state.results if 0.8 <= r['final_score'] < 1.0),
                    'Poor (<0.8)': sum(1 for r in st.session_state.results if r['final_score'] < 0.8),
                }
                
                dist_df = pd.DataFrame({
                    'Range': score_ranges.keys(),
                    'Count': score_ranges.values()
                })
                st.bar_chart(dist_df.set_index('Range'))
            
            with col2:
                st.subheader("Average Scores")
                avg_final = np.mean([r['final_score'] for r in st.session_state.results])
                avg_cross = np.mean([r['cross_encoder_score'] for r in st.session_state.results])
                avg_bm25 = np.mean([r['bm25_score'] for r in st.session_state.results])
                avg_intent = np.mean([r['intent_score'] for r in st.session_state.results])
                
                st.metric("Average Final Score", f"{avg_final:.2f}")
                st.metric("Average Cross-Encoder", f"{avg_cross:.2f}")
                st.metric("Average BM25", f"{avg_bm25:.2f}")
                st.metric("Average Intent", f"{avg_intent:.2f}")

# Memory cleanup
st.markdown("---")
st.subheader("🧹 Reset Application")
col1, col2, col3 = st.columns([1, 1, 3])
with col1:
    if st.button("🗑️ Clear Resumes Only", type="secondary", help="Clear only the loaded resumes"):
        st.session_state.resume_texts = []
        st.session_state.file_names = []
        st.session_state.results = []
        st.session_state.current_job_description = ""
        st.success("✅ Resumes cleared!")
        st.rerun()

with col2:
    if st.button("🧹 Clear Everything", type="primary", help="Clear all data and free memory"):
        st.session_state.resume_texts = []
        st.session_state.file_names = []
        st.session_state.results = []
        st.session_state.current_job_description = ""
        
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()
        st.success("✅ Everything cleared!")
        st.rerun()

# Footer
st.markdown("---")
st.markdown(
    """
    <div style='text-align: center; color: #666;'>
        🚀 Powered by BAAI/bge-large-en-v1.5 & Qwen3-1.7B | Built with Streamlit
    </div>
    """, 
    unsafe_allow_html=True
)