import gradio as gr
from langchain_community.document_loaders import PyPDFLoader, UnstructuredFileLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain.chains import LLMChain
from langchain_groq import ChatGroq
from langchain.prompts import PromptTemplate
from typing import List, Dict
import os
import tempfile
import re
from usage_logging import record_visit

# Initialize embeddings
embeddings = HuggingFaceEmbeddings()

# Initialize separate vector stores for resumes and culture docs
resume_store = Chroma(
    collection_name="resumes",
    embedding_function=embeddings,
    persist_directory="./chroma_db"
)

culture_store = Chroma(
    collection_name="culture_docs",
    embedding_function=embeddings,
    persist_directory="./chroma_db"
)

# Initialize LLM
llm = ChatGroq(
    api_key=os.environ["GROQ_API_KEY"],
    model_name="openai/gpt-oss-120b",
    temperature = 0,seed = 42
)

def anonymize_resume_text(text: str):
    """
    Heuristic redaction to remove common personal identifiers from resumes
    (email, phone, URLs, addresses, demographic fields, and likely name header).
    Returns: (sanitized_text, redaction_notes_list)
    """
    redactions = []
    sanitized = text

    # Email addresses
    sanitized2 = re.sub(r'[\w\.-]+@[\w\.-]+\.\w+', '[REDACTED_EMAIL]', sanitized)
    if sanitized2 != sanitized:
        redactions.append("Email addresses removed")
    sanitized = sanitized2

    # Phone numbers (broad heuristic)
    sanitized2 = re.sub(r'(\+?\d[\d\-\(\)\s]{7,}\d)', '[REDACTED_PHONE]', sanitized)
    if sanitized2 != sanitized:
        redactions.append("Phone numbers removed")
    sanitized = sanitized2

    # URLs
    sanitized2 = re.sub(r'(https?://\S+|www\.\S+)', '[REDACTED_URL]', sanitized)
    if sanitized2 != sanitized:
        redactions.append("URLs removed")
    sanitized = sanitized2

    # Physical addresses (heuristic)
    address_patterns = [
        r'\b\d{1,6}\s+\w+(?:\s+\w+){0,4}\s+(Street|St|Avenue|Ave|Road|Rd|Boulevard|Blvd|Lane|Ln|Drive|Dr|Court|Ct|Way|Parkway|Pkwy)\b\.?',
        r'\b(Apt|Apartment|Unit|Suite|Ste)\s*#?\s*\w+\b',
        r'\b\d{5}(?:-\d{4})?\b'  # US ZIP
    ]
    for pat in address_patterns:
        sanitized2 = re.sub(pat, '[REDACTED_ADDRESS]', sanitized, flags=re.IGNORECASE)
        if sanitized2 != sanitized:
            redactions.append("Address/location identifiers removed")
        sanitized = sanitized2

    # Explicit demographic fields
    demographic_patterns = [
        r'\b(gender|sex)\s*:\s*\w+\b',
        r'\b(age)\s*:\s*\d+\b',
        r'\b(dob|date of birth)\s*:\s*[\w\s,/-]+\b',
        r'\b(marital status)\s*:\s*\w+\b',
        r'\b(nationality)\s*:\s*\w+\b',
        r'\b(citizenship)\s*:\s*[\w\s,/-]+\b',
        r'\b(pronouns?)\s*:\s*[\w/]+\b',
    ]
    for pat in demographic_patterns:
        sanitized2 = re.sub(pat, '[REDACTED_DEMOGRAPHIC]', sanitized, flags=re.IGNORECASE)
        if sanitized2 != sanitized:
            redactions.append("Explicit demographic fields removed")
        sanitized = sanitized2

    # Likely name header masking (first line)
    lines = sanitized.splitlines()
    if lines:
        first_line = lines[0].strip()
        if re.fullmatch(r"[A-Za-z]+(?:\s+[A-Za-z]+){1,3}", first_line):
            lines[0] = "[REDACTED_NAME]"
            sanitized = "\n".join(lines)
            redactions.append("Likely name header removed")

    # Cleanup
    sanitized = re.sub(r'\n{3,}', '\n\n', sanitized).strip()
    redactions = sorted(set(redactions))

    return sanitized, redactions


def join_loaded_docs_text(docs):
    """Combine a list of LangChain Documents into a single text blob."""
    return "\n".join([d.page_content for d in docs if getattr(d, "page_content", None)])


def process_candidate_submission(resume_file, job_description: str) -> str:
    # Load and process resume
    if resume_file.name.endswith('.pdf'):
        loader = PyPDFLoader(resume_file.name)
    else:
        loader = UnstructuredFileLoader(resume_file.name)
    
    resume_doc = loader.load()[0]
    sanitized_resume_text, _ = anonymize_resume_text(resume_doc.page_content)
    
    # Create proper prompt template
    prompt_template = PromptTemplate(
        input_variables=["resume_text", "job_description"],
        template="""
        Given the following resume and job description, create a professional cold email to the candidate:
        
        Resume:
        {resume_text}
        
        Job Description:
        {job_description}
        
        Generate a concise, compelling cold email to the candidate that highlights the candidate's relevant skills and experience, how they align with the job requirements and company. Include a strong call-to-action.
        Ensure the email is well-structured, error-free, and tailored to the specific candidate and job description. Do not include any text apart from the email content.
        """
    )
    
    chain = LLMChain(
        llm=llm,
        prompt=prompt_template
    )
    
    response = chain.run({
        "resume_text": sanitized_resume_text,
        "job_description": job_description
    })
    
    return response

def store_culture_docs(culture_files: List[tempfile._TemporaryFileWrapper]) -> str:
    """Store company culture documentation in the vector store"""
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=100
    )
    
    all_docs = []
    for file in culture_files:
        if file.name.endswith('.pdf'):
            loader = PyPDFLoader(file.name)
        else:
            loader = UnstructuredFileLoader(file.name)
        docs = loader.load()
        splits = text_splitter.split_documents(docs)
        all_docs.extend(splits)
    
    culture_store.add_documents(all_docs)
    return f"Successfully stored {len(all_docs)} culture document chunks"

def store_resumes(resume_files: List[tempfile._TemporaryFileWrapper]) -> str:
    """Store resumes in the vector store with proper metadata"""
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=100
    )
    
    all_docs = []
    for file in resume_files:
        if file.name.endswith('.pdf'):
            loader = PyPDFLoader(file.name)
        else:
            loader = UnstructuredFileLoader(file.name)
        docs = loader.load()
        
        # Combine + anonymize before splitting
        raw_text = join_loaded_docs_text(docs)
        sanitized_text, redactions = anonymize_resume_text(raw_text)

        # Create a single Document to split
        from langchain.schema import Document
        base_doc = Document(page_content=sanitized_text, metadata={})
        
        # Extract filename without extension as resume ID
        resume_id = os.path.splitext(os.path.basename(file.name))[0]
        
        # Add metadata to each chunk
        splits = text_splitter.split_documents([base_doc])
        for split in splits:
            split.metadata["resume_id"] = resume_id
            split.metadata["source"] = "resume"
            split.metadata["sanitized"] = True
        
        all_docs.extend(splits)
    
    resume_store.add_documents(all_docs)
    return f"Successfully stored {len(resume_files)} resumes"

def verify_analysis(analysis_text: str, source_documents: List[str]) -> Dict:
    """Verify analysis against source documents with specific claim verification"""
    verification_prompt = PromptTemplate(
        input_variables=["analysis", "source_docs"],
        template="""
        You are a fact-checker. Compare the analysis below against the source documents.
        
        Analysis to verify:
        {analysis}
        
        Source documents:
        {source_docs}
        
        For each specific claim made in the analysis, verify if it's supported by the source documents.
        
        Format your response as:
        VERIFIED CLAIMS:
        ✓ [claim that is supported by source documents]
        ✗ [claim that is NOT supported by source documents]
        
        FACTUALITY SCORE: [0-100]
        """
    )
    
    chain = LLMChain(llm=llm, prompt=verification_prompt)
    result = chain.run({
        "analysis": analysis_text,
        "source_docs": "\n---\n".join(source_documents)
    })
    
    # Parse score and claims
    try:
        score_match = re.search(r'FACTUALITY SCORE:\s*(\d+)', result)
        score = int(score_match.group(1)) / 100 if score_match else 0.75
        
        # Extract verified and unverified claims
        verified_claims = re.findall(r'✓\s*(.+)', result)
        unverified_claims = re.findall(r'✗\s*(.+)', result)
        
    except:
        score = 0.75
        verified_claims = []
        unverified_claims = []
    
    return {
        "factuality_score": score, 
        "verified_claims": verified_claims,
        "unverified_claims": unverified_claims,
        "verification_result": result
    }

def self_correct_recommendation(original_recommendation: str, verification_issues: List[str], source_docs: List[str]) -> str:
    """Have LLM revise its recommendation based on verification feedback"""
    correction_prompt = PromptTemplate(
        input_variables=["original_rec", "issues", "source_docs"],
        template="""
        Your original hiring recommendation contained some unverified claims. Please revise it.
        
        Original Recommendation:
        {original_rec}
        
        Verification Issues Found:
        {issues}
        
        Source Documents:
        {source_docs}
        
        Provide a REVISED recommendation that:
        1. Removes or corrects unverified claims
        2. Bases conclusions only on documented evidence
        3. Maintains the same format and decision structure
        
        REVISED HIRING RECOMMENDATION:
        """
    )
    
    chain = LLMChain(llm=llm, prompt=correction_prompt)
    return chain.run({
        "original_rec": original_recommendation,
        "issues": "\n".join(verification_issues),
        "source_docs": "\n---\n".join(source_docs)
    })

bias_audit_prompt = PromptTemplate(
    input_variables=["skills_analysis", "culture_analysis", "final_recommendation", "job_desc", "culture_docs"],
    template="""Review the following candidate evaluation for potential bias:

SKILLS ANALYSIS:
{skills_analysis}

CULTURE ANALYSIS:
{culture_analysis}

FINAL RECOMMENDATION:
{final_recommendation}

REFERENCE MATERIALS (source of truth):
Job Description:
{job_desc}

Culture Documents:
{culture_docs}

Check specifically for:
- Over-reliance on education pedigree or past employers over actual skills
- Penalizing nontraditional career paths
- Use of subjective or exclusionary language in cultural fit
- Reasoning not supported by job description or culture documents

Output format (exactly):
BIAS AUDIT RESULT:
- Bias Indicators: [List any concerns or 'None Detected']
- Transparency Note: [Short note for recruiter if concerns exist]
"""
)

def run_bias_audit(skills_analysis, culture_analysis, final_recommendation, job_desc, culture_docs):
    chain = LLMChain(llm=llm, prompt=bias_audit_prompt)
    return chain.run({
        "skills_analysis": skills_analysis,
        "culture_analysis": culture_analysis,
        "final_recommendation": final_recommendation,
        "job_desc": job_desc,
        "culture_docs": culture_docs
    })


def analyze_candidates(job_description: str) -> str:
    # First extract required skills from job description
    skills_prompt = PromptTemplate(
        input_variables=["job_description"],
        template="""
        Extract the key technical skills and requirements from this job description:
        
        {job_description}
        
        Return the skills as a comma-separated list.
        """
    )
    
    skills_chain = LLMChain(
        llm=llm,
        prompt=skills_prompt
    )
    
    skills = skills_chain.run({"job_description": job_description})
    
    # Get relevant culture documents based on job description
    relevant_culture_docs = culture_store.similarity_search(
        job_description,  # Using job description to find relevant culture aspects
        k=3  # Adjust based on how many culture chunks you want to consider
    )
    culture_context = "\n".join([doc.page_content for doc in relevant_culture_docs])
    
    # First analyze what cultural aspects we're looking for based on the role
    culture_requirements_prompt = PromptTemplate(
        input_variables=["job_description", "culture_docs"],
        template="""
        Based on this job description and our company culture documents, identify the key cultural attributes we should look for in candidates:

        Job Description:
        {job_description}

        Relevant Company Culture Context:
        {culture_docs}

        List the top 3-5 cultural attributes that would make someone successful in this role at our company:
        """
    )
    
    culture_req_chain = LLMChain(
        llm=llm,
        prompt=culture_requirements_prompt
    )
    
    cultural_requirements = culture_req_chain.run({
        "job_description": job_description,
        "culture_docs": culture_context
    })
    
    # Query resumes
    results = resume_store.similarity_search(
        job_description,
        k=10
    )
    
    # Group resume chunks by resume_id
    resume_groups = {}
    for doc in results:
        resume_id = doc.metadata.get("resume_id")
        if resume_id not in resume_groups:
            resume_groups[resume_id] = []
        resume_groups[resume_id].append(doc.page_content)
    
    # For each resume, compare against culture docs
    consolidated_analyses = []  # Initialize empty list for all analyses
    for resume_id, chunks in resume_groups.items():
        resume_text = "\n".join(chunks)
        
        # Compare this specific resume against culture docs
        culture_analysis_prompt = PromptTemplate(
            input_variables=["resume", "cultural_requirements"],
            template="""
            Analyze this candidate's potential culture fit based on their resume and company culture, both attached below.
            
            Resume:
            {resume}
            
            Key Cultural Requirements for this Role:
            {cultural_requirements}
            
            For this candidate, provide:
            1. Concise cultural fit assessment with pros and cons. Only provide a summary.
            2. Overall culture fit score (0-100%)
            3. Recommendation on cultural fit (e.g. "Strong fit", "Moderate fit", "Not a fit")
            4. Brief explanation of your recommendation.

            Review your response, ensure it is concise and within 200 words.
            """
        )
        
        culture_chain = LLMChain(
            llm=llm,
            prompt=culture_analysis_prompt
        )
        
        try:
            culture_fit = culture_chain.run({
                "resume": resume_text,
                "cultural_requirements": cultural_requirements
            })
            
            # Verify culture analysis
            culture_verification = verify_analysis(culture_fit, [resume_text, cultural_requirements])
            
            # Now analyze technical skills match
            skills_analysis_prompt = PromptTemplate(
                input_variables=["resume", "required_skills", "job_description"],
                template="""
                Analyze this candidate's technical skills match for the position, using the resume, required skills and job description provided below.
                
                Resume:
                {resume}
                
                Required Skills:
                {required_skills}
                
                Job Description:
                {job_description}
                
                For this candidate, provide:
                1. Concise skills match assessment with pros and cons. Only provide a summary.
                2. Overall skill match score (0-100%)
                3. Recommendation on skill fit (e.g. "Strong fit", "Moderate fit", "Not a fit")
                4. Brief explanation of your recommendation

                Review your response, ensure it is concise and within 200 words.
                """
            )
            
            skills_chain = LLMChain(
                llm=llm,
                prompt=skills_analysis_prompt
            )
            
            skills_fit = skills_chain.run({
                "resume": resume_text,
                "required_skills": skills,
                "job_description": job_description
            })

            # Verify skills analysis
            skills_verification = verify_analysis(skills_fit, [resume_text, skills, job_description])

            # Create final recommendation
            final_recommendation_prompt = PromptTemplate(
                input_variables=["skills_analysis", "culture_analysis", "job_description"],
                template="""
                Provide a final hiring recommendation, using the job description, technical skills analysis, and culture fit analysis provided below.

                Job Description:
                {job_description}

                Technical Skills Analysis:
                {skills_analysis}

                Culture Fit Analysis:
                {culture_analysis}

                Provide your recommendation in the following format:

                FINAL HIRING RECOMMENDATION:
                Decision: [PROCEED / DO NOT PROCEED]

                Rationale:
                [Concise explanation of the recommendation]

                Review your final recommendation, be cut throat and make a data driven decision. For senior technical roles, give more importance to skills over culture fit.
                Ensure this response is concise and within 200 words.
                """
            )

            recommendation_chain = LLMChain(
                llm=llm,
                prompt=final_recommendation_prompt
            )

            final_recommendation = recommendation_chain.run({
                "skills_analysis": skills_fit,
                "culture_analysis": culture_fit,
                "job_description": job_description
            })

            # Self-correction integration - INSERT HERE (after line 364)
            # Collect all unverified claims for potential self-correction
            all_issues = culture_verification["unverified_claims"] + skills_verification["unverified_claims"]
            
            # Self-correct if verification scores are low AND there are actual issues
            if (culture_verification["factuality_score"] < 0.95 or skills_verification["factuality_score"] < 0.95) and all_issues:
                corrected_recommendation = self_correct_recommendation(
                    final_recommendation, 
                    all_issues,
                    [resume_text, job_description, cultural_requirements]
                )
                final_recommendation = corrected_recommendation
                revision_note = "\n\n🔄 RECOMMENDATION REVISED: Corrected based on verification feedback"
            else:
                revision_note = ""

            # Bias audit (triangulates across skills, culture, and final recommendation)
            bias_audit = run_bias_audit(
                skills_analysis=skills_fit,
                culture_analysis=culture_fit,
                final_recommendation=final_recommendation,
                job_desc=job_description,
                culture_docs=culture_context
            )

            # Add verification warnings if factuality score < 0.95
            verification_notes = ""
            if culture_verification["factuality_score"] < 0.95 or skills_verification["factuality_score"] < 0.95:
                verification_notes = "\n\n🔍 FACT CHECK RESULTS:"
    
                if culture_verification["unverified_claims"]:
                    verification_notes += f"\n\nCULTURE ANALYSIS - Unverified claims:"
                    for claim in culture_verification["unverified_claims"][:3]:
                        verification_notes += f"\n✗ {claim}"
    
                if skills_verification["unverified_claims"]:
                    verification_notes += f"\n\nSKILLS ANALYSIS - Unverified claims:"
                    for claim in skills_verification["unverified_claims"][:3]:
                        verification_notes += f"\n✗ {claim}"
            
            # Append the analysis for this candidate to the consolidated analyses
            consolidated_analyses.append(f"""
            === Candidate Analysis (Resume ID: {resume_id}) ===
            
            CULTURE FIT ANALYSIS:
            {culture_fit}
            
            TECHNICAL SKILLS ANALYSIS:
            {skills_fit}

            HIRING RECOMMENDATION:
            {final_recommendation}{revision_note}{verification_notes}

            BIAS AUDIT:
            {bias_audit}
            
            ----------------------------------------
            """)
            
        except Exception as e:
            # If there's an error analyzing this candidate, add error message but continue with others
            consolidated_analyses.append(f"""
            === Candidate Analysis (Resume ID: {resume_id}) ===
            Error analyzing candidate: {str(e)}
            ----------------------------------------
            """)
            continue
    
    # Return all analyses joined together
    return "\n".join(consolidated_analyses)


def clear_databases():
    """Clear both resume and culture document databases"""
    global resume_store, culture_store
    
    status_messages = []
    
    # Clear resume store
    try:
        results = resume_store.get()
        if results and results['ids']:
            num_docs = len(results['ids'])
            resume_store._collection.delete(
                ids=results['ids']
            )
            status_messages.append(f"Cleared {num_docs} documents from resume database")
        else:
            status_messages.append("Resume database was already empty")
    except Exception as e:
        status_messages.append(f"Error clearing resume store: {e}")
        
    # Clear culture store
    try:
        results = culture_store.get()
        if results and results['ids']:
            num_docs = len(results['ids'])
            culture_store._collection.delete(
                ids=results['ids']
            )
            status_messages.append(f"Cleared {num_docs} documents from culture database")
        else:
            status_messages.append("Culture database was already empty")
    except Exception as e:
        status_messages.append(f"Error clearing culture store: {e}")
    
    return "\n".join(status_messages)




def create_interface():
    with gr.Blocks(theme='freddyaboulton/test-blue') as app:
        app.load(fn=record_visit, inputs=None, outputs=None)
        gr.Markdown("# AI Recruiter Assistant")
        gr.Markdown("""**Purpose**  
        This prototype demonstrates how AI can support recruiting workflows — including candidate evaluation and outreach — while embedding safeguards for fairness, transparency, and verification.
        It is designed as a **decision-support tool**, not an automated decision-maker.
        
        ⚠️ **Important Disclaimer**
  
        This tool does **not** replace recruiter judgment, legal review, or hiring policy compliance. Final hiring decisions must always be made by humans.
        
        💬 **Feedback Welcome**
        
        Please share feedback, issues, or improvement ideas via the **Community** tab.
        """)
        with gr.Tabs():

            # Recruiter View
            with gr.Tab("Candidate Assessment"):
                gr.Markdown("Clear existing culture documents and resumes from storage. Use this every time you are uploading new company documentation or do not want to select from the existing pool of resumes.")
                clear_btn = gr.Button("Clear All Databases")
                clear_status = gr.Textbox(label="Clear Status")
                gr.Markdown("💡 Tip: A sample resume, culture document and job description are available in the **Files** section of this space for testing.")
                gr.Markdown("Use this feature to upload company culture documents (values, principles, leadership philosophy). These documents will be used to assess the cultural fit of candidates.")
                with gr.Row():
                    culture_docs_upload = gr.File(
                        label="Upload Company Culture Documents",
                        file_count="multiple"
                    )
                    store_culture_btn = gr.Button("Store Culture Docs")
                    culture_status = gr.Textbox(label="Status")
                gr.Markdown("Use this feature to upload resumes in bulk (Word or PDF). Each resume is anonymized before analysis. These resumes will be used to assess the technical skills and culture fit of candidates.")
                with gr.Row():
                    resume_bulk_upload = gr.File(
                        label="Upload Resumes",
                        file_count="multiple"
                    )
                    store_resumes_btn = gr.Button("Store Resumes")
                    resume_status = gr.Textbox(label="Status")
                
                with gr.Row():
                    job_desc_recruiter = gr.Textbox(
                        label="Paste the job description for the role you are hiring for.",
                        lines=20
                    )
                with gr.Row():
                    analyze_btn = gr.Button("Analyze Candidates")
                with gr.Row():
                    analysis_output = gr.Textbox(
                        label="Analysis Results",
                        lines=30
                    )
                
                store_culture_btn.click(
                    store_culture_docs,
                    inputs=culture_docs_upload,
                    outputs=culture_status
                )
                
                store_resumes_btn.click(
                    store_resumes,
                    inputs=resume_bulk_upload,
                    outputs=resume_status
                )
                
                analyze_btn.click(
                    analyze_candidates,
                    inputs=job_desc_recruiter,
                    outputs=analysis_output
                )

                clear_btn.click(
                    clear_databases,
                    inputs=[],
                    outputs=clear_status
                )
             # Candidate View
            with gr.Tab("Cold Email Generator"):
                with gr.Row():
                    resume_upload = gr.File(label="Upload Resume")
                    job_desc_input = gr.Textbox(
                        label="Paste Job Description",
                        lines=10
                    )
                generate_btn = gr.Button("Generate Cold Email")
                email_output = gr.Textbox(
                    label="Generated Cold Email",
                    lines=10
                )
                
                generate_btn.click(
                    process_candidate_submission,
                    inputs=[resume_upload, job_desc_input],
                    outputs=email_output
                )
    
    return app

if __name__ == "__main__":
    app = create_interface()
    app.launch()