import gradio as gr
from transformers import pipeline
import spacy
from docx import Document
import tempfile
import os
try:
    nlp = spacy.load('en_core_web_sm')
except IOError:
    import subprocess
    subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"])
    nlp = spacy.load('en_core_web_sm')

#nlp = spacy.load('en_core_web_sm')
# Load models quietly
#nlp = spacy.load('en_core_web_sm')
device = 0 if torch.cuda.is_available() else -1
fill_mask = pipeline("fill-mask", model="bert-base-uncased", device=device)  # Use GPU if available

# Cell 3: Define functions
def extract_text_from_docx(file_path):
    try:
        doc = Document(file_path)
        return '\n'.join([para.text for para in doc.paragraphs])
    except:
        return ""

def save_updated_resume(text):
    try:
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.docx')
        doc = Document()
        for line in text.split('\n'):
            doc.add_paragraph(line)
        doc.save(temp_file.name)
        return temp_file.name
    except:
        return None

def bert_suggest_replacements(text):
    suggestions = []
    improved_text = text
    doc = nlp(text)

    for token in doc:
        if token.is_stop or not token.is_alpha or len(token.text) <= 3:
            continue

        masked_sentence = improved_text.replace(token.text, "[MASK]", 1)
        try:
            predictions = fill_mask(masked_sentence)
            best_suggestion = predictions[0]['token_str']

            if best_suggestion != token.text.lower():
                improved_text = improved_text.replace(token.text, best_suggestion, 1)
                suggestions.append(f"• Replace '{token.text}' with '{best_suggestion}'")
        except:
            continue

    return improved_text, suggestions

def process_resume(file):
    if file is None:
        return "Please upload a resume file.", "No suggestions yet.", None

    try:
        resume_text = extract_text_from_docx(file.name)
        if not resume_text:
            return "Could not read the resume. Please ensure it's a valid .docx file.", "No suggestions available.", None

        updated_text, suggestions = bert_suggest_replacements(resume_text)
        output_file_path = save_updated_resume(updated_text)

        suggestions_text = "Suggested Improvements:\n" + '\n'.join(suggestions) if suggestions else "No improvements suggested. Your resume looks good!"

        return updated_text, suggestions_text, output_file_path
    except:
        return "An error occurred. Please try again with a different file.", "No suggestions available.", None

# Cell 4: Create and launch Gradio interface
iface = gr.Interface(
    fn=process_resume,
    inputs=gr.File(label="Upload Your Resume (.docx only)", file_types=[".docx"]),
    outputs=[
        gr.Textbox(label="Optimized Resume Text", lines=10),
        gr.Textbox(label="ATS Improvement Suggestions", lines=10),
        gr.File(label="Download Optimized Resume")
    ],
    title="ATS Resume Optimizer",
    description="Upload your .docx resume for ATS optimization. We'll suggest improvements and provide an updated version.",
    theme="default",
    css="""
        .gradio-container {max-width: 800px; margin: auto;}
        .output-markdown {white-space: pre-wrap;}
    """
)

iface.launch(share=True)