import gradio as gr
import os
import docx2txt
import PyPDF2
import json
from openai import OpenAI

# Replace this with your actual API key
client = OpenAI(
    base_url="https://api.aimlapi.com/v1",
    api_key=os.getenv("AIML_API_KEY")
)

# -------- Resume Parsing --------
def extract_text_from_pdf(file):
    try:
        reader = PyPDF2.PdfReader(file)
        return "\n".join([page.extract_text() or "" for page in reader.pages])
    except Exception:
        return ""

def extract_text_from_docx(file):
    try:
        return docx2txt.process(file)
    except Exception:
        return ""

def parse_resume(file):
    ext = os.path.splitext(file.name)[1].lower()
    if ext == ".pdf":
        return extract_text_from_pdf(file)
    elif ext == ".docx":
        return extract_text_from_docx(file)
    else:
        return None

# -------- AI Screening --------
def call_ai_screening_agent(resume_texts, job_description, min_experience):
    prompt = f"""
You are an AI resume screening assistant. Compare each candidate's resume to the job description below.
Only include candidates who meet the minimum required years of experience: {min_experience}.
For each candidate, return:
- name (from filename or parsed text)
- strengths (2–5 key strengths)
- risks (0–3 potential risks)
- score (from 1 to 10)

Job Description:
\"\"\"{job_description}\"\"\"

Resumes:
{resume_texts}

Return a JSON list of candidates as described.
"""
    try:
        response = client.chat.completions.create(
            model="gpt-4-turbo",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.2
        )
        return json.loads(response.choices[0].message.content)
    except Exception as e:
        return {"error": str(e)}

# -------- Main App Logic --------
def process_resumes(files, job_description, min_experience):
    if not files:
        return "⚠️ Please upload at least one resume.", None
    if not job_description.strip():
        return "⚠️ Job description cannot be empty.", None

    parsed_resumes = []
    for file in files:
        content = parse_resume(file)
        if not content or len(content.strip()) == 0:
            continue
        parsed_resumes.append({"name": os.path.basename(file.name), "text": content})

    if not parsed_resumes:
        return "⚠️ Could not parse any valid resumes.", None

    resume_texts = "\n\n".join([f"{r['name']}:\n{r['text']}" for r in parsed_resumes])
    results = call_ai_screening_agent(resume_texts, job_description, min_experience)

    if isinstance(results, dict) and "error" in results:
        return f"⚠️ API Error: {results['error']}", None

    # Display top 3 candidates
    top_candidates = sorted(results, key=lambda x: x["score"], reverse=True)[:3]
    display_md = ""
    for candidate in top_candidates:
        display_md += f"### {candidate['name']}  \n"
        display_md += f"⭐ **Score**: {candidate['score']}  \n"
        display_md += "✅ **Strengths**:\n"
        for s in candidate['strengths']:
            display_md += f"- ✅ {s}\n"
        if candidate['risks']:
            display_md += "⚠️ **Risks**:\n"
            for r in candidate['risks']:
                display_md += f"- ⚠️ {r}\n"
        display_md += "\n---\n"

    return None, display_md

# -------- Gradio UI --------
with gr.Blocks(title="SmartHire - AI Job Screening Assistant") as demo:
    gr.Markdown("# 🤖 SmartHire — AI Job Screening Assistant")
    gr.Markdown("Upload resumes and paste a job description to find the best candidates automatically!")

    with gr.Row():
        resume_input = gr.File(file_types=[".pdf", ".docx"], file_count="multiple", label="Upload Resumes")
        experience_input = gr.Number(label="Minimum Years of Experience", value=0)

    job_desc_input = gr.Textbox(lines=8, placeholder="Paste the job description here...", label="Job Description")
    process_button = gr.Button("Run Screening")

    error_output = gr.Textbox(label="Warnings / Errors", visible=False)
    output_md = gr.Markdown()

    process_button.click(
        fn=process_resumes,
        inputs=[resume_input, job_desc_input, experience_input],
        outputs=[error_output, output_md]
    )

demo.launch()