from __future__ import annotations
import os
import re
import subprocess
import zipfile
import json
import torch
from typing import List

os.environ["OMP_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["VECLIB_MAXIMUM_THREADS"] = "1"

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)

# --- UPDATED: Using Deepseek-Coder-V2-Lite-Instruct for better performance ---
tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/Deepseek-Coder-V2-Lite-Instruct", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    "deepseek-ai/Deepseek-Coder-V2-Lite-Instruct",
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True
)

# ===============================
# Text Extraction (PDF/DOCX)
# ===============================
def extract_text(file_path: str) -> str:
    """Extract text from PDF or DOCX resumes."""
    if not file_path or not os.path.isfile(file_path):
        return ""

    lower_name = file_path.lower()
    try:
        if lower_name.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 lower_name.endswith('.docx'):
            with zipfile.ZipFile(file_path) as zf:
                with zf.open('word/document.xml') as docx_xml:
                    xml_bytes = docx_xml.read()
                    xml_text = xml_bytes.decode('utf-8', errors='ignore')
                    xml_text = re.sub(r'<w:p[^>]*>', '\n', xml_text, flags=re.I)
                    text = re.sub(r'<[^>]+>', ' ', xml_text)
                    return re.sub(r'\s+', ' ', text)
        else:
            return ""
    except Exception:
        return ""

# ===============================
# Name Extraction (Fallback)
# ===============================
def extract_name(text: str, filename: str) -> str:
    """Extract candidate's name from resume text or filename."""
    if text:
        lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
        for line in lines[:10]:
            if re.match(r'(?i)resume|curriculum vitae', line):
                continue
            words = line.split()
            if 1 < len(words) <= 4:
                if all(re.match(r'^[A-ZÀ-ÖØ-Þ][\w\-]*', w) for w in words):
                    return line
    base = os.path.basename(filename)
    base = re.sub(r'\.(pdf|docx|doc)$', '', base, flags=re.I)
    base = re.sub(r'[\._-]+', ' ', base)
    base = re.sub(r'(?i)\b(cv|resume)\b', '', base)
    return base.title().strip()

# ===============================
# Janus-Pro Parsing
# ===============================
def parse_with_deepseek(text: str) -> dict:
    """Use Deepseek-Coder-V2-Lite-Instruct to extract resume details in JSON format."""
    
    prompt = f"""
Extract the following information from the resume text provided below. Your response should be a valid JSON object.

Information to extract:
- Full Name: The candidate's full name.
- Email: The candidate's email address.
- Phone: The candidate's phone number.
- Skills: A list of technical and soft skills.
- Education: A list of academic degrees and institutions.
- Experience: A list of previous jobs, including company, title, and dates.

Resume Text:
{text}

Return only valid JSON in the following format:
{{
  "name": "Full Name",
  "email": "email@example.com",
  "phone": "+961-xxx-xxx",
  "skills": ["Skill1", "Skill2", "Skill3"],
  "education": ["Degree1 - Institution1", "Degree2 - Institution2"],
  "experience": ["Job1 - Company1 (Dates)", "Job2 - Company2 (Dates)"]
}}
"""
    
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=512)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    import re, json
    match = re.search(r"\{.*\}", response, re.S)
    if match:
        try:
            return json.loads(match.group())
        except:
            pass
    
    return {"name": "", "email": "", "phone": "", "skills": [], "education": [], "experience": []}