import os
import re
import io
from datetime import datetime
import PyPDF2
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForSeq2SeqLM
from groq import Groq
import gradio as gr
from docxtpl import DocxTemplate

# Set your API key for Groq
os.environ["GROQ_API_KEY"] = "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"
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

# --- PDF/Text Extraction Functions --- #
def extract_text_from_file(file_path):
    """Extracts text from PDF or TXT files based on file extension."""
    if file_path.endswith('.pdf'):
        return extract_text_from_pdf(file_path)
    elif file_path.endswith('.txt'):
        return extract_text_from_txt(file_path)
    else:
        raise ValueError("Unsupported file type. Only PDF and TXT files are accepted.")

def extract_text_from_pdf(pdf_file_path):
    """Extracts text from a PDF file."""
    with open(pdf_file_path, 'rb') as pdf_file:
        pdf_reader = PyPDF2.PdfReader(pdf_file)
        text = ''.join(page.extract_text() for page in pdf_reader.pages if page.extract_text())
    return text

def extract_text_from_txt(txt_file_path):
    """Extracts text from a .txt file."""
    with open(txt_file_path, 'r', encoding='utf-8') as txt_file:
        return txt_file.read()

# --- Skill Extraction with Llama Model --- #
def extract_skills_llama(text):
    """Extracts skills from the text using the Llama model via Groq API."""
    try:
        response = client.chat.completions.create(
            messages=[{"role": "user", "content": f"Extract skills from the following text: {text}"}],
            model="llama3-70b-8192",
        )
        skills = response.choices[0].message.content.split(', ')  # Expecting a comma-separated list
        return skills
    except Exception as e:
        raise RuntimeError(f"Error during skill extraction: {e}")

# --- Job Description Processing --- #
def process_job_description(job_description_text):
    """Processes the job description text and extracts relevant skills."""
    job_description_text = preprocess_text(job_description_text)
    return extract_skills_llama(job_description_text)

# --- Text Preprocessing --- #
def preprocess_text(text):
    """Preprocesses text for analysis (lowercase, punctuation removal)."""
    text = text.lower()
    text = re.sub(r'[^\w\s]', '', text)  # Remove punctuation
    return re.sub(r'\s+', ' ', text).strip()  # Remove extra whitespace

# --- Resume Similarity Calculation --- #
def calculate_resume_similarity(resume_text, job_description_text):
    """Calculates similarity score between resume and job description using a sentence transformer model."""
    model_name = "cross-encoder/stsb-roberta-base"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name)

    inputs = tokenizer(resume_text, job_description_text, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
        similarity_score = torch.sigmoid(outputs.logits).item()
    return similarity_score

# --- Communication Generation --- #
def communication_generator(resume_skills, job_description_skills, similarity_score, max_length=150):
    """Generates a communication response based on the extracted skills from the resume and job description."""
    model_name = "google/flan-t5-base"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

    # Assess candidate fit based on similarity score
    fit_status = "fit for the job" if similarity_score >= 0.7 else "not a fit for the job"

    # Create a more detailed communication message
    message = (
        f"After a thorough review of the candidate's resume, we found a significant alignment "
        f"between their skills and the job description requirements. The candidate possesses the following "
        f"key skills: {', '.join(resume_skills)}. These align well with the job requirements, particularly in areas such as "
        f"{', '.join(job_description_skills)}. The candidate’s diverse expertise suggests they would make a valuable addition to our team. "
        f"We believe the candidate is {fit_status}. If further evaluation is needed, please let us know how we can assist."
    )

    inputs = tokenizer(message, return_tensors="pt", padding=True, truncation=True)
    response = model.generate(**inputs, max_length=max_length, num_beams=4, early_stopping=True)

    return tokenizer.decode(response[0], skip_special_tokens=True)

# --- Sentiment Analysis --- #
def sentiment_analysis(text):
    """Analyzes the sentiment of the text."""
    model_name = "mrm8488/distiluse-base-multilingual-cased-v2-finetuned-stsb_multi_mt-es"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name)

    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
        predicted_sentiment = torch.argmax(outputs.logits).item()
    return ["Negative", "Neutral", "Positive"][predicted_sentiment]

# --- Resume Analysis Function --- #
def analyze_resume(resume_file, job_description_file):
    """Analyzes the resume and job description, returning similarity score, skills, and communication response."""
    # Extract resume text based on file type
    try:
        resume_text = extract_text_from_file(resume_file.name)
        job_description_text = extract_text_from_file(job_description_file.name)
    except ValueError as ve:
        return str(ve)

    # Analyze texts
    job_description_skills = process_job_description(job_description_text)
    resume_skills = extract_skills_llama(resume_text)
    similarity_score = calculate_resume_similarity(resume_text, job_description_text)
    communication_response = communication_generator(resume_skills, job_description_skills, similarity_score)
    sentiment = sentiment_analysis(resume_text)

    return (
        f"Similarity Score: {similarity_score:.2f}",
        communication_response,
        f"Sentiment: {sentiment}",
        ", ".join(resume_skills),
        ", ".join(job_description_skills),
    )

# --- Offer Letter Generation --- #
def generate_offer_letter(template_file, candidate_name, role, start_date, hours):
    """Generates an offer letter from a template."""
    try:
        start_date = datetime.strptime(start_date, "%Y-%m-%d").strftime("%B %d, %Y")
    except ValueError:
        return "Invalid date format. Please use YYYY-MM-DD."

    context = {
        'candidate_name': candidate_name,
        'role': role,
        'start_date': start_date,
        'hours': hours
    }

    doc = DocxTemplate(template_file)
    doc.render(context)

    offer_letter_path = f"{candidate_name.replace(' ', '_')}_offer_letter.docx"
    doc.save(offer_letter_path)

    return offer_letter_path

# --- Gradio Interface --- #
iface = gr.Interface(
    fn=analyze_resume,
    inputs=[
        gr.File(label="Upload Resume (PDF/TXT)"),
        gr.File(label="Upload Job Description (PDF/TXT)")
    ],
    outputs=[
        gr.Textbox(label="Similarity Score"),
        gr.Textbox(label="Communication Response"),
        gr.Textbox(label="Sentiment Analysis"),
        gr.Textbox(label="Extracted Resume Skills"),
        gr.Textbox(label="Extracted Job Description Skills"),
    ],
    title="Resume and Job Description Analyzer",
    description="This tool analyzes a resume against a job description to extract skills, calculate similarity, and generate communication responses."
)

iface.launch()