classquiz-ai/classquiz/routers/editor.py

# SPDX-FileCopyrightText: 2023 Marlon W (Mawoka)
#
# SPDX-License-Identifier: MPL-2.0


import asyncio
import html
import uuid
from typing import Optional

import asyncpg.exceptions
import bleach
from fastapi import APIRouter, HTTPException, Depends
from pydantic import BaseModel

from classquiz.config import settings, redis, storage, meilisearch, ALLOWED_TAGS_FOR_QUIZ, arq
from classquiz.db.models import Quiz, QuizInput, User, QuizQuestionType, StorageItem
from classquiz.auth import get_current_user
import os
from datetime import datetime
from uuid import UUID

from classquiz.helpers import get_meili_data, check_image_string, extract_image_ids_from_quiz
from classquiz.storage.errors import DeletionFailedError

settings = settings()

router = APIRouter()

allowed_image_extensions = [".gif", ".jpg", ".jpeg", ".png", ".svg", ".webp", ".jfif"]


class InitEditorResponse(BaseModel):
    token: str


class EditSessionData(BaseModel):
    quiz_id: UUID
    edit: bool
    user_id: UUID


async def delete_images_for_edit_id(edit_id: str):
    await asyncio.sleep(30)
    res = await redis.lrange(f"edit_session:{edit_id}:images", 0, -1)
    if len(res) != 0:
        for image_id in res:
            await storage.delete([image_id])


@router.post("/start", response_model=InitEditorResponse)
async def init_editor(edit: bool, quiz_id: Optional[UUID] = None, user: User = Depends(get_current_user)):
    if edit and quiz_id is not None and await Quiz.objects.get_or_none(id=quiz_id, user_id=user.id) is None:
        raise HTTPException(status_code=404, detail="Quiz not found")
    if not edit and quiz_id is not None:
        raise HTTPException(status_code=400, detail="You can't choose the id for your quiz")
    if edit and quiz_id is None:
        raise HTTPException(status_code=400, detail="Edit can't be true if quiz_id is None")
    if quiz_id is None:
        quiz_id = uuid.uuid4()
    edit_id = os.urandom(4).hex()
    await redis.sadd("edit_sessions", edit_id)
    await redis.set(
        f"edit_session:{edit_id}", EditSessionData(quiz_id=quiz_id, edit=edit, user_id=user.id).json(), ex=3600
    )
    return InitEditorResponse(token=edit_id)


@router.post("/finish")
async def finish_edit(edit_id: str, quiz_input: QuizInput):
    session_data = await redis.get(f"edit_session:{edit_id}")
    if session_data is None:
        raise HTTPException(status_code=401, detail="Edit ID not found!")
    session_data = EditSessionData.parse_raw(session_data)
    quiz_input.title = html.unescape(bleach.clean(quiz_input.title, tags=ALLOWED_TAGS_FOR_QUIZ, strip=True))
    quiz_input.description = html.unescape(bleach.clean(quiz_input.description, tags=ALLOWED_TAGS_FOR_QUIZ, strip=True))
    if quiz_input.background_color is not None:
        quiz_input.background_color = html.unescape(bleach.clean(quiz_input.background_color, tags=[], strip=True))

    for i, question in enumerate(quiz_input.questions):
        if question.type == QuizQuestionType.ABCD:
            for i2, answer in enumerate(question.answers):
                if answer.color is not None:
                    quiz_input.questions[i].answers[i2].color = html.unescape(
                        bleach.clean(answer.color, tags=[], strip=True)
                    )
                if answer.answer == "":
                    quiz_input.questions[i].answers[i2].answer = None
                if answer.answer is not None:
                    quiz_input.questions[i].answers[i2].answer = html.unescape(
                        bleach.clean(answer.answer, tags=ALLOWED_TAGS_FOR_QUIZ, strip=True)
                    )

    images_to_delete = []
    old_quiz_data: Quiz = await Quiz.objects.get_or_none(id=session_data.quiz_id, user_id=session_data.user_id)

    for i, question in enumerate(quiz_input.questions):
        image = question.image
        quiz_input.questions[i].question = html.unescape(
            bleach.clean(quiz_input.questions[i].question, tags=ALLOWED_TAGS_FOR_QUIZ, strip=True)
        )
        if image == "":
            question.image = None
        if image is not None and not check_image_string(image)[0]:
            raise HTTPException(status_code=400, detail="Image URL(s) aren't valid!")

    if quiz_input.cover_image == "":
        quiz_input.cover_image = None

    if quiz_input.cover_image is not None and not check_image_string(quiz_input.cover_image)[0]:
        raise HTTPException(status_code=400, detail="image url is not valid")

    if quiz_input.background_image is not None and not check_image_string(quiz_input.background_image)[0]:
        raise HTTPException(status_code=400, detail="image url is not valid")

    if session_data.edit:
        await arq.enqueue_job("quiz_update", old_quiz_data, old_quiz_data.id, _defer_by=2)
        quiz = old_quiz_data
        meilisearch.index(settings.meilisearch_index).update_documents([await get_meili_data(quiz)])
        if not quiz_input.public:
            meilisearch.index(settings.meilisearch_index).delete_document(str(quiz.id))
        else:
            meilisearch.index(settings.meilisearch_index).add_documents([await get_meili_data(quiz)])
        quiz.title = quiz_input.title
        quiz.public = quiz_input.public
        quiz.description = quiz_input.description
        quiz.updated_at = datetime.now()
        quiz.questions = quiz_input.dict()["questions"]
        quiz.cover_image = quiz_input.cover_image
        quiz.background_color = quiz_input.background_color
        quiz.background_image = quiz_input.background_image
        for image in images_to_delete:
            if image is not None:
                try:
                    await storage.delete([image])
                except DeletionFailedError:
                    pass
        await redis.srem("edit_sessions", edit_id)
        await redis.delete(f"edit_session:{edit_id}")
        await redis.delete(f"edit_session:{edit_id}:images")
        await quiz.update()
        return quiz
    else:
        quiz = Quiz(
            **quiz_input.dict(),
            user_id=session_data.user_id,
            id=session_data.quiz_id,
            created_at=datetime.now(),
            updated_at=datetime.now(),
        )

        await redis.delete("global_quiz_count")
        if quiz_input.public:
            meilisearch.index(settings.meilisearch_index).add_documents([await get_meili_data(quiz)])
        try:
            await redis.srem("edit_sessions", edit_id)
            await redis.delete(f"edit_session:{edit_id}")
            await redis.delete(f"edit_session:{edit_id}:images")
            await quiz.save()
        except asyncpg.exceptions.UniqueViolationError:
            raise HTTPException(status_code=400, detail="The quiz already exists")
        new_images = extract_image_ids_from_quiz(quiz)
        for image in new_images:
            item = await StorageItem.objects.get_or_none(id=uuid.UUID(image))
            if item is None:
                continue
            await quiz.storageitems.add(item)