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

# --------------------
# Load Model
# --------------------
MODEL_NAME = "sravya-abburi/ResumeParserBERT"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForTokenClassification.from_pretrained(MODEL_NAME)

ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")

# --------------------
# Extract Text
# --------------------
def extract_text(file_path: str) -> str:
    """Extract text from PDF/DOCX resumes."""
    if file_path.lower().endswith(".pdf"):
        try:
            result = subprocess.run(
                ["pdftotext", "-layout", file_path, "-"],
                stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=False
            )
            return result.stdout.decode("utf-8", errors="ignore")
        except:
            return ""
    elif file_path.lower().endswith(".docx"):
        try:
            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)
        except:
            return ""
    return ""

# --------------------
# Parse Resume
# --------------------
def parse_resume(file_path: str, filename: str = None) -> dict:
    """Extract Name, Skills, Education, Experience from resume."""
    text = extract_text(file_path)
    entities = ner_pipeline(text)

    name, skills, education, experience = [], [], [], []
    for ent in entities:
        label = ent["entity_group"].upper()
        word = ent["word"].strip()

        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))
    }

# --------------------
# Example
# --------------------
if __name__ == "__main__":
    resume_path = "resume.pdf"  # Change to test file
    result = parse_resume(resume_path)
    print(json.dumps(result, indent=2))