# ******************************************************
# |docname| - Provide a scalable synchronous chat system
# ******************************************************
# For effective peer instructio in a mixed or online class we must support some
# kind of chat system.  This may spill over into other use cases but for now the
# peer instruction pages are the main use case.  To provide a near real time chat
# environment we need to use websockets.  This is a new area for us in summer 2021
# creating a websocket chat system is super simple for a toy environment. One server
# process can track all users and their connections, but in a production environment
# with multiple workers it gets much more complicated.  `this article <https://tsh.io/blog/how-to-scale-websocket/>`_
# does a nice job of explaining the issues and solutions.
#
# Our architecture is as follows:
# We will use redis pubsub model to support a fast production environment.
#
# We have one end point called ``websocket_endpoint`` that browsers open a websocket with
# this connection is only for the page to RECEIVE messages.  The endpoint is also a redis
# subscriber to the "peermessages" channel.
#
# We have a second endpoint called ``send_message`` that accepts a json formatted message
# package.  This could be a broadcast text message, or could be a special control mesage
# that allows the instructor to move the students through the peer process.  The ``send_message``
# endpoint is the redis producer.
#
# The producer sends the message into the redis queue and then all copies of the consumer
# look at the message.  If the recipient of that message is connected to that instance
# then the message is sent to the recipient and all other instance ignore that message.
# If the message is a broadcast message then all instances of the consumer forward that
# message to all connected parties.

import json
import os
import time

#
# Third-party imports
# -------------------
from datetime import datetime
from typing import Dict, Optional, Any
from multi_await import multi_await  # type: ignore

import aioredis

# Local application imports
# -------------------------
from fastapi import (
    APIRouter,
    Cookie,
    Query,
    WebSocket,  # Depends,; noqa F401
    status,
)
from fastapi.templating import Jinja2Templates

from ..applogger import rslogger
from ..config import settings
from ..crud import create_useinfo_entry
from ..models import UseinfoValidation
from ..schemas import PeerMessage

# from ..session import auth_manager

# Routing
# =======
# See `APIRouter config` for an explanation of this approach.
# Remove the "discuss" prefix until PR #2640 of FastAPI is merged
router = APIRouter(
    tags=["discuss"],
)

templates = Jinja2Templates(directory=f"{settings._book_server_path}/templates/discuss")


class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}

    async def connect(self, user: str, websocket: WebSocket):
        await websocket.accept()
        self.active_connections[user] = websocket

    def disconnect(self, sockid: str):
        del self.active_connections[sockid]

    async def send_personal_message(
        self,
        receiver: str,
        message: Dict[str, Any],
    ):
        to = receiver
        if to in self.active_connections:
            try:
                rslogger.debug(
                    f"{os.getpid()}: sending {message} to {to} on {self.active_connections[to]}"
                )
                await self.active_connections[to].send_json(message)
            except Exception as e:
                rslogger.error(f"{os.getpid()}: Error sending to {to} is {e}")
                del self.active_connections[to]
        else:
            rslogger.error(
                f"{os.getpid()}: {to} is not connected here {self.active_connections}"
            )

    async def broadcast(self, message: str):
        rslogger.debug(f"{os.getpid()}: {self.active_connections=} {message=}")
        to_remove = []
        for key, connection in self.active_connections.items():
            rslogger.debug(f"{os.getpid()}: sending to {connection}")
            res = None
            try:
                res = await connection.send_json(message)
            except Exception as e:
                rslogger.debug(f"{os.getpid()}: Failed to send {e}")
                to_remove.append(key)
            rslogger.debug(f"{os.getpid()}: result of send = {res}")
        for key in to_remove:
            del self.active_connections[key]


# this is good for prototyping, but we will need to integrate with
# Redis or a DB for production where we have multiple servers
manager = ConnectionManager()
local_users = set()


async def get_cookie_or_token(
    websocket: WebSocket,
    access_token: Optional[str] = Cookie(None),
    user: Optional[str] = Query(None),
):
    rslogger.debug(f"{os.getpid()}: HELLO {access_token=} or {user=}")
    if access_token is None and user is None:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
    return access_token or user


# It seems that ``@router.websocket`` is much better than the documented
# ``websocket_route``
@router.websocket("/chat/{uname}/ws")
async def websocket_endpoint(websocket: WebSocket, uname: str):
    """
    This endpoint is called to establish a websocket connection between
    The browser and the server.  The websocket is persistent as long
    as the process that runs this endpoint is alive.

    Note: This function must return which means it must be throroughly async
    Using a non async library like plain redis-py will not work as the subscriber
    will block
    """
    rslogger.debug(f"{os.getpid()}: IN WEBSOCKET {uname=}")
    username = uname
    # local_users is a global/module variable shared by all  requests served
    # by the same worker process.
    local_users.add(username)
    await manager.connect(username, websocket)
    r = aioredis.from_url(settings.redis_uri)
    subscriber = r.pubsub()

    async def my_get_message():
        return await subscriber.get_message(timeout=1.0, ignore_subscribe_messages=True)

    await subscriber.subscribe("peermessages")

    # multi_await `Docs and links here <http://www.hydrogen18.com/blog/python-await-multiple.html>`_ acts like
    # the generic sockets select statement, returning a result from get as soon
    # as ANY of our async functions have a result available.
    async with multi_await() as mawait:
        mawait.add(my_get_message)
        mawait.add(websocket.receive_text)

        while True:
            # Wait for something to happen
            # get returns two lists one of completions and one of failures
            complete, failures = await mawait.get()
            # i == 0 --> pubsub message
            # i == 1 --> websocket message
            pmess, wsres = complete
            psfail, wsfail = failures

            if wsfail is not None:
                rslogger.error(wsfail)
                # The fail is more than likely a runtime error from a page close or refresh
                # RuntimeError('Cannot call "receive" once a disconnect message has been received.')
                manager.disconnect(username)
                return

            if psfail is not None:
                rslogger.error(psfail)
                # probably do not want to return

            # handle message from the pubsub queue
            if pmess is not None:
                rslogger.debug(f"{os.getpid()}: {pmess=}")
                if pmess["type"] == "message":
                    # This is a message sent into the channel, our stuff is in
                    # the ``data`` field of the redis message
                    data = json.loads(pmess["data"])
                else:
                    rslogger.error(
                        f"{os.getpid()}: unknown message type {pmess['type']}"
                    )
                    continue
                if data["broadcast"]:
                    await manager.broadcast(data)
                else:
                    # because **every** connection runs this same loop
                    # we only want to send a non-broadcast message if
                    # it is to ourself.
                    partner = await r.hget("partnerdb", data["from"])
                    if partner:
                        partner = partner.decode("utf8")
                    else:
                        try:
                            mess = {
                                "type": "text",
                                "from": data["from"],
                                "message": "Could not find a partner for you",
                                "time": time.time(),
                                "broadcast": False,
                                "course_name": data["course_name"],
                                "div_id": data["div_id"],
                            }
                            await manager.send_personal_message(data["from"], mess)
                        except KeyError:
                            rslogger.error(
                                f"Not enough data to construct a message: {data}"
                            )

                        rslogger.error(
                            f"{os.getpid()}: Failed to find a partner for {data['from']}"
                        )
                    if partner == username:
                        await manager.send_personal_message(partner, data)
                        # log the message
                        # todo - we should not log messages that are 'control' messages
                        # These individual control messages update partner and answer
                        if data["type"] != "control":
                            await create_useinfo_entry(
                                UseinfoValidation(
                                    event="sendmessage",
                                    act=f"to:{partner}:{data['message']}",
                                    div_id=data["div_id"],
                                    course_id=data["course_name"],
                                    sid=data["from"],
                                    timestamp=datetime.utcnow(),
                                )
                            )
                    else:
                        rslogger.debug(f"{os.getpid()}: {partner=} is not {username}")

            if wsres is not None:
                rslogger.debug(
                    f"{os.getpid()}: We don't expect in coming websock messages but got: {wsres}"
                )
                # TODO: now that we have multi-await working it *may* be more
                # efficient to go back to sending all peer messages by websocket
                # and putting them into the subscriber queue here rather than
                # having a separate endpoint.


@router.post("/send_message")
async def send_message(packet: PeerMessage):
    r = await aioredis.from_url(settings.redis_uri)
    r.publish("peermessages", packet.json())
