import asyncio
import datetime
import json
import logging
import random
import traceback
from functools import partial
from typing import TYPE_CHECKING

import aiohttp.web

from nawah import data as Data
from nawah.classes import app_encoder
from nawah.config import Config

from .._call import call
from .._config import _compile_anon_session, _compile_anon_user
from .._val import camel_to_upper
from ._utils import _close_session

if TYPE_CHECKING:
    from asyncio.futures import Future

    from nawah.types import NawahEnv

logger = logging.getLogger('nawah')


async def _websocket_handler(request: aiohttp.web.Request):
    conn = Data.create_conn()
    logger.debug('Websocket connection starting with client at \'%s\'', request.remote)
    ws = aiohttp.web.WebSocketResponse()
    await ws.prepare(request)

    session_lock_key = random.random()

    while Config.sys.session_lock != session_lock_key:
        if Config.sys.session_lock is None:
            Config.sys.session_lock = session_lock_key

    Config.sys.session_counter += 1

    env: 'NawahEnv' = {
        'id': str(Config.sys.session_counter),
        'conn': conn,
        'REMOTE_ADDR': request.remote or 'localhost',
        'ws': ws,
        'init': False,
        'last_call': datetime.datetime.utcnow(),
        'quota': {
            'counter': Config.quota_anon_min,
            'last_check': datetime.datetime.utcnow(),
        },
        'args': {},
    }
    Config.sys.sessions[str(Config.sys.session_counter)] = env
    Config.sys.session_lock = None

    try:
        env['HTTP_USER_AGENT'] = request.headers['user-agent']
        env['HTTP_ORIGIN'] = request.headers['origin']
    except:
        env['HTTP_USER_AGENT'] = ''
        env['HTTP_ORIGIN'] = ''

    logger.debug(
        'Websocket connection #\'%s\' ready with client at \'%s\'',
        env['id'],
        env['REMOTE_ADDR'],
    )

    await ws.send_str(
        app_encoder.encode(
            {
                'status': 200,
                'msg': 'Connection ready',
                'args': {'code': 'CORE_CONN_READY'},
            }
        )
    )

    async for msg in ws:
        if 'conn' not in env:
            await ws.close()
            break
        logger.debug(
            'Received new message from session #\'%s\': %s', env['id'], msg.data[:256]
        )
        if msg.type == aiohttp.WSMsgType.TEXT:
            logger.debug(
                'Config.sys.ip_quota on session #\'%s\': %s',
                env['id'],
                Config.sys.ip_quota,
            )
            logger.debug(
                'session_quota: on session #\'%s\': %s', env['id'], env['quota']
            )
            # Check for IP quota
            if str(request.remote) not in Config.sys.ip_quota:
                Config.sys.ip_quota[str(request.remote)] = {
                    'counter': Config.quota_ip_min,
                    'last_check': datetime.datetime.utcnow(),
                }
            else:
                if (
                    datetime.datetime.utcnow()
                    - Config.sys.ip_quota[str(request.remote)]['last_check']
                ).seconds > 59:
                    Config.sys.ip_quota[str(request.remote)][
                        'last_check'
                    ] = datetime.datetime.utcnow()
                    Config.sys.ip_quota[str(request.remote)][
                        'counter'
                    ] = Config.quota_ip_min
                else:
                    # If hit quota, deny call
                    if Config.sys.ip_quota[str(request.remote)]['counter'] - 1 <= 0:
                        logger.warning(
                            'Denying Websocket request from \'%s\' for hitting IP quota',
                            request.remote,
                        )
                        asyncio.create_task(
                            _handle_msg(
                                env=env,
                                msg=msg,
                                decline_quota='ip',
                            )
                        )
                        continue
                    # Else, update quota counter
                    Config.sys.ip_quota[str(request.remote)]['counter'] -= 1
            # Check for session quota
            if (datetime.datetime.utcnow() - env['quota']['last_check']).seconds > 59:
                env['quota']['last_check'] = datetime.datetime.utcnow()
                env['quota']['counter'] = (
                    (Config.quota_anon_min - 1)
                    if not env['session']
                    or env['session']['token'] == Config.anon_token
                    else (Config.quota_auth_min - 1)
                )
                asyncio.create_task(_handle_msg(env=env, msg=msg))
            else:
                # If hit quota, deny call
                if env['quota']['counter'] - 1 <= 0:
                    asyncio.create_task(
                        _handle_msg(
                            env=env,
                            msg=msg,
                            decline_quota='session',
                        )
                    )
                    continue
                # Else, update quota counter
                env['quota']['counter'] -= 1
                asyncio.create_task(_handle_msg(env=env, msg=msg))

    if 'id' in env:
        await _close_session(env['id'])

    return ws


async def _handle_msg(
    env: 'NawahEnv',
    msg: aiohttp.WSMessage,
    decline_quota: str = None,
):
    try:
        env['last_call'] = datetime.datetime.utcnow()
        try:
            env['session']['token']
        except Exception:
            anon_user = _compile_anon_user()
            anon_session = _compile_anon_session()
            anon_session['user'] = anon_user
            env['session'] = anon_session
        res = json.loads(msg.data)

        # Check if msg should be denied for quota hit
        if decline_quota == 'ip':
            await env['ws'].send_str(
                app_encoder.encode(
                    {
                        'status': 429,
                        'msg': 'You have hit calls quota from this IP',
                        'args': {
                            'call_id': res['call_id']
                            if 'call_id' in res.keys()
                            else None,
                            'code': 'CORE_REQ_IP_QUOTA_HIT',
                        },
                    }
                )
            )
            return

        if decline_quota == 'session':
            await env['ws'].send_str(
                app_encoder.encode(
                    {
                        'status': 429,
                        'msg': 'You have hit calls quota',
                        'args': {
                            'call_id': res['call_id']
                            if 'call_id' in res.keys()
                            else None,
                            'code': 'CORE_REQ_SESSION_QUOTA_HIT',
                        },
                    }
                )
            )
            return

        logger.debug('Decoded request: %s', app_encoder.encode(res))

        if 'endpoint' not in res.keys():
            await env['ws'].send_str(
                app_encoder.encode(
                    {
                        'status': 400,
                        'msg': 'Request missing endpoint',
                        'args': {
                            'call_id': res['call_id']
                            if 'call_id' in res.keys()
                            else None,
                            'code': 'CORE_REQ_NO_ENDPOINT',
                        },
                    }
                )
            )
            return

        if env['init'] is False:
            if res['endpoint'] != 'conn/verify':
                await env['ws'].send_str(
                    app_encoder.encode(
                        {
                            'status': 1008,
                            'msg': 'Request token is not accepted',
                            'args': {
                                'call_id': res['call_id']
                                if 'call_id' in res.keys()
                                else None,
                                'code': 'CORE_REQ_NO_VERIFY',
                            },
                        }
                    )
                )
                await env['ws'].close()
                return

            # [TODO] Add condition to now allow a connection from a client with client_app = __sys
            if Config.client_apps and (
                'doc' not in res
                or 'app' not in res['doc']
                or res['doc']['app'] not in Config.client_apps
                or (
                    Config.client_apps[res['doc']['app']].type == 'web'
                    and env['HTTP_ORIGIN']
                    not in Config.client_apps[res['doc']['app']].origin
                )
            ):
                await env['ws'].send_str(
                    app_encoder.encode(
                        {
                            'status': 1008,
                            'msg': 'Request token is not accepted',
                            'args': {
                                'call_id': res['call_id']
                                if 'call_id' in res.keys()
                                else None,
                                'code': 'CORE_REQ_NO_VERIFY',
                            },
                        }
                    )
                )
                await env['ws'].close()
                return

            env['init'] = True
            # [TODO] Add condition to now allow a connection from a client with client_app = __sys
            if not Config.client_apps:
                env['client_app'] = '__public'
            else:
                env['client_app'] = res['doc']['app']
            logger.debug('Connection on session #\'%s\' is verified', env['id'])
            # if Config.analytics_events['app_conn_verified']:
            #     asyncio.create_task(
            #         call(
            #             'analytic/create',
            #             skip_events=[Event.PERM],
            #             env=env,
            #             doc={
            #                 'event': 'CONN_VERIFIED',
            #                 'subevent': env['client_app'],
            #                 'args': {
            #                     'REMOTE_ADDR': env['REMOTE_ADDR'],
            #                     'HTTP_USER_AGENT': env['HTTP_USER_AGENT'],
            #                 },
            #             },
            #         )
            #     )
            await env['ws'].send_str(
                app_encoder.encode(
                    {
                        'status': 200,
                        'msg': 'Connection established',
                        'args': {
                            'call_id': res['call_id']
                            if 'call_id' in res.keys()
                            else None,
                            'code': 'CORE_CONN_OK',
                        },
                    }
                )
            )
            return

        if res['endpoint'] == 'conn/close':
            logger.debug(
                f'Received connection close instructions on session #\'{env["id"]}\''
            )
            await env['ws'].close()
            return

        if res['endpoint'] == 'heart/beat':
            logger.debug(f'Received connection heartbeat on session #\'{env["id"]}\'')
            await env['ws'].send_str(
                app_encoder.encode(
                    {
                        'status': 200,
                        'msg': 'Heartbeat received',
                        'args': {
                            'call_id': res['call_id']
                            if 'call_id' in res.keys()
                            else None,
                            'code': 'CORE_HEARTBEAT_OK',
                        },
                    }
                )
            )
            return

        res['endpoint'] = res['endpoint']

        if (
            res['endpoint'] in ['session/auth', 'session/reauth']
            and str(env['session']['_id']) != 'f00000000000000000000012'
        ):
            await env['ws'].send_str(
                app_encoder.encode(
                    {
                        'status': 400,
                        'msg': 'You are already authed',
                        'args': {
                            'call_id': res['call_id']
                            if 'call_id' in res.keys()
                            else None,
                            'code': 'CORE_SESSION_ALREADY_AUTHED',
                        },
                    }
                )
            )
            return

        if (
            res['endpoint'] == 'session/signout'
            and str(env['session']['_id']) == 'f00000000000000000000012'
        ):
            await env['ws'].send_str(
                app_encoder.encode(
                    {
                        'status': 400,
                        'msg': 'Singout is not allowed for \'__ANON\' user',
                        'args': {
                            'call_id': res['call_id']
                            if 'call_id' in res.keys()
                            else None,
                            'code': 'CORE_SESSION_ANON_SIGNOUT',
                        },
                    }
                )
            )
            return

        if 'query' not in res.keys():
            res['query'] = []
        if 'doc' not in res.keys():
            res['doc'] = {}
        if 'call_id' not in res.keys():
            res['call_id'] = ''

        request = {
            'call_id': res['call_id'],
            # 'sid': res['sid'] or False,
            'query': res['query'],
            'doc': res['doc'],
            'path': res['endpoint'].split('/'),
        }

        if len(request['path']) != 2:
            await env['ws'].send_str(
                app_encoder.encode(
                    {
                        'status': 400,
                        'msg': 'Endpoint path is invalid',
                        'args': {
                            'call_id': request['call_id'],
                            'code': 'CORE_REQ_INVALID_PATH',
                        },
                    }
                )
            )
            return

        module = request['path'][0]

        if module not in Config.modules:
            await env['ws'].send_str(
                app_encoder.encode(
                    {
                        'status': 400,
                        'msg': 'Endpoint module is invalid',
                        'args': {
                            'call_id': request['call_id'],
                            'code': 'CORE_REQ_INVALID_MODULE',
                        },
                    }
                )
            )
            return

        if request['path'][1] not in Config.modules[module].funcs:
            await env['ws'].send_str(
                app_encoder.encode(
                    {
                        'status': 400,
                        'msg': 'Endpoint is invalid',
                        'args': {
                            'call_id': request['call_id'],
                            'code': 'CORE_REQ_INVALID_METHOD',
                        },
                    }
                )
            )
            return

        if Config.modules[module].funcs[request['path'][1]].get_func:
            await env['ws'].send_str(
                app_encoder.encode(
                    {
                        'status': 400,
                        'msg': 'Endpoint is a GET Function',
                        'args': {
                            'call_id': request['call_id'],
                            'code': 'CORE_REQ_GET_FUNC',
                        },
                    }
                )
            )
            return

        # if not request['sid']:
        request['sid'] = 'f00000000000000000000012'

        query = request['query']
        doc = request['doc']
        # await _process_file_obj(doc=doc, env=env)
        call_task = asyncio.create_task(
            call(
                f'{module}/{request["path"][1]}',
                skip_events=[],
                env=env,
                query=query,
                doc=doc,
            )
        )
        call_task.add_done_callback(
            partial(_call_callback, env=env, call_id=request['call_id'])
        )

    except Exception as e:
        logger.error(f'An error occurred. Details: {traceback.format_exc()}')
        if Config.debug:
            await env['ws'].send_str(
                app_encoder.encode(
                    {
                        'status': 500,
                        'msg': f'Unexpected error has occurred [{str(e)}]',
                        'args': {
                            'code': 'CORE_SERVER_ERROR',
                            'err': str(e),
                            'call_id': request['call_id'],
                        },
                    }
                )
            )
        else:
            await env['ws'].send_str(
                app_encoder.encode(
                    {
                        'status': 500,
                        'msg': 'Unexpected error has occurred',
                        'args': {
                            'code': 'CORE_SERVER_ERROR',
                            'call_id': request['call_id'],
                        },
                    }
                )
            )


def _call_callback(call_task: 'Future', env: 'NawahEnv', call_id: str):
    try:
        results = call_task.result()
        results['args']['call_id'] = call_id
        asyncio.create_task(env['ws'].send_str(app_encoder.encode(results)))
    except Exception as e:
        status = getattr(e, 'status', 500)
        msg = 'Unexpected error has occurred'
        code = camel_to_upper(e.__class__.__name__)
        if e.args:
            msg = e.args[0]
            if len(e.args) == 2:
                if isinstance(e.args[1], dict):
                    if 'args' in e.args[1]:
                        if 'code' in e.args[1]['args']:
                            code = e.args[1]['args']['code']
                    if 'status' in e.args[1]:
                        status = e.args[1]['status']

        asyncio.create_task(
            env['ws'].send_str(
                app_encoder.encode(
                    {
                        'status': status,
                        'msg': msg,
                        'args': {'code': code, 'call_id': call_id},
                    }
                )
            )
        )
