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)

except Exception as e:
    print(f"Error loading model/vectorizer: {e}")
    model = None
    vectorizer = None

# Function to extract features
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

# Function to predict password strength
def predict_password_strength(password):
    if not model or not vectorizer:
        return [["", "", "Model or vectorizer not loaded correctly"]]

    try:
        # 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)))

        # 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."
        else:
            text = "Unknown strength level."

        # Return the result as a list of lists to match the DataFrame format
        return [[password, prediction, text]]
    except Exception as e:
        return [[password, "", f"Error during prediction: {e}"]]

# Gradio Interface
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()