from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
import subprocess, zipfile, re, os

# === Load pretrained HF model ===
MODEL_NAME = "sravya-abburi/ResumeParserBERT"  # or "Kiet/autotrain-resume_parser-1159242747"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForTokenClassification.from_pretrained(MODEL_NAME)

# Use CPU for stability (device=-1) to avoid GPU memory issues from other parts of the app
ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple", device=-1)

# === Extract text from PDF/DOCX ===
def extract_text(file_path: str) -> str:
    """Extract text from PDF or DOCX resumes."""
    if file_path.lower().endswith(".pdf"):
        result = subprocess.run(
            ["pdftotext", "-layout", file_path, "-"],
            stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=False
        )
        return result.stdout.decode("utf-8", errors="ignore")
    elif file_path.lower().endswith(".docx"):
        with zipfile.ZipFile(file_path) as zf:
            with zf.open("word/document.xml") as docx_xml:
                xml_text = docx_xml.read().decode("utf-8", errors="ignore")
                xml_text = re.sub(r"<w:p[^>]*>", "\n", xml_text, flags=re.I)
                return re.sub(r"<[^>]+>", " ", xml_text)
    return ""

# === Parse resume with NER ===
def parse_resume(file_path: str, filename: str = None) -> dict:
    """Parse resume and extract Name, Skills, Education, Experience."""
    text = extract_text(file_path)
    entities = ner_pipeline(text)

    name, skills, education, experience = [], [], [], []

    for ent in entities:
        word = ent["word"].strip()
        label = ent["entity_group"].upper()

        # Skip empty or placeholder tokens
        if not word or word.startswith("LABEL_"):
            continue

        if label == "NAME":
            name.append(word)
        elif label == "SKILL":
            skills.append(word)
        elif label in ["EDUCATION", "DEGREE"]:
            education.append(word)
        elif label in ["EXPERIENCE", "JOB", "ROLE"]:
            experience.append(word)

    return {
        "name": " ".join(dict.fromkeys(name)),
        "skills": ", ".join(dict.fromkeys(skills)),
        "education": ", ".join(dict.fromkeys(education)),
        "experience": ", ".join(dict.fromkeys(experience))
    }