# ########################################################################################################################
# auth.py
# Auth Module
#
# Class for users authentication
#
#   All REST and downstream service errors are handled through log_and_raise_service
#   All Auth REST facing methods return rest_ok
#       {status:'OK, [user_status:<user_status>], [auth_status:<auth_status>, auth_msg:<auth_error_msg>]}
#
# Internal methods:
#   get_routes
#   startup
#   shutdown
#   check_auth
#
#
# REST methods for users:
#   auth_info
#   check_user
#   check_otp
#   login
#   register
#   recover_passwd
#   set_passwd
#   change_passwd
#
# REST methods for admin:
#   get_users
#   update_user
#   reset_user
#   suspend_user
#   change_org
#
# ########################################################################################################################

from datetime import datetime, timezone
import bcrypt
from starlette.routing import Route
from dl2050utils.core import oget
from dl2050utils.com import send_otp
from dl2050utils.restutils import OrjsonResponse, log_and_raise_service, enforce_required_args, rest_ok, mk_key, mk_jwt_token

# ####################################################################################################
# Defenitions and helper functions
# ####################################################################################################

# user_status
NOT_REGISTERED = 0
WAIT_SET_PASSWD = 1
SUSPENDED = 2
SUSPENDED_LICENCE = 3
SUSPENDED_FAILS = 4
SUSPENDED_REQUESTS = 5
ACTIVE = 10

# auth_status
AUTHORIZED = 100
NOT_AUTHORIZED = 101
WRONG_PASSWD = 102
REGISTER_NOT_ALLOWED = 103
USER_ALREADY_REGISTERED = 109
OTP_ERROR = 104
OTP_EXPIRED = 105
OTP_OK = 106
PASSWORD_CHANGED = 107
UNKNOWN_USER = 108
PASSWD_NOT_STRONG = 110

AUTH_MSGS = {
    'AUTHORIZED': 'User authorized',
    'NOT_AUTHORIZED': 'User not authorized',
    'UNKNOWN_USER': 'Unknown user',
    'REGISTER_NOT_ALLOWED': 'Register not allowed',
    'USER_ALREADY_REGISTERED': 'User already registered',
    'WRONG_PASSWD': 'Wrong password',
    'PASSWORD_CHANGED': 'Password changed',
    'PASSWD_NOT_STRONG': 'Password not strong enough',
    'OTP_ERROR': 'Wrong OTP',
    'OTP_EXPIRED': 'OTP expired',
    'OTP_OK': 'OTP OK',
}

MAX_FAILS = 5
MAX_REQUESTS = 1e6

def increase(d, f, reset=False):
    if d is None or f not in d: return
    if d[f] is None or reset: d[f] = 0
    d[f] += 1
    d['ts_access'] = datetime.now(timezone.utc)
    if f=='logins': d['ts_login'] = datetime.now(timezone.utc)

def get_age_in_seconds(ds):
    s1 = int(datetime.now(timezone.utc).strftime("%s"))
    s2 = int(datetime.strptime(ds,'%Y-%m-%d %H:%M:%S').strftime("%s"))
    return s1-s2

def get_user_info(user):
    return {
        'uid': user['id'],
        'email': user['email'],
        'status': user['ustatus'],
        'role': user['role'],
        'org': user['org'],
        'dep': user['org']
    }

def check_passwd_strengh(passwd):
    if type(passwd)!=str: return False
    if len(passwd)<6: return False
    return True

# ####################################################################################################
# Auth class
# ####################################################################################################

class Auth():
    def __init__(self, cfg, LOG, NOTIFY, db):
        self.cfg,self.LOG,self.NOTIFY,self.db = cfg,LOG,NOTIFY,db

    # ####################################################################################################
    # Helper methods
    # ####################################################################################################

    def ret(self, user_status=None, auth_status=None, payload=None):
        d = {}
        if user_status: d['user_status'] = user_status
        if auth_status: d['auth_status'] = auth_status
        if type(payload)==dict: d = {**d, **payload}
        return rest_ok(d)

    async def db_update(self, tbl, kc, d, prefix=''):
        res = await self.db.update(tbl, kc, d)
        if res:
            log_and_raise_service(self.LOG, label='AUTH', label2=prefix, msg='DB update error')

    async def enforce_find_user(self, d, label2='', msg='User not found'):
        user = await self.db.select_one('users', d)
        if user is None:
            log_and_raise_service(self.LOG, label='AUTH', label2=label2, msg=msg)
        return user

    async def enforce_admin_perm(self, request, label2='', msg='Not authorized', requires_su=False):
        user = await self.enforce_find_user({'id':request.user.display_name}, label2=label2, msg='Admin user not found')
        role,org = user['role'],user['org']
        if role not in ['su', 'admin'] or (requires_su==True and role!='su'):
            log_and_raise_service(self.LOG, label='AUTH',label2=label2, msg=msg)
        return role,org

    async def increase_fails(self, user):
        increase(user, 'fails')
        if user['fails']>MAX_FAILS:
            user['ustatus'] = SUSPENDED_FAILS
            self.LOG(4, 0, label='AUTH', label2='login', msg={'user':user['email'], 'msg':'Too many login fails - user suspended'})
        await self.db_update('users', 'email', user, 'increase_fails')
        if user['fails'] > MAX_FAILS:
            return SUSPENDED_FAILS
        return 0

    async def gen_otp(self, user, return_code=WAIT_SET_PASSWD):
        if self.otp_mode=='phone':
            if oget(user,['phone']) is None:
                log_and_raise_service(self.LOG, label='AUTH', label2='gen_otp', msg='phone not defined')
            user['email'] = ''
        else:
            if oget(user,['email']) is None:
                log_and_raise_service(self.LOG, label='AUTH', label2='gen_otp', msg='email not defined')
            user['phone'] = ''
        otp = mk_key(n=4)
        self.LOG(2, 0, label='GEN_OTP', label2='New OTP', msg={'email': user['email'], 'OTP': otp})
        err = send_otp(self.NOTIFY, self.otp_mode, self.product, user['email'], user['phone'], otp)
        if err is not None:
            self.LOG(4, 0, label='AUTH', label2='gen_otp', msg=err)
            log_and_raise_service(self.LOG, label='AUTH', label2='gen_otp', msg='email error when sending OTP')
        user['otp'],user['ts_passwd'],user['ustatus'] =  otp, datetime.now(timezone.utc),WAIT_SET_PASSWD
        await self.db_update('users', 'email', user, 'gen_otp')
        return rest_ok({'user_status': return_code})

    # ####################################################################################################
    # register_requests
    # ####################################################################################################

    async def register_requests(self, url, uid, t):
        """ Used by App to rate limit users """
        if url in ['/api/auth/is_auth', '/api/get_meta'] or uid==-1: return
        await self.db.insert('metrics', {'ts': datetime.now(timezone.utc), 'uid':uid, 'url':url, 't':t})
        user = await self.db.select_one('users', {'id': uid})
        increase(user, 'requests')
        if user['requests'] > MAX_REQUESTS:
            user['ustatus'] = SUSPENDED_REQUESTS
        await self.db_update('users', 'email', user, 'register_requests')

    # ####################################################################################################
    # Internal methods: used by App to manage Auth service
    # ####################################################################################################

    def get_routes(self):
        """" Used by App to include Auth routes in the App routes """
        return [
            Route('/api/auth/auth_info', endpoint=self.auth_info, methods=['GET']),
            Route('/api/auth/check_user', endpoint=self.check_user, methods=['POST']),
            Route('/api/auth/login', endpoint=self.login, methods=['POST']),
            Route('/api/auth/check_otp', endpoint=self.check_otp, methods=['POST']),
            Route('/api/auth/set_passwd', endpoint=self.set_passwd, methods=['POST']),
            Route('/api/auth/recover_passwd', endpoint=self.recover_passwd, methods=['POST']),
            Route('/api/auth/change_passwd', endpoint=self.change_passwd, methods=['POST']),
            Route('/api/auth/get_users', endpoint=self.get_users, methods=['POST']),
            Route('/api/auth/register', endpoint=self.register, methods=['POST']),
            Route('/api/auth/update_user', endpoint=self.update_user, methods=['POST']),
            Route('/api/auth/reset_user', endpoint=self.reset_user, methods=['POST']),
            Route('/api/auth/suspend_user', endpoint=self.suspend_user, methods=['POST']),
            Route('/api/auth/change_org', endpoint=self.change_org, methods=['POST']),
        ]

    async def startup(self):
        """" Used by App to configure and startup the Auth service """
        try:
            self.product = self.cfg['app']['product']
            self.auth_secret = self.cfg['rest']['auth_secret']
            self.allow_register = oget(self.cfg, ['rest', 'allow_register'], False)
            self.strong_passwd = oget(self.cfg, ['rest', 'strong_passwd'], False)
            self.otp_timeout = oget(self.cfg, ['rest', 'otp_timeout'], False)
            self.otp_mode = oget(self.cfg, ['rest', 'otp_mode'], False)
        except Exception as e:
            self.LOG(4, 0, label='AUTH', label2='startup', msg={'error_msg': str(e)})
            return True
        self.LOG(2, 0, label='AUTH', label2='startup', msg=f'OK (connected to DB:{self.db.dbname})')
        return False

    async def shutdown(self):
        """" Used by App to shutdown the Auth service """
        self.LOG(2, 0, label='AUTH', label2='shutdown', msg='OK')
        return False

    async def check_auth(self, request):
        """ Used by App to check for valid user token, through the valid decription of request.user.display_name
            Response must be a dict, not json, since it is an internal method """
        user = await self.enforce_find_user({'id':request.user.display_name}, label2='check_auth')
        return get_user_info(user)

    # ####################################################################################################
    # REST methods for users: all to return json
    # ####################################################################################################

    async def auth_info(self, request):
        """ Returns Auth info to be used on front-end flow """
        return rest_ok({
            'allow_register': self.allow_register,
            'strong_passwd': self.strong_passwd,
            'otp_timeout': self.otp_timeout,
            'otp_mode': self.otp_mode
        })

    async def check_user(self, request):
        """ Returns the users_status, searching by user email. To be used before asking for password
            If Auth allows self register and the user doesnt exists, starts the OTP process """
        data = await request.json()
        [email] = enforce_required_args(self.LOG, data, ['email'], label='AUTH', label2='check_user', as_list=True)
        user = await self.db.select_one('users', {'email':email})
        if user is None:
            return rest_ok({'user_status':NOT_REGISTERED})
        if user['ustatus']==NOT_REGISTERED and not self.allow_register:
            return await self.gen_otp(user, WAIT_SET_PASSWD)
        if user['ustatus']==ACTIVE:
            increase(user, 'logins')
            await self.db_update('users', 'email', user, 'check_user')
            return rest_ok({'user_status':ACTIVE})
        return rest_ok({'user_status':user['ustatus']})

    async def check_otp(self, request):
        """ Returns the users_status regarding the OTP, searching by user email and OTP. To be used before asking for OTP
            Checks for OTP validity, timeout period and limits number of fails """
        data = await request.json()
        [email,otp] = enforce_required_args(self.LOG, data, ['email','otp'], label='AUTH', label2='check_otp', as_list=True)
        user = await self.db.select_one('users', {'email': email})
        if user is None:
            return rest_ok({'user_status':NOT_REGISTERED})
        if user['ustatus']!=WAIT_SET_PASSWD is None:
            return rest_ok({'user_status':OTP_OK})
        if otp!=user['otp']:
            self.LOG(3, 0, label='AUTH', label2='check_otp', msg={'user':email, 'err_message':'OTP error'})
            if await self.increase_fails(user):
                return rest_ok({'user_status':SUSPENDED_FAILS})                
            return rest_ok({'user_status':OTP_ERROR})
        delta = get_age_in_seconds(user['ts_passwd'])
        if delta > self.otp_timeout:
            self.LOG(3, 0, label='AUTH', label2='OTP', msg=f'OTP expired (delta={delta})')
            return await self.gen_otp(user, OTP_EXPIRED)
        return rest_ok({'user_status':OTP_OK})

    async def login(self, request):
        data = await request.json()
        [email,passwd] = enforce_required_args(self.LOG, data, ['email', 'passwd'], label='AUTH', label2='login', as_list=True)
        user = await self.db.select_one('users', {'email':email})
        if user is None:
            self.LOG(3, 0, label='AUTH', label2='login', msg={'user':email, 'error_message':'user not registered'})
            return rest_ok({'user_status':NOT_REGISTERED})
        if user['ustatus'] in [NOT_REGISTERED, WAIT_SET_PASSWD, SUSPENDED, SUSPENDED_LICENCE]:
            self.LOG(3, 0, label='AUTH', label2='login', msg={'user':email, 'error_message': 'user not active'})
            return rest_ok({'user_status':user['ustatus']})
        if not bcrypt.checkpw(passwd.encode(), user['passwd'].encode()):
            self.LOG(3, 0, label='AUTH', label2='login', msg={'user': email, 'error_message': 'wrong password'})
            if await self.increase_fails(user):
                return rest_ok({'user_status': SUSPENDED_FAILS})                
            return rest_ok({'user_status': WRONG_PASSWD})
        jwt_token = mk_jwt_token(user['id'], email, self.auth_secret)
        increase(user, 'logins', True)
        user['fails'] = 0
        await self.db_update('users', 'email', user, 'login')
        return rest_ok({'auth_status':'AUTHORIZED', 'auth_msg':AUTH_MSGS['AUTHORIZED'], 'jwt_token': jwt_token, 'user_info': get_user_info(user)})
        # return rest_ok({'user_status': AUTHORIZED, 'jwt_token': jwt_token, 'user_info': get_user_info(user)})

    async def recover_passwd(self, request):
        data = await request.json()
        [email] = enforce_required_args(self.LOG, data, ['email'], label='AUTH', label2='recover_passwd', as_list=True)
        user = await self.enforce_find_user({'email':email}, label2='recover_passwd')
        user['ustatus'] = WAIT_SET_PASSWD
        await self.db_update('users', 'email', user, 'recover_passwd')
        await self.gen_otp(user, WAIT_SET_PASSWD)
        return rest_ok({'user_status': WAIT_SET_PASSWD})

    async def set_passwd(self, request):
        data = await request.json()
        [email,otp,passwd] = enforce_required_args(self.LOG, data, ['email','otp','passwd'], label='AUTH', label2='recover_passwd', as_list=True)
        user = await self.enforce_find_user({'email':email}, label2='set_passwd')
        if user['ustatus']!=WAIT_SET_PASSWD or otp!=user['otp']:
            self.LOG(3, 0, label='AUTH', label2='set_passwd', msg={'user': email, 'error_message': 'OTP error'})
            if await self.increase_fails(user):
                return rest_ok({'user_status': SUSPENDED_FAILS})                
            return rest_ok({'user_status': OTP_ERROR})
        delta = get_age_in_seconds(user['ts_passwd'])
        if delta > self.otp_timeout:
            self.LOG(3, 0, label='AUTH', label2='OTP', msg=f'OTP expired (delta={delta})')
            return await self.gen_otp(user, OTP_EXPIRED)
        hashed = bcrypt.hashpw(passwd.encode(), bcrypt.gensalt())
        user['passwd'],user['ts_passwd'],user['otp'],user['ustatus'] = hashed.decode(),datetime.now(timezone.utc),'',ACTIVE
        increase(user, 'logins')
        await self.db_update('users', 'email', user, 'set_passwd')
        jwt_token = mk_jwt_token(user['id'], user['email'], self.auth_secret)
        return rest_ok({'user_status': AUTHORIZED, 'jwt_token': jwt_token})

    async def register(self, request):
        org = None
        data = await request.json()
        if not self.allow_register:
            role,org = await self.enforce_admin_perm(request, label2='register')
            if role=='admin': org = user['org']
        [email] = enforce_required_args(self.LOG, data, ['email'], label='AUTH', label2='recover_passwd', as_list=True)
        user = await self.db.select_one('users', {'email':email})
        if user is not None:
            return rest_ok({'user_status': USER_ALREADY_REGISTERED})
        user = {
            'email': email,
            'ustatus': NOT_REGISTERED,
            'role': 'user',
            'org' :org,
            'ts_insert': datetime.now(timezone.utc),
            'logins' :0,
            'fails': 0,
            'requests': 0
        }
        proceed_to_set_passwd = self.allow_register and user['role']=='user'
        user['ustatus'] = WAIT_SET_PASSWD if proceed_to_set_passwd else NOT_REGISTERED
        await self.db.insert('users', user)
        if proceed_to_set_passwd:
            return await self.gen_otp(user, WAIT_SET_PASSWD)
        return rest_ok({'status': 'OK'})

    async def change_passwd(self, request):
        user = await self.enforce_find_user({'id':request.user.display_name}, label2='change_passwd')
        data = await request.json()
        [passwd,new_passwd] = enforce_required_args(self.LOG, data, ['passwd','new_passwd'], label='AUTH', label2='recover_passwd', as_list=True)
        if not bcrypt.checkpw(passwd.encode(), user['passwd'].encode()):
            return rest_ok({'user_status': WRONG_PASSWD})
        if not check_passwd_strengh(passwd):
            return rest_ok({'user_status': PASSWD_NOT_STRONG})
        hashed = bcrypt.hashpw(new_passwd.encode(), bcrypt.gensalt())
        user['passwd'],user['otp'],user['ustatus'] = hashed.decode(),'',ACTIVE
        await self.db_update('users', 'email', user, 'change_passwd')
        return rest_ok({'user_status': PASSWORD_CHANGED})

    # ####################################################################################################
    # REST methods for admin: all to return json
    # ####################################################################################################

    async def get_users(self, request):
        data = await request.json()
        role,org = await self.enforce_admin_perm(request, label2='get_users')
        data = await request.json()
        QARGS = ['cols','filters','sfilters','sort','ascending','offset','limit']
        kwargs = {k:data[k] for k in data if k in QARGS}
        filters =  kwargs['filters'] if 'filters' in kwargs and kwargs['filters'] is not None else []
        if role=='admin': filters.append({'col':'org', 'val':org})
        kwargs['filters'] = filters
        res = await self.db.select('users', **kwargs)
        if res is None or not 'data' in res or not 'nrows' in res:
            log_and_raise_service(self.LOG, label='AUTH', label2='get_users', msg='DB access error')
        return OrjsonResponse({'status': 'OK', 'data':res['data'], 'nrows':res['nrows']})

    async def update_user(self, request):
        role,org = await self.enforce_admin_perm(request, label2='update_user')
        data = await request.json()
        user = await self.enforce_find_user({'id': data['id']}, label2='update_user')
        if role=='admin' and user['org']!=org:
            return rest_ok({'status': 'NOT AUTHORIZED'})
        await self.db_update('users', 'id', data, 'update_user', 'update_user')
        return rest_ok({'status': 'OK'})

    async def reset_user(self, request):
        role,org = await self.enforce_admin_perm(request, label2='reset_user')
        data = await request.json()
        user = await self.enforce_find_user({'id': data['id']}, label2='reset_user')
        if role=='admin' and user['org']!=org:
            return rest_ok({'status': 'NOT AUTHORIZED'})
        user['ustatus'],user['fails'],user['otp'] = 0,0,None
        await self.db_update('users', 'id', user, 'reset_user', 'reset_user')
        return rest_ok({'status': 'OK'})

    async def suspend_user(self, request):
        """ Suspends a user changing ustatus to SUSPENDED or to SUSPENDED_LICENCE if data['ustatus']==SUSPENDED_LICENCE """
        role,org = await self.enforce_admin_perm(request, label2='suspend_user')
        data = await request.json()
        user = await self.enforce_find_user({'id': data['id']}, label2='suspend_user')
        if role=='admin' and user['org']!=org:
            return rest_ok({'status': 'NOT AUTHORIZED'})
        ustatus = int(data['ustatus']) if 'ustatus' in data and int(data['ustatus'])==SUSPENDED_LICENCE else SUSPENDED
        user['ustatus'] = ustatus
        await self.db_update('users', 'id', user, 'suspend_user')
        return rest_ok({'status': 'OK'})

    async def change_org(self, request):
        _,_ = await self.enforce_admin_perm(request, label2='change_org', requires_su=True)
        user = await self.enforce_find_user({'id':request.user.display_name}, label2='change_org')
        data = await request.json()
        user['org'] = data['org']
        jwt_token = mk_jwt_token(user['id'], user['email'], self.auth_secret)
        _ = await self.db.update('users', 'id', user)
        return rest_ok({'status': 'OK', 'jwt_token': jwt_token, 'user_info': get_user_info(user)})
