import json
import os
import gradio as gr
import pandas as pd
from openai import OpenAI
from transformers import pipeline
from dotenv import load_dotenv

load_dotenv()
api_key = os.getenv("OPENAI_API_KEY")
client = OpenAI(api_key=api_key)

pipe = pipeline("text-classification", model="cardiffnlp/twitter-roberta-base-sentiment-latest")

def process_csv(file):
    df = pd.read_csv(file)
    if "Feedback" not in df.columns or "Employee" not in df.columns:
        return None, "❌ Error: CSV must contain 'Employee' and 'Feedback' columns."
    df["Sentiment"] = df["Feedback"].apply(lambda x: pipe(x)[0]["label"])
    return {"df": df}, "✅ Your CSV file is processed!"

def predict_attrition_risk(employee_name: str, sentiment: str, explanation: str):
    risk_mapping = {"positive": "Low Risk", "neutral": "Medium Risk", "negative": "High Risk"}
    risk_level = risk_mapping.get(sentiment.lower(), "Unknown Sentiment")

    return f"**{employee_name}**: {risk_level} \n\n📝 {explanation}"


def analyze_attrition_with_llm(df_dict, hr_query):
    if df_dict is None or "df" not in df_dict:
        return "❌ Error: No processed data. Upload a CSV first."
    
    df = df_dict["df"]
    employees_data = {row["Employee"].strip(): row["Sentiment"] for _, row in df.iterrows()}
    
    prompt = f"HR Query: {hr_query}\nEmployees Data: {json.dumps(employees_data, indent=2)}"
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        functions=[
            {
                "name": "predict_attrition_risk",
                "description": "Predicts attrition risk based on sentiment and provides explanation.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "employee_name": {"type": "string", "description": "Employee's name"},
                        "sentiment": {"type": "string", "description": "Extracted sentiment"},
                        "explanation": {"type": "string", "description": "Brief explanation"}
                    },
                    "required": ["employee_name", "sentiment", "explanation"]
                }
            }
        ],
        function_call="auto"
    )
    
    print(response)  # Debugging: Print response

    # Handle function call response properly
    if response.choices[0].message.function_call:
        function_args = json.loads(response.choices[0].message.function_call.arguments)
        return predict_attrition_risk(**function_args)
    else:
        return "🤖 No response generated."

with gr.Blocks() as demo:
    gr.Markdown("<h1>AI-Driven Employee Attrition Risk Analysis</h1>")
    file_input = gr.File(label="Upload Employee Feedback CSV", file_types=[".csv"])
    process_button = gr.Button("Process CSV file")
    process_message = gr.Markdown()
    hr_input = gr.Textbox(label="Employee Name or HR Query")
    analyze_button = gr.Button("Check Attrition Risk")
    output_text = gr.Markdown()
    df_state = gr.State()

    process_button.click(process_csv, inputs=file_input, outputs=[df_state, process_message])
    analyze_button.click(analyze_attrition_with_llm, inputs=[df_state, hr_input], outputs=output_text)

demo.launch(share=True)