classes/base_models.py

import nh3

from constants import (
    MAX_COMMENT_LENGTH,
    MAX_FILE_NAME_LENGTH,
    MAX_ID_LENGTH,
    MAX_MESSAGE_LENGTH,
    MAX_RESPONSE_LENGTH,
)
from pydantic import BaseModel, Field, field_validator
from typing import Literal, Set
from uuid import UUID


class IdentifierBase(BaseModel):
    user_id: str = Field(
        pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
    )
    # Participant ID could be in ProfileBase instead. It doesn't really matter.
    participant_id: str = Field(
        pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
    )
    session_id: str = Field(
        pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
    )


class ProfileBase(BaseModel):
    consent: bool
    age_group: Literal["0-18", "18-24", "25-34", "35-44", "45-54", "55-64", "65+"]
    gender: Literal["M", "F"]
    roles: Set[
        Literal["patient", "clinician", "computer-scientist", "researcher", "other"]
    ] = Field(min_length=1, max_length=5)


class ChatRequest(IdentifierBase, ProfileBase):
    conversation_id: str = Field(
        pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
    )
    model_type: Literal[
        "champ",
        "openai",
        "google-conservative",
        "google-creative",
        "qwen",
        "fake",
        "skills_wiki",
    ]
    lang: Literal["en", "fr"]
    human_message: str = Field(min_length=1, max_length=MAX_MESSAGE_LENGTH)

    @field_validator("human_message")
    def sanitize_human_message(cls, human_message: str):
        """Remove HTML tags to prevent XSS"""
        return nh3.clean(human_message)


class FeedbackRequest(IdentifierBase, ProfileBase):
    message_index: int = Field(ge=0, le=10_000)
    rating: Literal["like", "dislike", "mixed"]
    comment: str = Field(min_length=0, max_length=MAX_COMMENT_LENGTH)
    reply_content: str = Field(min_length=1, max_length=MAX_RESPONSE_LENGTH)
    reply_id: UUID

    @field_validator("comment")
    def sanitize_comment(cls, comment: str):
        """Remove HTML tags to prevent XSS"""
        return nh3.clean(comment)

    @field_validator("reply_content")
    def sanitize_reply_content(cls, reply_content: str):
        """Remove HTML tags to prevent XSS"""
        return nh3.clean(reply_content)


class CommentRequest(IdentifierBase, ProfileBase):
    comment: str = Field(min_length=1, max_length=MAX_COMMENT_LENGTH)

    @field_validator("comment")
    def sanitize_comment(cls, comment: str):
        """Remove HTML tags to prevent XSS"""
        return nh3.clean(comment)


class DeleteFileRequest(IdentifierBase, ProfileBase):
    file_name: str = Field(
        # Matches is_valid_filename in helpers/file_helper.py: block only
        # control characters, path separators, and HTML/attribute-injection
        # characters — everything else a real filename can contain (commas,
        # apostrophes, colons, etc.) is allowed, so anything accepted on
        # upload can also be deleted. The "." / ".." exclusion lives in the
        # validator below — pydantic-core's regex engine has no look-ahead.
        pattern=r'^[^\x00-\x1f\\/<>"]+$',
        min_length=1,
        max_length=MAX_FILE_NAME_LENGTH,
    )

    @field_validator("file_name")
    def reject_dot_names(cls, file_name: str):
        if file_name in {".", ".."}:
            raise ValueError('file_name cannot be "." or ".."')
        return file_name


class ClearConversationRequest(BaseModel):
    old_session_id: str = Field(
        pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
    )
    new_session_id: str = Field(
        pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
    )