import os
import io
import streamlit as st
import docx
from transformers import pipeline
import numpy as np
from scipy.spatial.distance import cosine
import time

# Set page title
st.set_page_config(page_title="Resume Analyzer and Company Suitability Checker")

#####################################
# Preload Models
#####################################
@st.cache_resource(show_spinner=True)
def load_models(summarization_model="google/pegasus-xsum", similarity_model="sentence-transformers/all-MiniLM-L6-v2"):
    """Load models at startup"""
    with st.spinner("Loading AI models... This may take a minute on first run."):
        models = {}
        # Load summarization model
        models['summarizer'] = pipeline("summarization", model=summarization_model)
        
        # Load feature extraction model for similarity
        models['feature_extractor'] = pipeline("feature-extraction", model=similarity_model)
        
        return models

# Preload models immediately when app starts
models = load_models()

#####################################
# Function: Extract Text from File
#####################################
def extract_text_from_file(file_obj):
    """
    Extract text from .docx files.
    Returns the extracted text or an error message if extraction fails.
    """
    filename = file_obj.name
    ext = os.path.splitext(filename)[1].lower()
    text = ""

    if ext == ".docx":
        try:
            document = docx.Document(file_obj)
            text = "\n".join(para.text for para in document.paragraphs if para.text.strip())
        except Exception as e:
            text = f"Error processing DOCX file: {e}"
    elif ext == ".txt":
        try:
            text = file_obj.getvalue().decode("utf-8")
        except Exception as e:
            text = f"Error processing TXT file: {e}"
    else:
        text = "Unsupported file type. Please upload a .docx or .txt file."
    return text

#####################################
# Function: Summarize Resume Text
#####################################
def summarize_resume_text(resume_text, models):
    """
    Generates a concise summary of the resume text using the selected summarization model.
    """
    start_time = time.time()
    
    summarizer = models['summarizer']
    
    # Handle long text
    max_input_length = 1024  # Model limit
    
    if len(resume_text) > max_input_length:
        # Process in chunks if text is too long
        chunks = [resume_text[i:i+max_input_length] for i in range(0, min(len(resume_text), 3*max_input_length), max_input_length)]
        summaries = []
        
        for chunk in chunks:
            chunk_summary = summarizer(chunk, max_length=100, min_length=30, do_sample=False)[0]['summary_text']
            summaries.append(chunk_summary)
        
        candidate_summary = " ".join(summaries)
        if len(candidate_summary) > max_input_length:
            candidate_summary = summarizer(candidate_summary[:max_input_length], max_length=150, min_length=40, do_sample=False)[0]['summary_text']
    else:
        candidate_summary = summarizer(resume_text, max_length=150, min_length=40, do_sample=False)[0]['summary_text']
    
    execution_time = time.time() - start_time
    
    return candidate_summary, execution_time

#####################################
# Function: Compare Candidate Summary to Company Prompt
#####################################
def compute_suitability(candidate_summary, company_prompt, models):
    """
    Compute the similarity between candidate summary and company prompt.
    Returns a score in the range [0, 1] and execution time.
    """
    start_time = time.time()
    
    feature_extractor = models['feature_extractor']
    
    # Extract features (embeddings)
    candidate_features = feature_extractor(candidate_summary)
    company_features = feature_extractor(company_prompt)
    
    # Convert to numpy arrays and flatten if needed
    candidate_vec = np.mean(np.array(candidate_features[0]), axis=0)
    company_vec = np.mean(np.array(company_features[0]), axis=0)
    
    # Compute cosine similarity (1 - cosine distance)
    similarity = 1 - cosine(candidate_vec, company_vec)
    
    execution_time = time.time() - start_time
    
    return similarity, execution_time

#####################################
# Main Streamlit Interface
#####################################
st.title("Resume Analyzer and Company Suitability Checker")
st.markdown(
    """
Upload your resume file in **.docx** or **.txt** format. The app performs the following tasks:
1. Extracts text from the resume.
2. Uses a transformer-based model to generate a concise candidate summary.
3. Compares the candidate summary with a company profile to produce a suitability score.
"""
)

# File uploader
uploaded_file = st.file_uploader("Upload your resume (.docx or .txt)", type=["docx", "txt"])

# Company description text area
company_prompt = st.text_area(
    "Enter the company description or job requirements:",
    height=150,
    help="Enter a detailed description of the company culture, role requirements, and desired skills.",
)

# Show model selection in sidebar
st.sidebar.header("Model Settings")

# Model dropdowns - we're now only allowing one model of each type to be selected
summarization_model = st.sidebar.selectbox(
    "Summarization Model",
    ["google/pegasus-xsum", "facebook/bart-large-cnn", "t5-small", "sshleifer/distilbart-cnn-12-6"],
    index=0,
    help="Select the model to use for summarizing the resume text."
)

similarity_model = st.sidebar.selectbox(
    "Similarity Model",
    ["sentence-transformers/all-MiniLM-L6-v2", "sentence-transformers/all-mpnet-base-v2", 
     "sentence-transformers/paraphrase-MiniLM-L3-v2", "sentence-transformers/multi-qa-mpnet-base-dot-v1"],
    index=0,
    help="Select the model to use for comparing candidate summary with company profile."
)

# Reload models if changed
if st.sidebar.button("Reload Models"):
    st.cache_resource.clear()
    models = load_models(summarization_model, similarity_model)
    st.sidebar.success("Models reloaded successfully!")

# Process button
if uploaded_file is not None and company_prompt and st.button("Analyze Resume"):
    with st.spinner("Processing..."):
        # Extract text from resume
        resume_text = extract_text_from_file(uploaded_file)
        
        if resume_text.startswith("Error") or resume_text == "Unsupported file type. Please upload a .docx or .txt file.":
            st.error(resume_text)
        else:
            # Display extracted text
            with st.expander("Extracted Text"):
                st.text(resume_text)
            
            # Generate summary
            summary, summarization_time = summarize_resume_text(resume_text, models)
            
            # Display summary
            st.subheader("Candidate Summary")
            st.write(summary)
            st.info(f"Summarization completed in {summarization_time:.2f} seconds")
            
            # Only compute similarity if company description is provided
            if company_prompt:
                similarity_score, similarity_time = compute_suitability(summary, company_prompt, models)
                
                # Display similarity score
                st.subheader("Suitability Assessment")
                st.markdown(f"**Matching Score:** {similarity_score:.2%}")
                st.info(f"Similarity computation completed in {similarity_time:.2f} seconds")
                
                # Provide interpretation
                if similarity_score >= 0.85:
                    st.success("Excellent match! This candidate's profile is strongly aligned with the company requirements.")
                elif similarity_score >= 0.70:
                    st.success("Good match! This candidate shows strong potential for the position.")
                elif similarity_score >= 0.50:
                    st.warning("Moderate match. The candidate meets some requirements but there may be gaps.")
                else:
                    st.error("Low match. The candidate's profile may not align well with the requirements.")