import functools

from psycopg2 import IntegrityError
from flask import (
    Blueprint, flash, g, redirect, render_template, request, session, url_for, current_app
)
from werkzeug.security import check_password_hash, generate_password_hash

from AstroSpace.db import get_conn

bp = Blueprint('auth', __name__, url_prefix='/auth')

@bp.route('/register', methods=('GET', 'POST'))
def register():
    if request.method == 'POST':
        username = request.form['username'].lower()
        password = request.form['password']
        db = get_conn()
        error = None

        if not username:
            error = 'Username is required.'
        elif not password:
            error = 'Password is required.'

        if error is None:
            try:
                with db.cursor() as cur:
                    cur.execute("LOCK TABLE users IN EXCLUSIVE MODE")
                    cur.execute("SELECT COUNT(*) AS user_count FROM users")
                    user_count = cur.fetchone()["user_count"]

                    if user_count >= current_app.config['MAX_USERS']:
                        error = 'Sorry maximum number of users reached :('
                    else:
                        is_first_user = user_count == 0
                        cur.execute(
                            "INSERT INTO users (username, password, admin) VALUES (%s, %s, %s)",
                            (username, generate_password_hash(password), is_first_user),
                        )
                if error is None:
                    db.commit()
                    return redirect(url_for("auth.login"))
                db.rollback()
            except IntegrityError:
                db.rollback()
                error = f"User {username} is already registered."

        flash(error)

    return render_template('auth/auth.html', title='Register', WebName = current_app.config["TITLE"])

@bp.route('/login', methods=('GET', 'POST'))
def login():
    if request.method == 'POST':
        username = request.form['username'].lower()
        password = request.form['password']
        db = get_conn()
        error = None
        with db.cursor() as cur:
            cur.execute(
                "SELECT * FROM users WHERE username = %s", (username,) )
            user = cur.fetchone()
        
        if user is None:
            error = 'Incorrect username.'
        elif not check_password_hash(user['password'], password):
            error = 'Incorrect password.'

        if error is None:
            session.clear()
            session['user_id'] = user['id']
            return redirect(url_for('index'))

        flash(error)

    return render_template('auth/auth.html', title='Log In', WebName = current_app.config["TITLE"])

@bp.before_app_request
def load_logged_in_user():
    user_id = session.get('user_id')

    if user_id is None:
        g.user = None
    else:
        cur = get_conn().cursor()
        cur.execute(
            "SELECT * FROM users WHERE id = %s", (user_id,)
        )
        g.user = cur.fetchone()

@bp.route('/logout', methods=('POST',))
def logout():
    session.clear()
    return redirect(url_for('index'))

def login_required(view):
    @functools.wraps(view)
    def wrapped_view(**kwargs):
        if g.user is None:
            return redirect(url_for('auth.login'))

        return view(**kwargs)

    return wrapped_view
