import os
import re
import torch  # Explicitly imported if you want to use torch directly
import tempfile
from io import BytesIO

import streamlit as st
from PIL import Image
from transformers import pipeline
from pdf2image import convert_from_bytes

#####################################
# Load the OCR Pipeline (Uses Torch)
#####################################
try:
    # Ensure your transformers library is updated (>=4.x)
    ocr_pipeline = pipeline("image-to-text", model="YouLiXiya/tinyllava-v1.0-1.1b-hf")
    st.write("Model loaded successfully!")
except Exception as e:
    st.error(f"Error loading model: {e}")
    st.stop()

#####################################
# Utility: Convert PDF to Images
#####################################
def convert_pdf_to_images(pdf_bytes):
    try:
        images = convert_from_bytes(pdf_bytes)
        return images
    except Exception as e:
        st.error(f"PDF conversion error: {e}")
        return []

#####################################
# Pipeline: Extract Text with OCR Pipeline
#####################################
def extract_text_from_file(file_obj):
    file_extension = os.path.splitext(file_obj.name)[1].lower()
    full_text = ""

    if file_extension == ".pdf":
        file_bytes = file_obj.read()
        images = convert_pdf_to_images(file_bytes)
        for img in images:
            result = ocr_pipeline(img)
            if isinstance(result, list) and "text" in result[0]:
                full_text += result[0]["text"] + "\n"
    else:
        try:
            img = Image.open(file_obj)
            result = ocr_pipeline(img)
            if isinstance(result, list) and "text" in result[0]:
                full_text = result[0]["text"]
        except Exception as e:
            full_text = f"Error processing image: {e}"

    return full_text

#####################################
# Information Extraction Functions
#####################################
def extract_resume_info(text):
    info = {
        "Name": None,
        "Age": None,
        "Job Experience": None,
        "Skills": None,
        "Expected Industry/Direction": None,
    }

    # Extract name (e.g., "Name: John Doe")
    name_match = re.search(r"[Nn]ame[:\-]\s*([A-Za-z\s]+)", text)
    if name_match:
        info["Name"] = name_match.group(1).strip()
    else:
        # Heuristic: pick the first sequence of capitalized words
        potential_names = re.findall(r'\b[A-Z][a-z]+(?:\s+[A-Z][a-z]+)+\b', text)
        if potential_names:
            info["Name"] = potential_names[0]

    # Extract age
    age_match = re.search(r"[Aa]ge[:\-]\s*(\d{1,2})", text)
    if age_match:
        info["Age"] = age_match.group(1)

    # Extract job experience (years)
    exp_match = re.search(r"(\d+)\s+(?:years|yrs)\s+(?:of\s+)?experience", text, re.IGNORECASE)
    if exp_match:
        info["Job Experience"] = exp_match.group(1) + " years"
    else:
        exp_line = re.search(r"(Experience|Background)[:\-]\s*(.*)", text, re.IGNORECASE)
        if exp_line:
            info["Job Experience"] = exp_line.group(2).strip()

    # Extract skills (e.g., "Skills: Python, Java, SQL")
    skills_match = re.search(r"[Ss]kills[:\-]\s*(.+)", text)
    if skills_match:
        skills_text = skills_match.group(1)
        skills = [s.strip() for s in re.split(r",|\n", skills_text) if s.strip()]
        info["Skills"] = skills

    # Extract expected industry/direction
    industry_match = re.search(r"(Industry|Interest|Direction)[:\-]\s*(.+)", text, re.IGNORECASE)
    if industry_match:
        info["Expected Industry/Direction"] = industry_match.group(2).strip()

    return info

#####################################
# Main Processing Logic
#####################################
def process_resume(file_obj):
    if file_obj is None:
        return None, None

    resume_text = extract_text_from_file(file_obj)
    resume_info = extract_resume_info(resume_text)
    return resume_text, resume_info

#####################################
# Streamlit UI
#####################################
st.title("Resume Extraction and Information Parsing")
st.markdown("""
Upload a resume file (in PDF or image format) and the app will extract its text and parse critical candidate information.
""")

uploaded_file = st.file_uploader("Upload Resume (PDF or Image)", type=["pdf", "png", "jpg", "jpeg"])

if st.button("Extract Info"):
    if uploaded_file is None:
        st.error("Please upload a file first.")
    else:
        with st.spinner("Processing..."):
            resume_text, resume_info = process_resume(uploaded_file)
        st.subheader("Extracted Resume Text")
        st.text_area("", resume_text, height=200)
        st.subheader("Parsed Resume Information")
        st.json(resume_info)