from __future__ import absolute_import
from dostadmin import app
from dostadmin import db_model
from flask_login import current_user
from flask_principal import identity_loaded, RoleNeed, UserNeed
from flask_wtf import FlaskForm
from wtforms import StringField, PasswordField, SubmitField
from wtforms.validators import DataRequired, Length

app.login_manager.login_view = "/login"


@app.login_manager.user_loader
def load_user(user_id):
    return db_model.AdminUser.query.get(user_id)


class LoginForm(FlaskForm):
    username = StringField(
        "Username",
        [
            DataRequired(message="Invalid username/password"),
            Length(min=2, max=25, message="Invalid username/password"),
        ],
    )
    password = PasswordField(
        "Password",
        [
            DataRequired(message="Invalid username/password"),
            Length(min=6, max=30, message="Invalid username/password"),
        ],
    )
    submit = SubmitField("Login")


@identity_loaded.connect_via(app)
def on_identity_loaded(sender, identity):
    # Set the identity user object
    identity.user = current_user

    # Add the UserNeed to the identity
    if hasattr(current_user, "id"):
        identity.provides.add(UserNeed(current_user.id))

    # Assuming the User model has a list of roles, update the
    # identity with the roles that the user provides
    if hasattr(current_user, "roles"):
        for role in current_user.roles:
            identity.provides.add(RoleNeed(role.name))