import pandas as pd
import gradio as gr
from retriever import get_relevant_passages
from reranker import rerank

# Load and clean CSV
df = pd.read_csv("assessments.csv")

def clean_df(df):
    df = df.copy()
    df["url"] = "https://www.shl.com" + df.iloc[:, 1]
    df["remote_support"] = df.iloc[:, 2].map(lambda x: "Yes" if x == "T" else "No")
    df["adaptive_support"] = df.iloc[:, 3].map(lambda x: "Yes" if x == "T" else "No")
    df["test_type"] = df.iloc[:, 4].astype(str).str.split("\\n")
    df["description"] = df.iloc[:, 5]
    df["duration"] = df.iloc[:, 8].astype(str).str.extract(r'(\d+)').astype(float)
    return df[["url", "adaptive_support", "remote_support", "description", "duration", "test_type"]]

df_clean = clean_df(df)

def recommend(query):
    top_k_df = get_relevant_passages(query, df_clean, top_k=20)
    candidates = top_k_df.to_dict(orient="records")
    return rerank(query, candidates)

iface = gr.Interface(
    fn=recommend,
    inputs=gr.Textbox(label="Enter Job Description", lines=4),
    outputs="json",
    title="SHL Assessment Recommender",
    description="Paste a job description to get the most relevant SHL assessments."
)

if __name__ == "__main__":
    iface.launch()