from functools import wraps

import basicauth
from flask import Response, request

from llm_server import opts


def check_auth(token):
    if not opts.admin_token:
        # The admin token is not set/enabled.
        # Default: deny all.
        return False
    password = None
    if token.startswith('Basic '):
        try:
            _, password = basicauth.decode(token)
        except:
            return False
    elif token.startswith('Bearer '):
        password = token.split('Bearer ', maxsplit=1)
        del password[0]
        password = ''.join(password)
    return password == opts.admin_token


def authenticate():
    """Sends a 401 response that enables basic auth"""
    return Response(
        'AUTHENTICATION REQUIRED', 401,
        {'WWW-Authenticate': 'Basic realm="Login Required"'})


def requires_auth(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        auth = request.headers.get('Authorization')
        if not auth or not check_auth(auth):
            return authenticate()
        return f(*args, **kwargs)

    return decorated