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
import gc
from transformers import AutoModelForCausalLM, AutoTokenizer
import time
import faiss
import re

# Fix for older PyTorch versions that don't have get_default_device
if not hasattr(torch, 'get_default_device'):
    def get_default_device():
        if torch.cuda.is_available():
            return torch.device('cuda')
        else:
            return torch.device('cpu')
    torch.get_default_device = get_default_device

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

# Set page configuration
st.set_page_config(
    page_title="AI-driven Candidate Matcher",
    page_icon="🎯",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Sidebar configuration
with st.sidebar:
    st.title("⚙️ Configuration")
    
    # Advanced options
    st.subheader("Display Options")
    top_k = st.selectbox("Number of results to display", options=[1, 2, 3, 4, 5], index=4)
    
    # LLM Settings
    st.subheader("LLM Settings")
    st.info("💡 Intent analysis using Qwen3-1.7B is always enabled")
    
    st.markdown("---")
    st.markdown("### 🤖 Pipeline Overview")
    st.markdown("**5-Stage Advanced Pipeline:**")
    st.markdown("1. FAISS Recall (Top 50)")
    st.markdown("2. Cross-Encoder Re-ranking (Top 20)")
    st.markdown("3. BM25 Keyword Matching")
    st.markdown("4. LLM Intent Analysis")
    st.markdown("5. Combined Scoring")
    
    st.markdown("### 📈 Scoring Formula")
    st.markdown("**Final Score = Cross-Encoder (0-0.7) + BM25 (0.1-0.2) + Intent (0-0.1)**")

# 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 'file_uploader_key' not in st.session_state:
    st.session_state.file_uploader_key = 0
if 'job_description_key' not in st.session_state:
    st.session_state.job_description_key = 1000
# No need for Qwen3-14B model since we're not generating explanations

# Separate smaller model for intent analysis
try:
    if 'qwen3_intent_tokenizer' not in st.session_state:
        st.session_state.qwen3_intent_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-1.7B")
    if 'qwen3_intent_model' not in st.session_state:
        st.session_state.qwen3_intent_model = AutoModelForCausalLM.from_pretrained(
            "Qwen/Qwen3-1.7B",
            torch_dtype="auto",
            device_map="auto"
        )
except Exception as e:
    st.warning(f"⚠️ Could not load Qwen3-1.7B: {str(e)}")
    st.session_state.qwen3_intent_tokenizer = None
    st.session_state.qwen3_intent_model = None

@st.cache_resource
def load_embedding_model():
    """Load and cache the BGE embedding model"""
    try:
        with st.spinner("🔄 Loading BAAI/bge-large-en-v1.5 model..."):
            # Try with explicit device specification
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
            model = SentenceTransformer('BAAI/bge-large-en-v1.5', device=device)
            st.success("✅ Embedding model loaded successfully!")
            return model
    except Exception as e:
        st.error(f"❌ Error loading embedding model: {str(e)}")
        try:
            # Fallback: try with a smaller model
            st.warning("🔄 Trying fallback model: all-MiniLM-L6-v2...")
            model = SentenceTransformer('all-MiniLM-L6-v2')
            st.success("✅ Fallback embedding model loaded!")
            return model
        except Exception as e2:
            st.error(f"❌ Fallback also failed: {str(e2)}")
            return None

@st.cache_resource
def load_cross_encoder():
    """Load and cache the Cross-Encoder model"""
    try:
        with st.spinner("🔄 Loading Cross-Encoder ms-marco-MiniLM-L6-v2..."):
            from sentence_transformers import CrossEncoder
            # Try with explicit device specification and logistic scoring
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
            model = CrossEncoder('cross-encoder/ms-marco-MiniLM-L6-v2', device=device)
            st.success("✅ Cross-Encoder model loaded successfully!")
            return model
    except Exception as e:
        st.error(f"❌ Error loading Cross-Encoder model: {str(e)}")
        try:
            # Fallback: try without device specification but with logistic scoring
            st.warning("🔄 Trying Cross-Encoder without device specification...")
            model = CrossEncoder('cross-encoder/ms-marco-MiniLM-L6-v2')
            st.success("✅ Cross-Encoder model loaded (fallback)!")
            return model
        except Exception as e2:
            st.error(f"❌ Cross-Encoder fallback also failed: {str(e2)}")
            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):
        # Load models
        self.embedding_model = load_embedding_model()
        self.cross_encoder = load_cross_encoder()
    
    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, final_top_k=5):
        """Advanced pipeline: FAISS recall -> Cross-encoder -> BM25 -> LLM intent -> Final ranking"""
        if not resume_texts:
            return []
            
        # Stage 1: FAISS Recall (Top 50)
        st.write("🔍 **Stage 1**: FAISS Recall - Finding top 50 candidates...")
        top_50_indices = self.faiss_recall(resume_texts, job_description, top_k=50)
        
        # Stage 2: Cross-Encoder Re-ranking (Top 20)
        st.write("🎯 **Stage 2**: Cross-Encoder Re-ranking - Selecting top 20...")
        top_20_results = self.cross_encoder_rerank(resume_texts, job_description, top_50_indices, top_k=20)
        
        # Stage 3: BM25 Keyword Matching
        st.write("🔤 **Stage 3**: BM25 Keyword Matching...")
        top_20_with_bm25 = self.add_bm25_scores(resume_texts, job_description, top_20_results)
        
        # Stage 4: LLM Intent Analysis (using Qwen3-1.7B)
        st.write("🤖 **Stage 4**: LLM Intent Analysis...")
        top_20_with_intent = self.add_intent_scores(resume_texts, job_description, top_20_with_bm25)
        
        # Stage 5: Final Combined Ranking
        st.write(f"🏆 **Stage 5**: Final Combined Ranking - Selecting top {final_top_k}...")
        final_results = self.calculate_final_scores(top_20_with_intent)
        
        return final_results[:final_top_k]  # Return top K as selected by user
    
    def faiss_recall(self, resume_texts, job_description, top_k=50):
        """Stage 1: Use FAISS for initial recall to find top 50 resumes"""
        try:
            # Get job embedding
            job_embedding = self.get_embedding(job_description)
            
            st.write(f"🔄 Generating embeddings for {len(resume_texts)} resumes...")
            
            # Get resume embeddings
            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))
            
            progress_bar.empty()
            
            st.write("🔍 Building FAISS index and searching...")
            
            # Create 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)
            
            # Search for top K
            job_embedding = job_embedding.reshape(1, -1).astype('float32')
            scores, indices = index.search(job_embedding, min(top_k, len(resume_texts)))
            
            # Show completion message
            st.write(f"✅ FAISS recall completed! Found top {min(top_k, len(resume_texts))} candidates.")
            
            return indices[0].tolist()
            
        except Exception as e:
            st.error(f"Error in FAISS recall: {str(e)}")
            # Fallback: return all indices
            return list(range(min(top_k, len(resume_texts))))
    
    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"""
        try:
            if not self.cross_encoder:
                st.error("Cross-encoder not loaded!")
                return [(idx, 0.0) for idx in top_50_indices[:top_k]]
            
            st.write(f"🔄 Processing {len(top_50_indices)} candidates with Cross-Encoder...")
            
            # 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:
                st.warning("No valid pairs found for cross-encoder!")
                return [(idx, 0.0) for idx in top_50_indices[:top_k]]
            
            st.write(f"📊 Cross-Encoder analyzing {len(pairs)} resume-job pairs...")
            
            # Get cross-encoder scores
            progress_bar = st.progress(0)
            scores = []
            
            def safe_sigmoid(x):
                """Safe sigmoid function that handles overflow"""
                if x >= 0:
                    exp_neg_x = np.exp(-x)
                    return 1 / (1 + exp_neg_x)
                else:
                    exp_x = np.exp(x)
                    return exp_x / (1 + exp_x)
            
            # 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]
                # Get raw logits from cross-encoder
                batch_scores = self.cross_encoder.predict(batch)
                # Apply sigmoid to convert logits to [0,1] range
                batch_scores_sigmoid = [safe_sigmoid(score) for score in batch_scores]
                scores.extend(batch_scores_sigmoid)
                progress_bar.progress(min(1.0, (i + batch_size) / len(pairs)))
            
            progress_bar.empty()
            
            # Combine indices with scores and sort
            indexed_scores = list(zip(valid_indices, scores))
            indexed_scores.sort(key=lambda x: x[1], reverse=True)
            
            # Normalize scores to 0-0.7 range (highest score becomes 0.7)
            if scores and len(scores) > 0:
                max_score = max(scores)
                min_score = min(scores)
                
                if max_score > min_score:
                    # Scale to 0-0.7 range
                    normalized_indexed_scores = []
                    for idx, score in indexed_scores:
                        normalized_score = 0.7 * (score - min_score) / (max_score - min_score)
                        normalized_indexed_scores.append((idx, normalized_score))
                    indexed_scores = normalized_indexed_scores
                else:
                    # All scores are the same, give them all 0.35 (middle value)
                    indexed_scores = [(idx, 0.35) for idx, _ in indexed_scores]
            
            # Show completion message
            st.write(f"✅ Cross-Encoder completed! Selected top {min(top_k, len(indexed_scores))} candidates.")
            st.write(f"📊 Cross-Encoder scores normalized to 0-0.7 range (highest: {indexed_scores[0][1]:.3f})")
            
            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"""
        try:
            st.write(f"🔄 Calculating BM25 keyword scores for {len(top_20_results)} candidates...")
            
            # 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))
            
            st.write(f"✅ BM25 keyword matching completed!")
            
            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"""
        try:
            results_with_intent = []
            progress_bar = st.progress(0)
            
            for i, (idx, cross_score, bm25_score) in enumerate(top_20_with_bm25):
                candidate_name = st.session_state.file_names[idx] if idx < len(st.session_state.file_names) else f"Resume_{idx}"
                intent_score, intent_text = self.analyze_intent(resume_texts[idx], job_description)
                
                # Print the intent analysis result
                st.write(f"📋 **{candidate_name}**: Intent = **{intent_text}** (Score: {intent_score:.1f})")
                
                results_with_intent.append((idx, cross_score, bm25_score, intent_score))
                progress_bar.progress((i + 1) / len(top_20_with_bm25))
            
            progress_bar.empty()
            return results_with_intent
            
        except Exception as e:
            st.error(f"Error adding intent scores: {str(e)}")
            return [(idx, cross_score, bm25_score, 0.1) for idx, cross_score, bm25_score in top_20_with_bm25]
    
    def analyze_intent(self, resume_text, job_description):
        """Analyze candidate's intent using LLM"""
        try:
            # Truncate texts
            resume_snippet = resume_text[:1500] if len(resume_text) > 1500 else resume_text
            job_snippet = job_description[:800] if len(job_description) > 800 else job_description
            
            prompt = f"""You are a helpful HR assistant. Look at this candidate's resume and job posting.

The candidate is likely a good fit if they have ANY of these:
- Related work experience (even if different industry)
- Relevant technical skills
- Educational background that could apply
- Any transferable skills
- Similar job titles or responsibilities

Be generous in your assessment. Most candidates who made it this far are potentially suitable.

Answer "Yes" for most candidates unless they are completely unrelated.
Answer "No" only if absolutely no connection exists.

Job Posting:
{job_snippet}

Candidate Resume:
{resume_snippet}

Is this candidate suitable? Answer:"""
            
            response = generate_qwen3_response(
                prompt,
                st.session_state.qwen3_intent_tokenizer,
                st.session_state.qwen3_intent_model,
                max_new_tokens=20
            )
            
            # Debug: print the raw response
            print(f"Raw LLM response: '{response}'")
            
            # Parse response - look for the answer directly
            response_lower = response.lower().strip()
            if 'yes' in response_lower:
                return 0.1, "Yes"
            elif 'no' in response_lower:
                return 0.0, "No"
            else:
                # If no clear answer, default to "Yes" to be more lenient
                print(f"Unclear response, defaulting to Yes: '{response}'")
                return 0.1, "Yes"
                
        except Exception as e:
            st.warning(f"Error analyzing intent: {str(e)}")
            return 0.1, "Yes"  # Default to "Yes" instead of "Maybe"
    
    def calculate_final_scores(self, results_with_all_scores):
        """Stage 5: Calculate final combined scores"""
        try:
            st.write(f"🔄 Computing final combined scores for {len(results_with_all_scores)} candidates...")
            
            final_results = []
            
            for idx, cross_score, bm25_score, intent_score in results_with_all_scores:
                # Cross-encoder scores are already in [0,1] range with logistic scoring
                normalized_cross = cross_score
                
                # Final Score = Cross-Encoder (0-0.7) + BM25 (0.1-0.2) + Intent (0-0.1)
                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)
            
            st.write(f"✅ Final ranking completed! Candidates sorted by combined score.")
            
            return final_results
            
        except Exception as e:
            st.error(f"Error calculating final scores: {str(e)}")
            return []
    
    def generate_simple_explanation(self, score, semantic_score, bm25_score):
        """Generate simple explanation for the match (fallback)"""
        if score > 0.8:
            quality = "excellent"
        elif score > 0.6:
            quality = "strong"
        elif score > 0.4:
            quality = "moderate"
        else:
            quality = "limited"
        
        explanation = f"This candidate shows {quality} alignment with the position (score: {score:.2f}). "
        
        if semantic_score > bm25_score:
            explanation += f"The resume demonstrates strong conceptual relevance ({semantic_score:.2f}) suggesting good experience fit. "
        else:
            explanation += f"The resume has high keyword match ({bm25_score:.2f}) indicating direct skill alignment. "
        
        return explanation
    


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-driven Candidate Matcher")
st.markdown("*Advanced 5-stage pipeline using BAAI/bge-large-en-v1.5 embeddings, Cross-Encoder re-ranking, and Qwen3-1.7B intent analysis*")

# Privacy Statement
with st.expander("🔒 Privacy & Data Security", expanded=False):
    st.markdown("""
    ### Data Privacy Statement
    
    **Your privacy is our top priority. We are committed to protecting the confidentiality of all resume data.**
    
    #### 🛡️ Data Handling
    - **No Permanent Storage**: Resume content is processed in memory only and never stored permanently
    - **Session-Based**: All data is automatically cleared when you close the browser or reset the application
    - **Local Processing**: All AI analysis happens locally within this application environment
    - **No External Transmission**: Resume data is never sent to external services or third parties
    
    #### 🔐 Security Measures
    - **Temporary Files**: Uploaded files are processed in secure temporary locations and immediately deleted
    - **Memory Management**: Automatic cleanup of resume data from system memory
    - **No Logging**: Resume content is never logged or cached anywhere
    - **Secure Processing**: All text extraction and analysis occurs within isolated processing environments
    
    #### 👤 User Control
    - **Clear Data Options**: Multiple options available to clear resume data and free memory
    - **Session Management**: Complete control over when and how your data is processed
    - **Transparent Processing**: Full visibility into what data is being analyzed
    
    **We recommend reviewing your organization's data handling policies before uploading sensitive resume information.**
    """)

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...",
    key=st.session_state.job_description_key
)

# 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.session_state.file_uploader_key += 1
            st.session_state.job_description_key += 1
            st.rerun()

input_method = st.radio(
    "Choose input method:",
    ["📁 Upload Files", "🗂️ Load from CSV 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",
        key=st.session_state.file_uploader_key
    )
    
    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)}")

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

# Run Advanced Pipeline Analysis
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"):
    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:
        with st.spinner("🚀 Running Advanced Pipeline Analysis..."):
            try:
                # Run the advanced pipeline
                pipeline_results = screener.advanced_pipeline_ranking(
                    st.session_state.resume_texts, job_description, final_top_k=top_k
                )
                
                # 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]
                    
                    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': [],
                        'text': text,
                        'text_preview': text[:500] + "..." if len(text) > 500 else text,
                        'explanation': None  # Will be filled with simple explanation
                    })
                
                # Add simple explanations
                for result in results:
                    result['explanation'] = screener.generate_simple_explanation(
                        result['final_score'], 
                        result['cross_encoder_score'], 
                        result['bm25_score']
                    )
                
                # 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.")
                
            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.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})"
            })
        
        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.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,
                "AI_Explanation": result['explanation'],
                "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.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-0.7)")
                    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']:.2f})", help="Job seeking likelihood (0-0.1)")
                
                with col2:
                    st.write("**💡 AI-Generated Match Analysis:**")
                    st.info(result['explanation'])
                    
                    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.session_state.file_uploader_key += 1
        st.session_state.job_description_key += 1
        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.session_state.file_uploader_key += 1
        st.session_state.job_description_key += 1
        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, Cross-Encoder/ms-marco-MiniLM-L6-v2 & Qwen3-1.7B | Built with Streamlit
    </div>
    """, 
    unsafe_allow_html=True
)