generateur_v3/backend/api/routes/room/routes.py

from pydantic import Field
from pydantic import create_model_from_typeddict
from pydantic.error_wrappers import ValidationError
from pydantic import validate_arguments
from services.database import generate_unique_code
from sqlmodel import col
from sqlmodel import select
from pydantic import BaseModel
from typing import Any, Callable, Dict, List, Optional
from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect, status, Query
from config import ALGORITHM, SECRET_KEY
from database.auth.crud import get_user_from_clientId_db
from database.auth.models import User
from database.db import get_session
from database.room.crud import check_room, create_room_db, userInRoom
from sqlmodel import Session
from database.room.models import Anonymous, MemberSerializer, MemberWithRelations, MemberRead, Room, RoomCreate, RoomRead, Member
from services.auth import get_current_user_optional
from fastapi.exceptions import HTTPException
from jose import jwt, exceptions
import inspect
router = APIRouter(tags=["room"])


class RoomAndMember(BaseModel):
    room: RoomRead
    member: MemberRead


class Waiter(BaseModel):
    username: str
    waiter_id: str


def serialize_member(member: Member) -> MemberRead | Waiter:
    member_obj = member.user or member.anonymous
    if member.waiting == False:
        return MemberRead(username=member_obj.username, reconnect_code=getattr(member_obj, "reconnect_code", ""), isUser=member.user_id != None, isAdmin=member.is_admin)
    if member.waiting == True:
        return Waiter(username=member_obj.username, waiter_id=member.waiter_code)


@router.post('/room', response_model=RoomAndMember)
def create_room(room: RoomCreate, username: Optional[str] = Query(default=None, max_length=20), user: User | None = Depends(get_current_user_optional), db: Session = Depends(get_session)):
    room_obj = create_room_db(room=room, user=user, username=username, db=db)
    return {'room': room_obj['room'], "member": serialize_member(room_obj['member'])}


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

    def add(self, group: str, ws: WebSocket):
        if group not in self.active_connections:
            self.active_connections[group] = []

        if ws not in self.active_connections[group]:
            self.active_connections[group].append(ws)

    def remove(self, group: str, ws: WebSocket):
        if group in self.active_connections:
            if ws in self.active_connections[group]:
                self.active_connections[group].remove(ws)

    async def broadcast(self, message, group: str, exclude: list[WebSocket] = []):
        if group in self.active_connections:
            for connection in self.active_connections[group]:
                if connection not in exclude:
                    await connection.send_json(message)


def make_event_decorator(eventsDict):
    def _(name: str, conditions: List[Callable | bool] = []):
        def add_event(func):
            model = validate_arguments(func).model
            eventsDict[name] = {"func": func,
                                "conditions": conditions, "model": model}
            return func
        return add_event
    return _


class Event(BaseModel):
    func: Callable
    conditions: List[Callable | bool]
    model: BaseModel


def dict_model(model: BaseModel, exclude: List[str]):
    value = {}
    for n, f in model:
        if n not in exclude:
            value[n] = f
    return value


def dict_all(obj: Any):
    if isinstance(obj, dict):
        value = {}
        for k, v in obj.items():
            if isinstance(v, dict):
                v = dict_all(v)
                value[k] = dict(v)
            elif isinstance(v, BaseModel):
                value[k] = dict(v)
            else:
                try:
                    value[k] = dict(v)
                except:
                    value[k] = v
        return value
    return dict(obj)


class Consumer:
    events: Dict[str, Event] = {}
    sendings: Dict[str, Any] = {}
    event = make_event_decorator(events)
    sending = make_event_decorator(sendings)

    def __init__(self, ws: WebSocket):
        self.ws: WebSocket = ws
        #self.events: Dict[str, Callable] = {}

    async def connect(self):
        pass

    async def validation_error_handler(self, e: ValidationError):
        errors = e.errors()
        await self.ws.send_json({"type": "error", "data": {"detail": [{ers['loc'][-1]: ers['msg']} for ers in errors]}})

    async def send(self, payload):
        type = payload.get('type', None)
        print('TYPE', type, self.member.id)
        if type is not None:
            event_wrapper = self.sendings.get(type, None)
            if event_wrapper is not None:
                handler = event_wrapper.get('func')
                conditions = event_wrapper.get('conditions')

                is_valid = all([(await c(self)) if inspect.iscoroutinefunction(c) else c(self) if inspect.isfunction(c) else c == True if isinstance(c, bool) else True for c in conditions])

                if handler is not None and is_valid:
                    model = event_wrapper.get("model")

                    data = payload.get('data') or {}
                    try:
                        validated_payload = model(self=self, **data)
                    except ValidationError as e:
                        await self.ws.send_json({"type": "error", "data": {"msg": "Oops there was an error"}})
                        return

                    validated_payload = dict_model(validated_payload,
                                                   exclude=["v__duplicate_kwargs", "args", 'kwargs', "self"])
                    try:
                        parsed_payload = handler(
                        self, **validated_payload)
                    
                            
                        await self.ws.send_json({'type': type, "data": dict_all(parsed_payload)})
                        return
                    except Exception as e:
                        
                        print('NOPE', self.member.id, e)
                        return
                return
        print('pls')
        await self.ws.send_json(payload)
        print('sent')

    async def receive(self, data):
        event = data.get('type', None)
        if event is not None:
            event_wrapper = self.events.get(event, None)
            if event_wrapper is not None:
                handler = event_wrapper.get('func')
                conditions = event_wrapper.get('conditions')

                is_valid = all([(await c(self)) if inspect.iscoroutinefunction(c) else c(self) if inspect.isfunction(c) else c == True if isinstance(c, bool) else True for c in conditions])

                if handler is not None and is_valid:
                    model = event_wrapper.get("model")

                    payload = data.get('data') or {}
                    try:
                        validated_payload = model(self=self, **payload)
                    except ValidationError as e:
                        await self.validation_error_handler(e)
                        return

                    await handler(**{k: v for k, v in validated_payload.dict().items() if k not in ["v__duplicate_kwargs", "args", 'kwargs']})

    async def disconnect(self):
        pass

    async def run(self):
        await self.connect()
        try:
            while True:
                data = await self.ws.receive_json()
                await self.receive(data)
        except WebSocketDisconnect:
            await self.disconnect()


class ConsumerManager:
    def __init__(self):
        self.active_connections: Dict[str, List[Consumer]] = {}

    def add(self, group: str, ws: Consumer):

        if group not in self.active_connections:
            self.active_connections[group] = []
        print("adding", ws, self.active_connections[group])
        if ws not in self.active_connections[group]:
            print('ACTUALLY ADDING')
            self.active_connections[group].append(ws)

    def remove(self, group: str, ws: Consumer):
        if group in self.active_connections:
            if ws in self.active_connections[group]:
                self.active_connections[group].remove(ws)

    async def broadcast(self, message, group: str, exclude: list[Consumer] = []):
        if group in self.active_connections:
            print(self.active_connections[group], exclude)
            for connection in list(set(self.active_connections[group])):
                if connection not in exclude:
                    print('SEND TO', connection, message)
                    await connection.send(message)


manager = ConsumerManager()


class Token(BaseModel):
    token: str


def get_user_from_token(token: str, db: Session):
    try:
        decoded = jwt.decode(token=token, key=SECRET_KEY,
                             algorithms=[ALGORITHM])
    except exceptions.ExpiredSignatureError:
        return False

    clientId = decoded.get('sub')
    return get_user_from_clientId_db(clientId=clientId, db=db)


def get_member_from_user(user_id: int, room_id: int, db: Session):
    member = db.exec(select(Member).where(Member.room_id ==
                     room_id, Member.user_id == user_id)).first()
    return member


def get_member_from_token(token: str, room_id: int, db: Session):
    user = get_user_from_token(token, db)
    if user is False:
        return False
    if user is None:
        return None
    member = get_member_from_user(user.id, room_id, db)
    return member


def get_member_from_anonymous(anonymous_id: int, room_id: int, db: Session):
    member = db.exec(select(Member).where(Member.room_id ==
                     room_id, Member.anonymous_id == anonymous_id)).first()
    return member


def get_member_reconnect_code(reconnect_code: str, room_id: int, db: Session):
    anonymous = get_anonymous_from_code(reconnect_code, db)
    if anonymous is None:
        return None
    member = get_member_from_anonymous(anonymous.id, room_id, db)
    return member


def get_anonymous_from_code(reconnect_code: str, db: Session):
    anonymous = db.exec(select(Anonymous).where(
        Anonymous.reconnect_code == reconnect_code)).first()
    return anonymous


def connect_member(member: Member, db: Session):
    member.online = True
    db.add(member)
    db.commit()
    db.refresh(member)
    return member


def disconnect_member(member: Member, db: Session):
    if member.waiting == False:
        member.online = False
        db.add(member)
        db.commit()
        db.refresh(member)
        return member
    else:
        db.delete(member)
        db.commit()
        return member


def validate_username(username: str, room: Room,  db: Session = Depends(get_session)):
    print('VALIDATE', username)
    if len(username) > 20:
        return None
    members = select(Member.anonymous_id).where(
        Member.room_id == room.id, Member.anonymous_id != None)
    anonymous = select(Anonymous).where(
        col(Anonymous.id).in_(members), Anonymous.username == username)
    username_anonymous = db.exec(anonymous).first()
    return None if username_anonymous is not None else username


def create_anonymous_member(username: str, room: Room, db: Session):
    username = validate_username(username)
    if username is None:
        return None
    reconnect_code = generate_unique_code(
        Anonymous, s=db, field_name="reconnect_code")
    anonymous = Anonymous(username=username, reconnect_code=reconnect_code)
    member = Member(room=room, anonymous=anonymous)
    db.add(member)
    db.commit()
    db.refresh(member)
    return member


def create_user_member(user: User, room: Room, db: Session):
    member = get_member_from_user(user.id, room.id, db)
    if member is not None:
        return None
    member = Member(room=room, user=user)
    db.add(member)
    db.commit()
    db.refresh(member)
    return member


def create_anonymous_waiter(username: str, room: Room, db: Session):
    username = validate_username(username, room, db)
    if username is None:
        return None
    reconnect_code = generate_unique_code(
        Anonymous, s=db, field_name="reconnect_code")
    anonymous = Anonymous(username=username, reconnect_code=reconnect_code)

    waiter_code = generate_unique_code(Member, s=db, field_name="waiter_code")
    member = Member(room=room, anonymous=anonymous,
                    waiting=True, waiter_code=waiter_code)
    db.add(member)
    db.commit()
    db.refresh(member)
    return member


def create_user_waiter(user: User, room: Room, db: Session):
    member = get_member_from_user(user.id, room.id, db)
    if member is not None:
        return None
    waiter_code = generate_unique_code(Member, s=db, field_name="waiter_code")
    member = Member(room=room, user=user, waiting=True,
                    waiter_code=waiter_code)
    db.add(member)
    db.commit()
    db.refresh(member)
    return member


def get_waiter(waiter_code: str, db: Session):
    return db.exec(select(Member).where(Member.waiter_code == waiter_code)).first()


def accept_waiter(member: Member, db: Session):
    member.waiting = False
    member.waiter_code = None
    db.add(member)
    db.commit()
    db.refresh(member)
    return member

def refuse_waiter(member: Member, db: Session):
    db.delete(member)
    db.commit()
    return None


class RoomConsumer(Consumer):

    def __init__(self, ws: WebSocket, room: Room, manager: ConsumerManager, db: Session):
        self.room = room
        self.ws = ws
        self.manager = manager
        self.db = db
        self.member = None
    # WS Utilities

    async def connect(self):
        await self.ws.accept()

    async def direct_send(self, type: str, payload: Any):
        await self.ws.send_json({'type': type, "data": payload})

    async def send_to_all_room(self, type: str, payload: Any, exclude: bool = False):
        await self.manager.broadcast({'type': type, "data": payload}, f'{self.room.id}__member', [exclude == True and self.ws])
        await self.manager.broadcast({'type': type, "data": payload}, f'{self.room.id}__admin', [exclude == True and self.ws])

    async def send_to_admin(self, type: str, payload: Any, exclude: bool = False):
        await self.manager.broadcast({'type': type, "data": payload}, f'{self.room.id}__admin', [exclude == True and self.ws])

    async def send_to_members(self, type: str, payload: Any, exclude: bool = False):
        await self.manager.broadcast({'type': type, "data": payload}, f'{self.room.id}__member', [exclude == True and self.ws])

    async def broadcast(self, type, payload, exclude= False):
        await self.manager.broadcast({"type": type, "data": payload}, self.room.id, exclude = [exclude == True and self])

    def add_to_admin(self):
        self.manager.add(f'{self.room.id}__admin', self.ws)

    def add_to_members(self):
        self.manager.add(f'{self.room.id}__members', self.ws)

    def add_to_group(self):
        if self.member.waiting == True:
            self.manager.add(f'waiter__{self.member.waiter_code}', self)
        self.manager.add(self.room.id, self)

    async def connect_self(self):
        if isinstance(self.member, Member):
            connect_member(self.member, self.db)
            await self.broadcast(type="connect", payload={"member": serialize_member(self.member).dict()})

    async def disconnect_self(self):
        if isinstance(self.member, Member):
            disconnect_member(self.member, self.db)
            await self.broadcast(type="disconnect", payload={"member": serialize_member(self.member).dict()})

    # DB Utilities

    # Received Events
    @Consumer.event('login')
    async def login(self, token: str | None = None, reconnect_code: str | None = None):
        if token is not None:
            member = get_member_from_token(token, self.room.id, self.db)
            if member == False:
                await self.direct_send(type="error", payload={"msg": "Token expired"})
                return
            if member is None:
                await self.direct_send(type="error", payload={"msg": "Utilisateur introuvable dans cette salle"})
                return
            self.member = member

            # await self.connect_self()
            self.add_to_group()
            await self.direct_send(type="loggedIn", payload={"member": serialize_member(self.member).dict()})

        elif reconnect_code is not None:
            member = get_member_reconnect_code(
                reconnect_code, self.room.id, db=self.db)
            if member is None:
                await self.direct_send(type="error", payload={"msg": "Utilisateur introuvable dans cette salle"})
                return

            self.member = member

            # await self.connect_self()
            self.add_to_group()
            await self.direct_send(type="loggedIn", payload={"member": serialize_member(self.member).dict()})
        if reconnect_code is None and token is None:
            await self.direct_send(type="error", payload={"msg": "Veuillez spécifier une méthode de connection"})

    @Consumer.event('join')
    async def join(self, token: str | None = None, username: str | None = None):
        if self.room.public == False:
            if token is not None:
                user = get_user_from_token(token, self.db)

                if user is None:
                    await self.direct_send(type="error", payload={"msg": "Utilisateur introuvable"})
                    return
                if user is False:
                    await self.direct_send(type="error", payload={"msg": "Token expired"})
                    return

                waiter = create_user_waiter(user, self.room, self.db)

                if waiter.waiting is False:
                    self.member = waiter
                    # await self.connect_self()
                    self.add_to_group()
                    await self.direct_send(type="loggedIn", payload={"member": serialize_member(self.member).dict()})
                    return

                self.member = waiter
                self.add_to_group()
                await self.direct_send(type="waiting", payload={"waiter": serialize_member(self.member).dict()})
                await self.broadcast(type="waiter", payload={"waiter": serialize_member(self.member).dict()})

            elif username is not None:
                waiter = create_anonymous_waiter(username, self.room, self.db)
                if waiter is None:
                    await self.direct_send(type="error", payload={"msg": "Nom d'utilisateur invalide ou indisponible"})
                    return
                self.member = waiter
                self.add_to_group()
                await self.direct_send(type="waiting", payload={"waiter": serialize_member(self.member).dict()})
                await self.broadcast(type="waiter", payload={"waiter": serialize_member(self.member).dict()})
        else:
            if token is not None:
                user = get_user_from_token(token, self.db)
                if user is None:
                    return
                if user is False:
                    return

                member = create_user_member(user, self.room, self.db)
                if member is None:
                    return
                self.member = member
                self.add_to_group()
                await self.direct_send()
            elif username is not None:
                member = create_anonymous_member(username, self.room, self.db)
                if member is None:
                    await self.direct_send(type="error", data={"msg": "Nom d'utilisateur indisponible"})
                    return
                self.member = member
                self.add_to_group()
                await self.direct_send()

    def isAdmin(self):
        return self.member is not None and self.member.is_admin == True

    @Consumer.event('accept', conditions=[isAdmin])
    async def accept(self, waiter_id: str):
        waiter = get_waiter(waiter_id, self.db)
        member = accept_waiter(waiter, self.db)
        await self.manager.broadcast({"type": "accepted", "data": {'member': serialize_member(member).dict()}}, f"waiter__{waiter_id}")
        await self.broadcast(type="joined", payload={"member": serialize_member(member).dict()})
    
    @Consumer.event('refuse', conditions=[isAdmin])
    async def accept(self, waiter_id: str):
        waiter = get_waiter(waiter_id, self.db)
        member = refuse_waiter(waiter, self.db)
        await self.manager.broadcast({"type": "refused", "data": {'waiter_id': waiter_id}}, f"waiter__{waiter_id}")
        #await self.broadcast(type="joined", payload={"member": serialize_member(member).dict()})

    @Consumer.event('ping_room')
    async def proom(self):
        await self.broadcast(type='ping', payload={})

    def isMember(self):
        return self.member is not None and self.member.waiting == False

    # Sending Events
    @Consumer.sending("joined", conditions=[isMember])
    def joined(self, member: MemberRead):
        print('MEMBER', self.member, member)
        if (self.member.user is not None and member.username == self.member.user.username) or (self.member.anonymous and member.reconnect_code == self.member.anonymous.reconnect_code):
            raise ValueError("Nope")
        if self.member.is_admin == False:
            member.reconnect_code = ""
        return {"member": member}

    @Consumer.sending('waiter', conditions=[isAdmin])
    def waiter(self, waiter: Waiter):
        return {"waiter": waiter}

    @Consumer.sending('ping', conditions=[isMember])
    def ping(self):
        return {}

    async def disconnect(self):
        self.manager.remove(self.room.id, self)
        #await self.disconnect_self() 


@router.websocket('/ws/room/{room_id}')
async def room_ws(ws: WebSocket, room: Room | None = Depends(check_room), db: Session = Depends(get_session)):
    if room is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail='Room not found')
    consumer = RoomConsumer(ws=ws, room=room, manager=manager, db=db)
    await consumer.run()


class TestConsumer(Consumer):
    async def connect(self):
        await self.ws.accept()

    def test(self):
        return True

    @Consumer.event("test", conditions=[True, test])
    async def testering(self):
        await self.ws.send_json({"type": "success"})
        await self.send({"type": "test", "data": {"i": {"username": "lilian", "reconnect_code": "something", "isAdmin": False, "isUser": False}, "test": 12}})
        # await self.send({"type": "test", "data": {"i": {"username": "lilian", "reconnect_code": "something", "isAdmin": False, "isUser": False}}})
        return

    @Consumer.sending('test', conditions=[])
    def sendtest(self, i: MemberRead, test: int):
        print("i", i)
        print(i.reconnect_code)
        print(dict(i))
        i.reconnect_code = "nope"
        return {"i": i, "test": test}


@router.websocket('/ws/test')
async def test(ws: WebSocket):
    consumer = TestConsumer(ws)
    await consumer.run()