import gradio as gr
from transformers import pipeline
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

# Initialize more lightweight models
summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device=-1)  # Use CPU
sentence_model = SentenceTransformer('paraphrase-MiniLM-L6-v2')

# Simple in-memory user storage (replace with proper database in production)
users = {}

def get_embedding(text):
    return sentence_model.encode(text)

def calculate_job_match(job_description, user_data):
    job_embedding = get_embedding(job_description)
    user_embedding = get_embedding(user_data)
    similarity = cosine_similarity([job_embedding], [user_embedding])[0][0]
    return similarity

def register(username, password, email):
    if username in users:
        return "Username already exists"
    users[username] = {"password": password, "email": email, "user_data": ""}
    return "Registered successfully"

def login(username, password):
    if username not in users or users[username]["password"] != password:
        return "Invalid username or password"
    return "Logged in successfully"

def update_profile(username, email, user_data):
    if username not in users:
        return "User not found"
    users[username]["email"] = email
    users[username]["user_data"] = user_data
    return "Profile updated successfully"

def generate_text(prompt, max_length=150, min_length=50):
    summary = summarizer(prompt, max_length=max_length, min_length=min_length, do_sample=False)[0]['summary_text']
    return summary

def generate_cv(username, job_description):
    if username not in users:
        return "User not found"
    user_data = users[username]["user_data"]
    match_score = calculate_job_match(job_description, user_data)
    
    prompt = f"Generate a CV based on the following job description: {job_description}\nUser data: {user_data}\nJob match score: {match_score:.2f}"
    generated_cv = generate_text(prompt, max_length=300, min_length=100)
    
    return f"Generated CV:\n\n{generated_cv}\n\nJob Match Score: {match_score:.2f}"

def generate_cover_letter(username, job_description):
    if username not in users:
        return "User not found"
    user_data = users[username]["user_data"]
    match_score = calculate_job_match(job_description, user_data)
    
    prompt = f"Generate a cover letter based on the following job description: {job_description}\nUser data: {user_data}\nJob match score: {match_score:.2f}"
    cover_letter = generate_text(prompt, max_length=250, min_length=100)
    
    return f"Generated Cover Letter:\n\n{cover_letter}\n\nJob Match Score: {match_score:.2f}"

def prepare_interview(username, job_description):
    if username not in users:
        return "User not found"
    user_data = users[username]["user_data"]
    match_score = calculate_job_match(job_description, user_data)
    
    prompt = f"Generate 5 potential interview questions based on the following job description: {job_description}\nUser data: {user_data}\nJob match score: {match_score:.2f}"
    interview_questions = generate_text(prompt, max_length=200, min_length=100)
    
    return f"Potential Interview Questions:\n\n{interview_questions}\n\nJob Match Score: {match_score:.2f}"

with gr.Blocks() as demo:
    gr.Markdown("# Advanced Personalized CV Generator")
    
    with gr.Tab("Register"):
        register_username = gr.Textbox(label="Username")
        register_password = gr.Textbox(label="Password", type="password")
        register_email = gr.Textbox(label="Email")
        register_button = gr.Button("Register")
        register_output = gr.Textbox(label="Output")
        register_button.click(register, inputs=[register_username, register_password, register_email], outputs=register_output)
    
    with gr.Tab("Login"):
        login_username = gr.Textbox(label="Username")
        login_password = gr.Textbox(label="Password", type="password")
        login_button = gr.Button("Login")
        login_output = gr.Textbox(label="Output")
        login_button.click(login, inputs=[login_username, login_password], outputs=login_output)
    
    with gr.Tab("Update Profile"):
        update_username = gr.Textbox(label="Username")
        update_email = gr.Textbox(label="Email")
        update_user_data = gr.Textbox(label="Professional Information")
        update_button = gr.Button("Update Profile")
        update_output = gr.Textbox(label="Output")
        update_button.click(update_profile, inputs=[update_username, update_email, update_user_data], outputs=update_output)
    
    with gr.Tab("Generate CV"):
        cv_username = gr.Textbox(label="Username")
        cv_job_description = gr.Textbox(label="Job Description")
        cv_button = gr.Button("Generate CV")
        cv_output = gr.Textbox(label="Generated CV")
        cv_button.click(generate_cv, inputs=[cv_username, cv_job_description], outputs=cv_output)
    
    with gr.Tab("Generate Cover Letter"):
        cl_username = gr.Textbox(label="Username")
        cl_job_description = gr.Textbox(label="Job Description")
        cl_button = gr.Button("Generate Cover Letter")
        cl_output = gr.Textbox(label="Generated Cover Letter")
        cl_button.click(generate_cover_letter, inputs=[cl_username, cl_job_description], outputs=cl_output)
    
    with gr.Tab("Prepare for Interview"):
        int_username = gr.Textbox(label="Username")
        int_job_description = gr.Textbox(label="Job Description")
        int_button = gr.Button("Generate Interview Questions")
        int_output = gr.Textbox(label="Interview Questions")
        int_button.click(prepare_interview, inputs=[int_username, int_job_description], outputs=int_output)

demo.launch()