import streamlit as st
import pandas as pd
import json
import os
from pydantic import BaseModel, Field
from typing import List, Set, Dict, Any, Optional # Already have these, but commented for brevity if not all used
import time # Added for potential small delays if needed
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage # Not directly used in provided snippet
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser # Not directly used in provided snippet
from langchain_core.prompts import PromptTemplate # Not directly used in provided snippet
import gspread
import tempfile
import time
from google.oauth2 import service_account
import tiktoken

st.set_page_config(
    page_title="Candidate Matching App",
    page_icon="👨‍💻🎯",
    layout="wide"
)
os.environ["STREAMLIT_HOME"] = tempfile.gettempdir()
os.environ["STREAMLIT_DISABLE_TELEMETRY"] = "1"
# Define pydantic model for structured output
class Shortlist(BaseModel):
    fit_score: float = Field(description="A score between 0 and 10 indicating how closely the candidate profile matches the job requirements upto 3 decimal points.")
    candidate_name: str = Field(description="The name of the candidate.")
    candidate_url: str = Field(description="The URL of the candidate's LinkedIn profile.")
    candidate_summary: str = Field(description="A brief summary of the candidate's skills and experience along with its educational background.")
    candidate_location: str = Field(description="The location of the candidate.")
    justification: str = Field(description="Justification for the shortlisted candidate with the fit score")

# Function to calculate tokens
def calculate_tokens(text, model="gpt-4o-mini"):
    try:
        if "gpt-4" in model:
            encoding = tiktoken.encoding_for_model("gpt-4o-mini")
        elif "gpt-3.5" in model:
            encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
        else:
            encoding = tiktoken.get_encoding("cl100k_base")
        return len(encoding.encode(text))
    except Exception as e:
        return len(text) // 4

# Function to display token usage
def display_token_usage():
    if 'total_input_tokens' not in st.session_state:
        st.session_state.total_input_tokens = 0
    if 'total_output_tokens' not in st.session_state:
        st.session_state.total_output_tokens = 0
    
    total_input = st.session_state.total_input_tokens
    total_output = st.session_state.total_output_tokens
    total_tokens = total_input + total_output
    
    model_to_check = st.session_state.get('model_name', "gpt-4o-mini") # Use a default if not set

    if model_to_check == "gpt-4o-mini":
        input_cost_per_1k = 0.00015 # Adjusted to example rates ($0.15 / 1M tokens)
        output_cost_per_1k = 0.0006  # Adjusted to example rates ($0.60 / 1M tokens)
    elif "gpt-4" in model_to_check: # Fallback for other gpt-4
        input_cost_per_1k = 0.005 
        output_cost_per_1k = 0.015 # General gpt-4 pricing can vary
    else:  # Assume gpt-3.5-turbo pricing
        input_cost_per_1k = 0.0005  # $0.0005 per 1K input tokens
        output_cost_per_1k = 0.0015   # $0.0015 per 1K output tokens
    
    estimated_cost = (total_input / 1000 * input_cost_per_1k) + (total_output / 1000 * output_cost_per_1k)
    
    st.subheader("📊 Token Usage Statistics (for last processed job)")
    
    col1, col2, col3 = st.columns(3)
    with col1: st.metric("Input Tokens", f"{total_input:,}")
    with col2: st.metric("Output Tokens", f"{total_output:,}")
    with col3: st.metric("Total Tokens", f"{total_tokens:,}")
    st.markdown(f"**Estimated Cost:** ${estimated_cost:.4f}")
    return total_tokens

# Function to parse and normalize tech stacks
def parse_tech_stack(stack):
    if pd.isna(stack) or stack == "" or stack is None: return set()
    if isinstance(stack, set): return stack
    try:
        if isinstance(stack, str) and stack.startswith("{") and stack.endswith("}"):
            items = stack.strip("{}").split(",")
            return set(item.strip().strip("'\"") for item in items if item.strip())
        return set(map(lambda x: x.strip().lower(), str(stack).split(',')))
    except Exception as e:
        st.error(f"Error parsing tech stack: {e}")
        return set()

def display_tech_stack(stack_set):
    return ", ".join(sorted(list(stack_set))) if isinstance(stack_set, set) else str(stack_set)


def get_matching_candidates(job_stack, candidates_df):
    matched = []
    job_stack_set = parse_tech_stack(job_stack)
    for _, candidate in candidates_df.iterrows():
        candidate_stack = parse_tech_stack(candidate['Key Tech Stack'])
        common = job_stack_set & candidate_stack
        if len(common) >= 2: # Original condition
            matched.append({
                "Name": candidate["Full Name"], "URL": candidate["LinkedIn URL"],
                "Degree & Education": candidate["Degree & University"],
                "Years of Experience": candidate["Years of Experience"],
                "Current Title & Company": candidate['Current Title & Company'],
                "Key Highlights": candidate["Key Highlights"],
                "Location": candidate["Location (from most recent experience)"],
                "Experience": str(candidate["Experience"]), "Tech Stack": candidate_stack
            })
    return matched

def setup_llm():
    """Set up the LangChain LLM with structured output"""
    # Define the model to use
    model_name = "gpt-4o-mini"
    
    # Store model name in session state for token calculation
    if 'model_name' not in st.session_state:
        st.session_state.model_name = model_name
    
    # Create LLM instance
    llm = ChatOpenAI(
        model=model_name,
        temperature=0.3,
        max_tokens=None,
        timeout=None,
        max_retries=2,
    )
    
    # Create structured output
    sum_llm = llm.with_structured_output(Shortlist)
    
    # Create system prompt
    system = """You are an expert Tech Recruiter. For each candidate–job pair, follow these steps and show your chain of thought before giving a final Fit Score (0–10):

1. LOCATION CHECK (Hard Disqualification)  
   - If candidate’s location lies outside the job’s required location, immediately reject (Score 1–5) with reasoning “Location mismatch.”   

2. HARD DISQUALIFICATIONS (Auto-reject, Score 1–5)  
   - No VC-backed startup experience (Seed–Series C/D)  
   - Only Big Tech or corporate labs, with no startup follow-on  
   - < 3 years post-graduate SWE experience  
   - More than one role < 2 years (unless due to M&A or shutdown)  
   - Career centered on enterprise/consulting firms (e.g., Infosys, Wipro, Cognizant, Tata, Capgemini, Dell, Cisco)  
   - Visa dependency (H1B/OPT/TN) unless explicitly allowed   

3. EDUCATION & STARTUP EXPERIENCE SCORING  
   - **Tier 1 (Max points):** MIT, Stanford, CMU, UC Berkeley, Caltech, Harvard, IIT Bombay, IIT Delhi, Princeton, UIUC, UW, Columbia, UChicago, Cornell, UM-Ann Arbor, UT Austin, Waterloo, U Toronto  
   - **Tier 2 (Moderate points):** UC Davis, Georgia Tech, Purdue, UMass Amherst, etc.  
   - **Tier 3 (Low points):** Other or unranked institutions  
   - Assume CS degree for all; use university field to assign tier 
   - Validate startup’s funding stage via Crunchbase/Pitchbook; preferred investors include YC, Sequoia, a16z, Accel, Founders Fund, Lightspeed, Greylock, Benchmark, Index Ventures  

4. WEIGHTED FIT SCORE COMPONENTS (Qualified candidates only)  
   - Engineering & Problem Solving: 20%  
   - Product Experience (built systems end-to-end): 20%  
   - Startup Experience (time at VC-backed roles): 20%  
   - Tech Stack Alignment: 15%  
   - Tenure & Stability (≥ 2 years per role): 15%  
   - Domain Relevance (industry match): 10% :  

5. ADJACENT COMPANY MATCHING  
   - If startup funding can’t be verified, suggest similar-stage companies in the same market and justify  

**Output:**  
- **Chain of Thought:** bullet points for each step above  
- **Final Fit Score:** X.X/10 and classification  
  - 1–5: Poor Fit (Auto-reject)  
  - 6–7: Weak Fit (Auto-reject)  
  - 8.0–8.7: Moderate Fit (Auto-reject)  
  - 8.8–10: Strong Fit (Include in results)  
    """
    
    # Create query prompt
    query_prompt = ChatPromptTemplate.from_messages([
        ("system", system),
        ("human", """
    You are an expert Recruitor. Your task is to determine if the candidate matches the given job. 
Provide the score as a `float` rounded to exactly **three decimal places** (e.g., 8.943, 9.211, etc.). 
Avoid rounding to whole or one-decimal numbers. Every candidate should have a **unique** fit score.
          For this you will be provided with the follwing inputs of job and candidates:
    Job Details
    Company: {Company}
    Role: {Role}
    About Company: {desc}
    Locations: {Locations}
    Tech Stack: {Tech_Stack}
    Industry: {Industry}
    
    Candidate Details:
    Full Name: {Full_Name}
    LinkedIn URL: {LinkedIn_URL}
    Current Title & Company: {Current_Title_Company}
    Years of Experience: {Years_of_Experience}
    Degree & University: {Degree_University}
    Key Tech Stack: {Key_Tech_Stack}
    Key Highlights: {Key_Highlights}
    Location (from most recent experience): {cand_Location}
    Past_Experience: {Experience}
    Answer in the structured manner as per the schema.
    If any parameter is Unknown try not to include in the summary, only include those parameters which are known.
    The `fit_score` must be a float with **exactly three decimal digits** (e.g. 8.812, 9.006). Do not round to 1 or 2 decimals.
    """),
    ])
    
    # Chain the prompt and LLM
    cat_class = query_prompt | sum_llm
    
    return cat_class

def call_llm(candidate_data, job_data, llm_chain):
    try:
        job_tech_stack = ", ".join(sorted(list(job_data.get("Tech_Stack", set())))) if isinstance(job_data.get("Tech_Stack"), set) else job_data.get("Tech_Stack", "")
        candidate_tech_stack = ", ".join(sorted(list(candidate_data.get("Tech Stack", set())))) if isinstance(candidate_data.get("Tech Stack"), set) else candidate_data.get("Tech Stack", "")
        
        payload = {
            "Company": job_data.get("Company", ""), "Role": job_data.get("Role", ""),
            "desc": job_data.get("desc", ""), "Locations": job_data.get("Locations", ""),
            "Tech_Stack": job_tech_stack, "Industry": job_data.get("Industry", ""),
            "Full_Name": candidate_data.get("Name", ""), "LinkedIn_URL": candidate_data.get("URL", ""),
            "Current_Title_Company": candidate_data.get("Current Title & Company", ""),
            "Years_of_Experience": candidate_data.get("Years of Experience", ""),
            "Degree_University": candidate_data.get("Degree & Education", ""),
            "Key_Tech_Stack": candidate_tech_stack, "Key_Highlights": candidate_data.get("Key Highlights", ""),
            "cand_Location": candidate_data.get("Location", ""), "Experience": candidate_data.get("Experience", "")
        }
        payload_str = json.dumps(payload)
        input_tokens = calculate_tokens(payload_str, st.session_state.model_name)
        response = llm_chain.invoke(payload)
        # print(candidate_data.get("Experience", "")) # Kept for your debugging if needed

        response_str = f"candidate_name: {response.candidate_name} URL:{response.candidate_url} summ:{response.candidate_summary} loc: {response.candidate_location} just {response.justification} fit_score: {float(f'{response.fit_score:.3f}')}." # Truncated
        output_tokens = calculate_tokens(response_str, st.session_state.model_name)
        
        if 'total_input_tokens' not in st.session_state: st.session_state.total_input_tokens = 0
        if 'total_output_tokens' not in st.session_state: st.session_state.total_output_tokens = 0
        st.session_state.total_input_tokens += input_tokens
        st.session_state.total_output_tokens += output_tokens
        
        return {
            "candidate_name": response.candidate_name, "candidate_url": response.candidate_url,
            "candidate_summary": response.candidate_summary, "candidate_location": response.candidate_location,
            "fit_score": response.fit_score, "justification": response.justification
        }
    except Exception as e:
        st.error(f"Error calling LLM for {candidate_data.get('Name', 'Unknown')}: {e}")
        return {
            "candidate_name": candidate_data.get("Name", "Unknown"), "candidate_url": candidate_data.get("URL", ""),
            "candidate_summary": "Error processing candidate profile", "candidate_location": candidate_data.get("Location", "Unknown"),
            "fit_score": 0.0, "justification": f"Error in LLM processing: {str(e)}"
        }

def process_candidates_for_job(job_row, candidates_df, llm_chain=None):
    st.session_state.total_input_tokens = 0 # Reset for this job
    st.session_state.total_output_tokens = 0

    if llm_chain is None:
        with st.spinner("Setting up LLM..."): llm_chain = setup_llm()
    
    selected_candidates = []
    job_data = {
        "Company": job_row["Company"], "Role": job_row["Role"], "desc": job_row.get("One liner", ""),
        "Locations": job_row.get("Locations", ""), "Tech_Stack": job_row["Tech Stack"], "Industry": job_row.get("Industry", "")
    }
    
    with st.spinner("Finding matching candidates based on tech stack..."):
        matching_candidates = get_matching_candidates(job_row["Tech Stack"], candidates_df)
    
    if not matching_candidates:
        st.warning("No candidates with matching tech stack found for this job.")
        return []
    
    st.success(f"Found {len(matching_candidates)} candidates with matching tech stack. Evaluating with LLM...")
    
    candidates_progress = st.progress(0)
    candidate_status = st.empty() # For live updates
    
    for i, candidate_data in enumerate(matching_candidates):
        # *** MODIFICATION: Check for stop flag ***
        if st.session_state.get('stop_processing_flag', False):
            candidate_status.warning("Processing stopped by user.")
            time.sleep(1) # Allow message to be seen
            break 

        candidate_status.text(f"Evaluating candidate {i+1}/{len(matching_candidates)}: {candidate_data.get('Name', 'Unknown')}")
        response = call_llm(candidate_data, job_data, llm_chain)
        
        response_dict = {
            "Name": response["candidate_name"], "LinkedIn": response["candidate_url"],
            "summary": response["candidate_summary"], "Location": response["candidate_location"],
            "Fit Score": float(f"{response['fit_score']:.3f}"), "justification": response["justification"],
            "Educational Background": candidate_data.get("Degree & Education", ""),
            "Years of Experience": candidate_data.get("Years of Experience", ""),
            "Current Title & Company": candidate_data.get("Current Title & Company", "")
        }
        
        # *** MODIFICATION: Live output of candidate dicts - will disappear on rerun after processing ***
        if response["fit_score"] >= 8.800:
            selected_candidates.append(response_dict)
            # This st.markdown will be visible during processing and cleared on the next full script rerun
            # after this processing block finishes or is stopped.
            st.markdown(
    f"**Selected Candidate:** [{response_dict['Name']}]({response_dict['LinkedIn']}) "
    f"(Score: {response_dict['Fit Score']:.3f}, Location: {response_dict['Location']})"
)
        
        candidates_progress.progress((i + 1) / len(matching_candidates))

    candidates_progress.empty()
    candidate_status.empty()
    
    if not st.session_state.get('stop_processing_flag', False): # Only show if not stopped
        if selected_candidates:
            st.success(f"✅ LLM evaluation complete. Found {len(selected_candidates)} suitable candidates for this job!")
        else:
            st.info("LLM evaluation complete. No candidates met the minimum fit score threshold for this job.")
    
    return selected_candidates


def main():
    st.title("👨‍💻 Candidate Matching App")
    if 'processed_jobs' not in st.session_state: st.session_state.processed_jobs = {} # May not be used with new logic
    if 'Selected_Candidates' not in st.session_state: st.session_state.Selected_Candidates = {}
    if 'llm_chain' not in st.session_state: st.session_state.llm_chain = None # Initialize to None
    # *** MODIFICATION: Initialize stop flag ***
    if 'stop_processing_flag' not in st.session_state: st.session_state.stop_processing_flag = False


    st.write("This app matches job listings with candidate profiles...")
    
    with st.sidebar:
        st.header("API Configuration")
        api_key = st.text_input("Enter OpenAI API Key", type="password", key="api_key_input")
        if api_key:
            os.environ["OPENAI_API_KEY"] = api_key
            SERVICE_ACCOUNT_FILE = 'src/synapse-recruitment-e94255ca76fd.json' # Ensure this path is correct
            SCOPES = ['https://www.googleapis.com/auth/spreadsheets']
            creds = service_account.Credentials.from_service_account_file(SERVICE_ACCOUNT_FILE, scopes=SCOPES)
            # Initialize LLM chain once API key is set
            if st.session_state.llm_chain is None:
                 with st.spinner("Setting up LLM..."):
                    st.session_state.llm_chain = setup_llm()
            st.success("API Key set")
        else:
            st.warning("Please enter OpenAI API Key to use LLM features")
            st.session_state.llm_chain = None # Clear chain if key removed
            

    try:        
        gc = gspread.authorize(creds)
        job_sheet = gc.open_by_key('1BZlvbtFyiQ9Pgr_lpepDJua1ZeVEqrCLjssNd6OiG9k')
        candidates_sheet = gc.open_by_key('1u_9o5f0MPHFUSScjEcnA8Lojm4Y9m9LuWhvjYm6ytF4')
    except Exception as e:
        st.error(f"Failed to connect to Google Sheets. Please Ensure the API key is correct")
        st.stop()


    if not os.environ.get("OPENAI_API_KEY"):
        st.warning("⚠️ You need to provide an OpenAI API key in the sidebar to use this app.")
        st.stop()
    if st.session_state.llm_chain is None and os.environ.get("OPENAI_API_KEY"):
        with st.spinner("Setting up LLM..."):
            st.session_state.llm_chain = setup_llm()
        st.rerun() # Rerun to ensure LLM is ready for the main display logic

    try:
        job_worksheet = job_sheet.worksheet('paraform_jobs_formatted')
        job_data = job_worksheet.get_all_values()
        candidate_worksheet = candidates_sheet.worksheet('transformed_candidates_updated')
        candidate_data = candidate_worksheet.get_all_values()
        
        jobs_df = pd.DataFrame(job_data[1:], columns=job_data[0]).drop(["Link"], axis=1, errors='ignore')
        jobs_df1 = jobs_df[["Company","Role","One liner","Locations","Tech Stack","Workplace","Industry","YOE"]]
        jobs_df1 = jobs_df1.fillna("Unknown")
        candidates_df = pd.DataFrame(candidate_data[1:], columns=candidate_data[0]).fillna("Unknown")
        candidates_df.drop_duplicates(subset=['Full Name'], keep='first', inplace=True)
        
        with st.expander("Preview uploaded data"):
            st.subheader("Jobs Data Preview"); st.dataframe(jobs_df1.head(5))
        
        # Column mapping (simplified, ensure your CSVs have these exact names or adjust)
        # candidates_df = candidates_df.rename(columns={...}) # Add if needed
            
        display_job_selection(jobs_df, candidates_df, job_sheet) # job_sheet is 'sh'
            
    except Exception as e:
        st.error(f"Error processing files or data: {e}")
    st.divider()
    
def display_job_selection(jobs_df, candidates_df, sh):
    st.subheader("Select a job to view potential matches")
    job_options = [f"{row['Role']} at {row['Company']}" for _, row in jobs_df.iterrows()]

    if 'last_selected_job_index' not in st.session_state:
        st.session_state.last_selected_job_index = 0

    selected_job_index = st.selectbox(
        "Jobs:",
        range(len(job_options)),
        format_func=lambda x: job_options[x],
        key="job_selectbox"
    )

    # Clear previous job state when a new job is selected
    if selected_job_index != st.session_state.last_selected_job_index:
        old_job_key = st.session_state.last_selected_job_index
        job_processed_key = f"job_{old_job_key}_processed_successfully"
        job_is_processing_key = f"job_{old_job_key}_is_currently_processing"

        # Remove old job flags
        for key in [job_processed_key, job_is_processing_key, 'stop_processing_flag', 'total_input_tokens', 'total_output_tokens']:
            st.session_state.pop(key, None)

        # Clear selected candidates for old job if they exist
        if 'Selected_Candidates' in st.session_state:
            st.session_state.Selected_Candidates.pop(old_job_key, None)

        # Clear cache to avoid old data in UI
        st.cache_data.clear()

        # Update last selected job index
        st.session_state.last_selected_job_index = selected_job_index

        # Rerun to refresh UI and prevent stale data
        st.rerun()

    # Ensure Selected_Candidates is initialized for the new job
    if 'Selected_Candidates' not in st.session_state:
        st.session_state.Selected_Candidates = {}

    if selected_job_index not in st.session_state.Selected_Candidates:
        st.session_state.Selected_Candidates[selected_job_index] = []

    # Proceed with job details
    job_row = jobs_df.iloc[selected_job_index]
    job_row_stack = parse_tech_stack(job_row["Tech Stack"])

    col_job_details_display, _ = st.columns([2, 1])
    with col_job_details_display:
        st.subheader(f"Job Details: {job_row['Role']}")
        job_details_dict = {
            "Company": job_row["Company"],
            "Role": job_row["Role"],
            "Description": job_row.get("One liner", "N/A"),
            "Locations": job_row.get("Locations", "N/A"),
            "Industry": job_row.get("Industry", "N/A"),
            "Tech Stack": display_tech_stack(job_row_stack)
        }
        for key, value in job_details_dict.items():
            st.markdown(f"**{key}:** {value}")

    job_processed_key = f"job_{selected_job_index}_processed_successfully"
    job_is_processing_key = f"job_{selected_job_index}_is_currently_processing"
    st.session_state.setdefault(job_processed_key, False)
    st.session_state.setdefault(job_is_processing_key, False)

    sheet_name = f"{job_row['Role']} at {job_row['Company']}".strip()[:100]
    worksheet_exists = False
    existing_candidates_from_sheet = []
    try:
        cand_ws = sh.worksheet(sheet_name)
        worksheet_exists = True
        data = cand_ws.get_all_values()
        if len(data) > 1:
            existing_candidates_from_sheet = data
    except Exception:
        pass

    if not st.session_state[job_processed_key] or existing_candidates_from_sheet:
        col_find, col_stop = st.columns(2)
        with col_find:
            if st.button("Find Matching Candidates for this Job", key=f"find_btn_{selected_job_index}",
                         disabled=st.session_state[job_is_processing_key]):
                if not os.environ.get("OPENAI_API_KEY") or st.session_state.llm_chain is None:
                    st.error("OpenAI API key not set or LLM not initialized.")
                else:
                    st.session_state[job_is_processing_key] = True
                    st.session_state.stop_processing_flag = False
                    st.session_state.Selected_Candidates[selected_job_index] = []
                    st.rerun()
        with col_stop:
            if st.session_state[job_is_processing_key]:
                if st.button("STOP Processing", key=f"stop_btn_{selected_job_index}"):
                    st.session_state.stop_processing_flag = True
                    st.cache_data.clear()
                    st.warning("Stop request sent. Processing will halt shortly.")
                    st.rerun()

    if st.session_state[job_is_processing_key]:
        with st.spinner(f"Processing candidates for {job_row['Role']} at {job_row['Company']}..."):
            processed_list = process_candidates_for_job(job_row, candidates_df, st.session_state.llm_chain)
        st.session_state[job_is_processing_key] = False

        if not st.session_state.get('stop_processing_flag', False):
            if processed_list:
                processed_list.sort(key=lambda x: x.get("Fit Score", 0.0), reverse=True)
                st.session_state.Selected_Candidates[selected_job_index] = processed_list
                st.session_state[job_processed_key] = True
                try:
                    target_ws = sh.worksheet(sheet_name) if worksheet_exists else sh.add_worksheet(
                        title=sheet_name, rows=max(100, len(processed_list)+10), cols=20)
                    headers = list(processed_list[0].keys())
                    rows = [headers] + [[str(c.get(h, "")) for h in headers] for c in processed_list]
                    target_ws.clear()
                    target_ws.update('A1', rows)
                    st.success(f"Results saved to Google Sheet: '{sheet_name}'")
                except Exception as e:
                    st.error(f"Error writing to Google Sheet '{sheet_name}': {e}")
            else:
                st.info("No suitable candidates found after processing.")
                st.session_state.Selected_Candidates[selected_job_index] = []
                st.session_state[job_processed_key] = True
        else:
            st.info("Processing was stopped by user.")
            st.session_state[job_processed_key] = False
            st.session_state.Selected_Candidates[selected_job_index] = []
        st.session_state.pop('stop_processing_flag', None)
        st.rerun()

    should_display = False
    final_candidates = []
    if not st.session_state[job_is_processing_key]:
        if st.session_state[job_processed_key]:
            should_display = True
            final_candidates = st.session_state.Selected_Candidates.get(selected_job_index, [])
        elif existing_candidates_from_sheet:
            should_display = True
            headers = existing_candidates_from_sheet[0]
            for row in existing_candidates_from_sheet[1:]:
                cand = {headers[i]: row[i] if i < len(row) else None for i in range(len(headers))}
                try: cand['Fit Score'] = float(cand.get('Fit Score', 0))
                except: cand['Fit Score'] = 0.0
                final_candidates.append(cand)
            final_candidates.sort(key=lambda x: x.get('Fit Score', 0.0), reverse=True)

    if should_display:
        col_title, col_copyall = st.columns([3, 1])
        with col_title:
            st.subheader("Selected Candidates")
        with col_copyall:
            combined_text = ""
            for cand in final_candidates:
                combined_text += f"Name: {cand.get('Name','N/A')}\nLinkedIn URL: {cand.get('LinkedIn','N/A')}\n\n"

            import json
            html = f'''
            <button id="copy-all-btn">📋 Copy All</button>
            <script>
            const combinedText = {json.dumps(combined_text)};
            document.getElementById("copy-all-btn").onclick = () => {{
                navigator.clipboard.writeText(combinedText);
            }};
            </script>
            '''
            st.components.v1.html(html, height=60)

        if st.session_state.get(job_processed_key) and (
           st.session_state.get('total_input_tokens',0) > 0 or st.session_state.get('total_output_tokens',0) > 0):
            display_token_usage()

        for i, candidate in enumerate(final_candidates):
            score = candidate.get('Fit Score', 0.0)
            score_display = f"{score:.3f}" if isinstance(score, (int, float)) else score
            exp_title = f"{i+1}. {candidate.get('Name','N/A')} (Score: {score_display})"
            with st.expander(exp_title):
                text_copy = f"Candidate: {candidate.get('Name','N/A')}\nLinkedIn: {candidate.get('LinkedIn','N/A')}\n"
                btn = f"copy_btn_job{selected_job_index}_cand{i}"
                js = f'''
                <script>
                function copyToClipboard_{btn}() {{ navigator.clipboard.writeText(`{text_copy}`); }}
                </script>
                <button onclick="copyToClipboard_{btn}()">📋 Copy Details</button>
                '''
                cols = st.columns([0.82,0.18])
                with cols[1]: st.components.v1.html(js, height=40)
                with cols[0]:
                    st.markdown(f"**Summary:** {candidate.get('summary','N/A')}")
                    st.markdown(f"**Current:** {candidate.get('Current Title & Company','N/A')}")
                    st.markdown(f"**Education:** {candidate.get('Educational Background','N/A')}")
                    st.markdown(f"**Experience:** {candidate.get('Years of Experience','N/A')}")
                    st.markdown(f"**Location:** {candidate.get('Location','N/A')}")
                    if candidate.get('LinkedIn'):
                        st.markdown(f"**[LinkedIn Profile]({candidate['LinkedIn']})**")
                    if candidate.get('justification'):
                        st.markdown("**Justification:**")
                        st.info(candidate['justification'])

        if st.button("Reset and Process Again", key=f"reset_btn_{selected_job_index}"):
            st.session_state[job_processed_key] = False
            st.session_state.pop(job_is_processing_key, None)
            st.session_state.Selected_Candidates.pop(selected_job_index, None)
            st.cache_data.clear()
            try: sh.worksheet(sheet_name).clear()
            except: pass
            st.rerun()



if __name__ == "__main__":
    main()