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",
        "skills_wiki_short",
    ]
    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(
        # Pattern: Allows letters, numbers, -, _, spaces, and dots (but no double dots or starting dots or spaces)
        pattern=r"^[a-zA-Z0-9_()-][a-zA-Z0-9\s_()-]*(\.[a-zA-Z0-9\s_-]+)*$",
        min_length=1,
        max_length=MAX_FILE_NAME_LENGTH,
    )


class ReviewRequestBase(BaseModel):
    """Requests from the critique (/review) and evaluation (/evaluate) pages.

    These pages have no consent/profile flow, so they only carry the
    machine-generated user and session identifiers.
    """

    user_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
    )

    @field_validator("user_id", mode="before")
    def strip_user_id(cls, user_id: str):
        """Trim whitespace before the pattern check, so a stray leading/
        trailing space doesn't turn a persona login into a mismatched id."""
        return user_id.strip() if isinstance(user_id, str) else user_id


class ReviewAskRequest(ReviewRequestBase):
    """Ask ChatGPT the selected question."""

    question_id: int = Field(ge=1, le=12709)


class CritiqueHighlight(BaseModel):
    """One highlighted excerpt of the answer, rated good/bad/mixed.

    The rating alone is meaningful, so the comment may be empty.
    """

    highlighted_text: str = Field(min_length=1, max_length=MAX_RESPONSE_LENGTH)
    comment: str = Field(min_length=0, max_length=MAX_COMMENT_LENGTH)
    rating: Literal["good", "bad", "mixed"]

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


class CritiqueSubmitRequest(ReviewRequestBase):
    """A complete critique of one LLM answer, saved as a single record."""

    question_id: int = Field(ge=1, le=12709)
    answer_id: UUID
    answer: str = Field(min_length=1, max_length=MAX_RESPONSE_LENGTH)
    # The reviewer's grade of the LLM answer, out of 5.
    answer_grade: int = Field(ge=1, le=5)
    general_comment: str = Field(min_length=0, max_length=MAX_COMMENT_LENGTH)
    highlights: list[CritiqueHighlight] = Field(default_factory=list, max_length=200)

    @field_validator("answer", "general_comment")
    def sanitize(cls, value: str):
        """Remove HTML tags to prevent XSS"""
        return nh3.clean(value)


class EvaluationSubmitRequest(ReviewRequestBase):
    """A supervisor's evaluation of a submitted critique."""

    critique_id: UUID
    # The supervisor's grade of the critique, out of 5.
    critique_grade: int = Field(ge=1, le=5)
    comment: str = Field(min_length=0, 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 TraineeAnswerRequest(ReviewRequestBase):
    """A trainee's free-text answer to a practice scenario, plus confidence."""

    question_id: int = Field(ge=1, le=12709)
    answer: str = Field(min_length=1, max_length=MAX_RESPONSE_LENGTH)
    # How confident the trainee is in their answer, 1 (low) to 5 (high).
    confidence: int = Field(ge=1, le=5)

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


class TraineeGuidelineAnswerRequest(ReviewRequestBase):
    """A trainee's free-text answer to a guideline question, plus confidence."""

    # Guideline question ids are strings namespaced by document,
    # e.g. "anaphylaxis-3" (see guideline_questions.py).
    question_id: str = Field(min_length=1, max_length=64)
    answer: str = Field(min_length=1, max_length=MAX_RESPONSE_LENGTH)
    # How confident the trainee is in their answer, 1 (low) to 5 (high).
    confidence: int = Field(ge=1, le=5)

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


class TraineeFeedbackReviewRequest(ReviewRequestBase):
    """A trainee's review of the model feedback (highlights + general comment)."""

    # int for the CSV bank, str for guideline questions ("anaphylaxis-3").
    question_id: int | str
    feedback_id: UUID
    # Whether the trainee agrees with the LLM feedback. Required — the client
    # blocks submission until answered, and this enforces it server-side too.
    agree: bool
    general_comment: str = Field(min_length=0, max_length=MAX_COMMENT_LENGTH)
    highlights: list[CritiqueHighlight] = Field(default_factory=list, max_length=200)

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


class TraineeCommentRequest(ReviewRequestBase):
    """A free-form comment left from the trainee footer (about anything)."""

    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 SupervisorReviewRequest(BaseModel):
    """A supervisor's evaluation of one trainee attempt's LLM feedback.

    Not a ReviewRequestBase: there's no trainee session here, just the
    supervisor's own (unauthenticated, demo-only) identity.
    """

    supervisor_id: str = Field(
        pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=MAX_ID_LENGTH
    )
    feedback_id: UUID
    # The supervisor's grade of the LLM feedback, out of 5.
    feedback_grade: int = Field(ge=1, le=5)
    # The supervisor's grade of the trainee's answer itself, out of 5.
    answer_grade: int = Field(ge=1, le=5)
    # Comment on the LLM feedback's quality.
    feedback_comment: str = Field(min_length=0, max_length=MAX_COMMENT_LENGTH)
    # Comment on the trainee's answer itself.
    answer_comment: str = Field(min_length=0, max_length=MAX_COMMENT_LENGTH)

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


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
    )


class ChatMessage(BaseModel):
    role: Literal["user", "assistant", "system"]
    content: str