main.py

import streamlit as st
import os
from llama_index.core import SimpleDirectoryReader, Settings, VectorStoreIndex
from llama_index.embeddings.nebius import NebiusEmbedding
from llama_index.llms.nebius import NebiusLLM
from dotenv import load_dotenv
import tempfile
import shutil
import base64
from pathlib import Path
from PyPDF2 import PdfReader
import io

# Load .env from script directory first, then current working directory (covers streamlit + CLI)
_env_dir = Path(__file__).resolve().parent
load_dotenv(_env_dir / ".env")
load_dotenv()

# Use Nebius Token Factory API (where you get the API key). Studio has a different model catalog.
NEBIUS_TOKEN_FACTORY_API_BASE = "https://api.tokenfactory.nebius.com/v1"

def get_resoptimum_api_key() -> str:
    """Get Resoptimum API key from env only (no UI input for security). Returns stripped non-empty string or ''."""
    return (os.getenv("NEBIUS_API_KEY") or "").strip()

def run_rag_completion(
    documents,
    query_text: str,
    job_title: str,
    job_description: str,
    embedding_model: str = "BAAI/bge-en-icl",
    generative_model: str = "meta-llama/Llama-3.3-70B-Instruct"
) -> str:
    """Run RAG completion using Nebius models for resume optimization."""
    try:
        api_key = get_resoptimum_api_key()
        if not api_key:
            raise ValueError("Missing Resoptimum API key. Add NEBIUS_API_KEY=your_key to a .env file in this folder.")
        llm = NebiusLLM(
            model=generative_model,
            api_key=api_key,
            api_base=NEBIUS_TOKEN_FACTORY_API_BASE,
        )

        embed_model = NebiusEmbedding(
            model_name=embedding_model,
            api_key=api_key,
            api_base=NEBIUS_TOKEN_FACTORY_API_BASE,
        )
        
        Settings.llm = llm
        Settings.embed_model = embed_model
        
        # Step 1: Analyze the resume
        analysis_prompt = f"""
        Analyze this resume in detail. Focus on:
        1. Key skills and expertise
        2. Professional experience and achievements
        3. Education and certifications
        4. Notable projects or accomplishments
        5. Career progression and gaps
        
        Provide a concise analysis in bullet points.
        """
        
        index = VectorStoreIndex.from_documents(documents)
        resume_analysis = index.as_query_engine(similarity_top_k=5).query(analysis_prompt)
        
        # Step 2: Generate optimization suggestions
        optimization_prompt = f"""
        Based on the resume analysis and job requirements, provide specific, actionable improvements.
        
        Resume Analysis:
        {resume_analysis}
        
        Job Title: {job_title}
        Job Description: {job_description}
        
        Optimization Request: {query_text}
        
        Provide a direct, structured response in this exact format:

        ## Key Findings
        • [2-3 bullet points highlighting main alignment and gaps]

        ## Specific Improvements
        • [3-5 bullet points with concrete suggestions]
        • Each bullet should start with a strong action verb
        • Include specific examples where possible

        ## Action Items
        • [2-3 specific, immediate steps to take]
        • Each item should be clear and implementable

        Keep all points concise and actionable. Do not include any thinking process or analysis.
        """
        
        optimization_suggestions = index.as_query_engine(similarity_top_k=5).query(optimization_prompt)
        
        return str(optimization_suggestions)
    except Exception as e:
        raise

def display_pdf_preview(pdf_file):
    """Display PDF preview in the sidebar."""
    try:
        st.sidebar.subheader("Resume Preview")
        base64_pdf = base64.b64encode(pdf_file.getvalue()).decode('utf-8')
        pdf_display = f'<iframe src="data:application/pdf;base64,{base64_pdf}" width="100%" height="500" type="application/pdf"></iframe>'
        st.sidebar.markdown(pdf_display, unsafe_allow_html=True)
        return True
    except Exception as e:
        st.sidebar.error(f"Error previewing PDF: {str(e)}")
        return False

def main():
    st.set_page_config(page_title="Resume Optimizer", layout="wide", initial_sidebar_state="expanded")
    
    # Initialize session states (logic unchanged)
    if "messages" not in st.session_state:
        st.session_state.messages = []
    if "docs_loaded" not in st.session_state:
        st.session_state.docs_loaded = False
    if "temp_dir" not in st.session_state:
        st.session_state.temp_dir = None
    if "current_pdf" not in st.session_state:
        st.session_state.current_pdf = None

    # ========== UI ONLY: Glassmorphic dark theme CSS ==========
    st.markdown("""
    <style>
    /* Base: dark gradient + glowing background animation */
    @keyframes glowPulse {
        0%, 100% { opacity: 0.4; transform: scale(1); }
        50% { opacity: 0.7; transform: scale(1.05); }
    }
    @keyframes gradientShift {
        0% { background-position: 0% 50%; }
        50% { background-position: 100% 50%; }
        100% { background-position: 0% 50%; }
    }
    @keyframes textShimmer {
        0% { background-position: 0% 50%; }
        50% { background-position: 100% 50%; }
        100% { background-position: 0% 50%; }
    }
    @keyframes fadeInSlide {
        from { opacity: 0; transform: translateY(10px); }
        to { opacity: 1; transform: translateY(0); }
    }
    @keyframes letterWave {
        0%, 100% { transform: translateY(0); }
        50% { transform: translateY(-8px); }
    }

    .stApp {
        background: linear-gradient(160deg, #0f0f14 0%, #1a1a24 40%, #0d0d12 100%);
        background-size: 200% 200%;
        background-attachment: fixed;
        animation: gradientShift 18s ease infinite;
        position: relative;
        overflow: hidden;
    }
    .stApp::before {
        content: '';
        position: fixed;
        top: -50%;
        left: -50%;
        width: 200%;
        height: 200%;
        background: radial-gradient(ellipse at 30% 20%, rgba(124, 58, 237, 0.15) 0%, transparent 50%),
                    radial-gradient(ellipse at 70% 80%, rgba(59, 130, 246, 0.12) 0%, transparent 50%),
                    radial-gradient(ellipse at 50% 50%, rgba(139, 92, 246, 0.08) 0%, transparent 60%);
        pointer-events: none;
        animation: glowPulse 8s ease-in-out infinite;
        z-index: 0;
    }
    .stApp > div { position: relative; z-index: 1; }

    /* Main content area */
    .main .block-container {
        padding-top: 2rem;
        padding-bottom: 2rem;
        max-width: 1400px;
    }

    /* Sidebar: glassmorphic & FIXED position */
    [data-testid="stSidebar"] > div:first-child {
        background: rgba(26, 26, 36, 0.85);
        backdrop-filter: blur(25px);
        -webkit-backdrop-filter: blur(25px);
        border-right: 1px solid rgba(255, 255, 255, 0.1);
        box-shadow: 10px 0 30px rgba(0, 0, 0, 0.4);
    }
    
    /* Image styling: Glow and Beauty */
    .logo-container {
        text-align: center;
        padding: 20px 0;
        animation: fadeInSlide 1s ease-out;
    }
    .logo-container img {
        border-radius: 16px;
        box-shadow: 0 0 20px rgba(124, 58, 237, 0.3);
        transition: transform 0.3s ease, box-shadow 0.3s ease;
    }
    .logo-container img:hover {
        transform: scale(1.05);
        box-shadow: 0 0 30px rgba(124, 58, 237, 0.6);
    }

    [data-testid="stSidebar"] .stMarkdown { color: rgba(255,255,255,0.9); }
    [data-testid="stSidebar"] label { color: rgba(255,255,255,0.85) !important; }

    /* Sidebar inputs: glass input style */
    [data-testid="stSidebar"] .stTextInput input, [data-testid="stSidebar"] input {
        background: rgba(255, 255, 255, 0.08) !important;
        border: 1px solid rgba(255, 255, 255, 0.15) !important;
        border-radius: 12px !important;
        color: #e8e8ec !important;
        transition: all 0.2s;
    }
    [data-testid="stSidebar"] .stTextInput input:focus {
        border-color: rgba(139, 92, 246, 0.6) !important;
        background: rgba(255, 255, 255, 0.12) !important;
        box-shadow: 0 0 15px rgba(139, 92, 246, 0.2) !important;
    }

    /* Main: typography & Animations */
    .main-title { 
        font-size: 3.5rem !important;
        font-weight: 900 !important; 
        margin-bottom: 0.5rem !important;
        display: block !important;
        line-height: 1.2 !important;
    }
    .wave-char {
        display: inline-block;
        background: linear-gradient(90deg, #ffffff, #00ffff, #ffffff, #ff00ff, #ffffff);
        background-size: 200% auto;
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        animation: textShimmer 4s linear infinite, letterWave 2s ease-in-out infinite;
    }
    .main-title .emoji {
        display: inline-block;
        margin-right: 15px;
        font-size: 3.8rem;
        vertical-align: middle;
    }
    .main h2, .main h3 { 
        color: rgba(255,255,255,0.98) !important; 
        font-weight: 700 !important;
        animation: fadeInSlide 1s ease-out;
    }
    .main p, .main .stMarkdown { color: rgba(255,255,255,0.82) !important; }
    
    /* Primary button: gradient + hover lift */
    .stButton > button {
        background: linear-gradient(135deg, #8b5cf6 0%, #6d28d9 100%) !important;
        border: none !important;
        border-radius: 16px !important;
        padding: 0.8rem 2rem !important;
        font-size: 1.1rem !important;
        transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
        box-shadow: 0 8px 20px rgba(139, 92, 246, 0.3) !important;
    }
    .stButton > button:hover {
        transform: translateY(-3px) scale(1.02);
        box-shadow: 0 12px 25px rgba(139, 92, 246, 0.5) !important;
    }
    
    /* Unified Single-Card Styling with Background Animation */
    @keyframes cardGradient {
        0% { background-position: 0% 50%; }
        50% { background-position: 100% 50%; }
        100% { background-position: 0% 50%; }
    }
    @keyframes resultFadeIn {
        from { opacity: 0; transform: translateY(10px); }
        to { opacity: 1; transform: translateY(0); }
    }
    
    /* Structural Fix: Target the Horizontal Block (Columns Container) for the Unified Card */
    [data-testid="stHorizontalBlock"] {
        background: linear-gradient(135deg, rgba(30, 30, 45, 0.8), rgba(15, 15, 25, 0.9));
        background-size: 200% 200%;
        animation: cardGradient 15s ease infinite;
        backdrop-filter: blur(30px);
        -webkit-backdrop-filter: blur(30px);
        border: 1px solid rgba(255, 255, 255, 0.12);
        border-radius: 32px;
        padding: 3rem !important;
        box-shadow: 0 25px 60px rgba(0, 0, 0, 0.6);
        margin: 2rem 0 !important;
    }
    
    /* Ensure individual columns don't have their own cards */
    [data-testid="column"] {
        background: none !important;
        border: none !important;
        box-shadow: none !important;
        padding: 0 1.5rem !important;
    }

    /* Primary button micro-interactions */
    .stButton > button {
        background: linear-gradient(135deg, #8b5cf6 0%, #6d28d9 100%) !important;
        border: none !important;
        border-radius: 12px !important;
        padding: 0.8rem 2rem !important;
        font-weight: 700 !important;
        letter-spacing: 0.05em !important;
        text-transform: uppercase !important;
        transition: all 0.3s ease !important;
        box-shadow: 0 4px 15px rgba(139, 92, 246, 0.3) !important;
    }
    .stButton > button:hover {
        transform: scale(1.03) translateY(-2px) !important;
        box-shadow: 0 8px 25px rgba(139, 92, 246, 0.5) !important;
    }
    .stButton > button:active {
        transform: scale(0.98) !important;
    }

    /* Result content animation */
    .result-content {
        animation: resultFadeIn 0.8s ease-out forwards;
    }

    /* Divider */
    hr { border-color: rgba(255,255,255,0.1) !important; }
    </style>
    """, unsafe_allow_html=True)
    # ========== End UI-only CSS ==========
    
    # Header with Letter Wave Animation
    title_text = "Resume Optimizer"
    wave_html = '<h1 class="main-title"><span class="emoji">📝</span>'
    for i, char in enumerate(title_text):
        if char == " ":
            wave_html += '&nbsp;'
        else:
            delay = i * 0.1
            wave_html += f'<span class="wave-char" style="animation-delay: {delay}s">{char}</span>'
    wave_html += '</h1>'
    
    st.markdown(wave_html, unsafe_allow_html=True)
    st.caption("Powered by Resoptimum")
    
    # Sidebar for configuration
    with st.sidebar:
        st.markdown('<div class="logo-container">', unsafe_allow_html=True)
        st.image("./Resoptimum.png", width=220)
        st.markdown('</div>', unsafe_allow_html=True)

        # Model selection (Nebius Token Factory models: https://docs.tokenfactory.nebius.com)
        generative_model = st.selectbox(
            "Generative Model",
            [
                "meta-llama/Llama-3.3-70B-Instruct",
                "meta-llama/Meta-Llama-3.1-70B-Instruct",
                "meta-llama/Meta-Llama-3.1-8B-Instruct",
            ],
            index=0
        )
        
        st.divider()
        
        # Resume upload
        st.subheader("Upload Resume")
        uploaded_file = st.file_uploader(
            "Choose your resume (PDF)",
            type="pdf",
            accept_multiple_files=False
        )
        
        # Handle PDF upload and processing
        if uploaded_file is not None:
            if uploaded_file != st.session_state.current_pdf:
                st.session_state.current_pdf = uploaded_file
                try:
                    api_key = get_resoptimum_api_key()
                    if not api_key:
                        st.error(
                            "Missing Resoptimum API key. Add NEBIUS_API_KEY=your_key to a .env file in this folder."
                        )
                        st.stop()
                    
                    # Create temporary directory for the PDF
                    if st.session_state.temp_dir:
                        shutil.rmtree(st.session_state.temp_dir)
                    st.session_state.temp_dir = tempfile.mkdtemp()
                    
                    # Save uploaded PDF to temp directory
                    file_path = os.path.join(st.session_state.temp_dir, uploaded_file.name)
                    with open(file_path, "wb") as f:
                        f.write(uploaded_file.getbuffer())
                    
                    with st.spinner("Loading Resume..."):
                        documents = SimpleDirectoryReader(st.session_state.temp_dir).load_data()
                        st.session_state.docs_loaded = True
                        st.session_state.documents = documents
                        st.success("✓ Resume loaded successfully")
                        display_pdf_preview(uploaded_file)
                except Exception as e:
                    st.error(f"Error: {str(e)}")
    
    # Main content area
    col1, col2 = st.columns([1, 1.2], gap="large")
    
    with col1:
        st.subheader("📋 Job Information")
        job_title = st.text_input(
            "Job Title", 
            placeholder="e.g. Senior Software Engineer",
            help="Enter the specific title of the role"
        )
        job_description = st.text_area(
            "Job Description", 
            height=200, 
            placeholder="Paste the job requirements here...",
            help="Include key responsibilities and requirements"
        )
        
        st.divider()
        
        st.subheader("⚙️ Optimization Options")
        optimization_type = st.selectbox(
            "Select Optimization Type",
            [
                "ATS Keyword Optimizer",
                "Experience Section Enhancer",
                "Skills Hierarchy Creator",
                "Professional Summary Crafter",
                "Education Optimizer",
                "Technical Skills Showcase",
                "Career Gap Framing"
            ],
            index=0,
            help="Choose how you want to optimize your resume"
        )
        
        # Action button with micro-interaction logic
        if st.button("Optimize Resume"):
            if not st.session_state.docs_loaded:
                st.error("Please upload your resume first")
                st.stop()
            if not job_title or not job_description:
                st.error("Please provide both job title and description")
                st.stop()
            
            # Generate optimization prompt based on selection
            prompts = {
                "ATS Keyword Optimizer": "Identify and optimize ATS keywords. Focus on exact matches and semantic variations from the job description.",
                "Experience Section Enhancer": "Enhance experience section to align with job requirements. Focus on quantifiable achievements.",
                "Skills Hierarchy Creator": "Organize skills based on job requirements. Identify gaps and development opportunities.",
                "Professional Summary Crafter": "Create a targeted professional summary highlighting relevant experience and skills.",
                "Education Optimizer": "Optimize education section to emphasize relevant qualifications for this position.",
                "Technical Skills Showcase": "Organize technical skills based on job requirements. Highlight key competencies.",
                "Career Gap Framing": "Address career gaps professionally. Focus on growth and relevant experience."
            }
            
            with st.spinner("Analyzing resume and generating suggestions..."):
                try:
                    response = run_rag_completion(
                        st.session_state.documents,
                        prompts[optimization_type],
                        job_title,
                        job_description,
                        "BAAI/bge-en-icl",
                        generative_model
                    )
                    # Remove think tags from response
                    response = response.replace("<think>", "").replace("</think>", "")
                    st.session_state.messages.append({"role": "assistant", "content": response})
                except Exception as e:
                    st.error(f"Error: {str(e)}")
    
    with col2:
        st.subheader("🎯 Optimization Results")
        if not st.session_state.messages:
            st.info("Optimization results will appear here...")
        else:
            # Wrap results in div for fade-in animation
            results_html = "".join([f'<div class="result-content">{m["content"]}</div>' for m in st.session_state.messages])
            st.markdown(f'<div class="result-container">{results_html}</div>', unsafe_allow_html=True)
            

if __name__ == "__main__":
    main()