from flask import Flask, render_template, request, redirect, url_for, session
from werkzeug.security import generate_password_hash, check_password_hash
import os
from sklearn.datasets import load_iris
from sklearn.linear_model import LogisticRegression

app = Flask(__name__)
app.secret_key = os.environ.get("SECRET_KEY", "supersecretkey123")

# Train ML Model
iris = load_iris(as_frame=True)
X, y = iris.data.values, iris.target
species = iris.target_names
iris_model = LogisticRegression(max_iter=300)
iris_model.fit(X, y)

# In-memory DB
users = {}

@app.route("/")
def home():
    if "user" in session:
        return redirect(url_for("dashboard"))
    return redirect(url_for("login"))

@app.route("/signup", methods=["GET", "POST"])
def signup():
    if request.method == "POST":
        username = request.form["username"]
        password = request.form["password"]
        confirm = request.form.get("confirmPassword")

        if password != confirm:
            return render_template("signup.html", error="Passwords do not match")
        if username in users:
            return render_template("signup.html", error="Username already exists")

        users[username] = generate_password_hash(password)
        session["user"] = username
        return redirect(url_for("dashboard"))
    return render_template("signup.html")

@app.route("/login", methods=["GET", "POST"])
def login():
    if request.method == "POST":
        username = request.form["username"]
        password = request.form["password"]

        if username in users and check_password_hash(users[username], password):
            session["user"] = username
            return redirect(url_for("dashboard"))
        return render_template("login.html", error="Invalid username or password")
    return render_template("login.html")

@app.route("/dashboard")
def dashboard():
    if "user" not in session:
        return redirect(url_for("login"))
    return render_template("dashboard.html", username=session["user"])

@app.route("/predict", methods=["POST"])
def predict():
    if "user" not in session:
        return redirect(url_for("login"))

    try:
        sl = float(request.form["sepal_length"])
        sw = float(request.form["sepal_width"])
        pl = float(request.form["petal_length"])
        pw = float(request.form["petal_width"])
    except Exception:
        return render_template("dashboard.html", username=session["user"], error="Invalid input values", scroll_to="predictor")

    res_idx = iris_model.predict([[sl, sw, pl, pw]])[0]
    res = species[res_idx]
    return render_template("dashboard.html", username=session["user"], prediction=res, scroll_to="predictor")

@app.route("/logout")
def logout():
    session.pop("user", None)
    return redirect(url_for("login"))

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, debug=True)