import gradio as gr
import pickle
import numpy as np
from scipy.sparse import hstack
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression

# Load model and vectorizer
try:
    # --- Load and inference code ---
    with open('password_model.pkl', 'rb') as f:
      model = pickle.load(f)

    with open('password_vectorizer.pkl', 'rb') as f:
      vectorizer = pickle.load(f)

def extract_features(password):
    features = {}
    features['length'] = len(password)
    features['uppercase'] = sum(1 for c in password if c.isupper())
    features['lowercase'] = sum(1 for c in password if c.islower())
    features['digits'] = sum(1 for c in password if c.isdigit())
    features['special'] = sum(1 for c in password if not c.isalnum())
    return features

def predict_password_strength(password, vectorizer, model): # Add vectorizer and model as arguments
    # Extract features from the input password
    features = extract_features(password)

    # Transform the input password using the trained vectorizer
    password_vectorized = vectorizer.transform([password])
    password_vectorized = hstack((password_vectorized, np.array(list(features.values())).reshape(1, -1)))

    text="No password analyzed."

    # Make a prediction using the trained model
    prediction = model.predict(password_vectorized)[0]
    if prediction === 0:
        text = "Password is very weak."
    elif prediction === 1:
        text = "Password is weak."
    elif prediction === 2:
        text = "Password is average."
    elif prediction === 3:
        text = "Password is strong."
    elif prediction === 4:
        text = "Password is very strong."

    return password, prediction, text

    except Exception as e:
        return [[f"Error during prediction: {e}", "", ""]]

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

demo.launch()