import gradio as gr
import joblib
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.calibration import CalibratedClassifierCV

# Attempt to load the model and vectorizer
try:
    model = joblib.load("helix-psa.pkl")
    print("Model loaded successfully:", type(model))
except Exception as e:
    print("Error loading model:", e)

try:
    vectorizer = joblib.load("helix-psa-vectorizer.pkl")
    print("Vectorizer loaded successfully:", type(vectorizer))
except Exception as e:
    print("Error loading vectorizer:", e)

def predictPasswordStrength(password_input):
    password_tfidf = vectorizer.transform([password_input])
    
    # Make predictions
    predicted_proba = model.predict_proba(password_tfidf)
    predicted_class = int(model.predict(password_tfidf)[0])  # Convert to Python integer
    output = ''
    if predicted_class == 0:
        output = "The password is very weak..."
    elif predicted_class == 1:
        output = "The password is average."
    else:
        output = "The password is strong. But alas, it is not unbreakable."

    confidence = float(predicted_proba.max())

    return password_input, output, confidence

demo = gr.Interface(
    fn=predictPasswordStrength,
    inputs=[gr.Textbox('Hello123', label='Password', info='The password to check the strength of', max_lines=1)],
    outputs=[gr.Dataframe(row_count = (2, "dynamic"), col_count=(3, "fixed"), label="Generated Names", headers=["Password", "Prediction", "Confidence"])],
    title='Helix - Password Strength Analyzer',
    description='A password strength analyzer, trained on over 10 million different passwords.'
)

demo.launch()