SurfSense/surfsense_backend/app/proprietary/platforms/youtube/schemas.py

# ruff: noqa: N815 - field names intentionally mirror the Apify camelCase spec
"""Apify-compatible input/output models for the YouTube scraper.

The models mirror the public Apify "YouTube Scraper" actor spec so the endpoint
can be a drop-in. The MVP populates a reliably-sourced subset; every other field
is accepted (input) or emitted as ``None``/``[]`` (output) so parity is additive.

``VideoItem`` uses ``extra="allow"`` on purpose: the Apify spec explicitly says
"Other properties may be included", and it lets us grow the output shape without
breaking existing consumers.
"""

from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field

SubtitlesFormat = Literal["srt", "vtt", "xml", "plaintext"]
SortingOrder = Literal["relevance", "rating", "date", "views"]
DateFilter = Literal["hour", "today", "week", "month", "year"]
VideoTypeFilter = Literal["video", "movie"]
LengthFilter = Literal["under4", "between420", "plus20"]
SortVideosBy = Literal["NEWEST", "POPULAR", "OLDEST"]
VideoContentType = Literal["video", "shorts", "stream"]
SortCommentsBy = Literal["TOP_COMMENTS", "NEWEST_FIRST"]
CommentType = Literal["comment", "reply"]


class StartUrl(BaseModel):
    """A single direct URL entry (Apify passes ``{"url": ...}`` objects)."""

    model_config = ConfigDict(extra="allow")

    url: str


class YouTubeScrapeInput(BaseModel):
    """Full Apify input surface.

    ``maxResults*`` semantics follow Apify: ``0`` means "fetch none of that type"
    (see the ``startUrls`` note in the spec), NOT unlimited. A caller wanting an
    unbounded run leaves the value high; the request-time guard lives in the
    route, not here.
    """

    model_config = ConfigDict(extra="allow")

    # Discovery
    searchQueries: list[str] = Field(default_factory=list)
    startUrls: list[StartUrl] = Field(default_factory=list)

    # Per-type caps (independent, applied per search term and per channel)
    maxResults: int = Field(default=0, ge=0, le=999999)
    maxResultsShorts: int = Field(default=0, ge=0, le=999999)
    maxResultStreams: int = Field(default=0, ge=0, le=999999)

    # Subtitles
    downloadSubtitles: bool = False
    subtitlesLanguage: str = "en"
    subtitlesFormat: SubtitlesFormat = "srt"
    saveSubsToKVS: bool = False
    preferAutoGeneratedSubtitles: bool = False

    # Search filters
    sortingOrder: SortingOrder | None = None
    dateFilter: DateFilter | None = None
    videoType: VideoTypeFilter | None = None
    lengthFilter: LengthFilter | None = None
    isHD: bool | None = None
    hasSubtitles: bool | None = None
    hasCC: bool | None = None
    is3D: bool | None = None
    isLive: bool | None = None
    isBought: bool | None = None
    is4K: bool | None = None
    is360: bool | None = None
    hasLocation: bool | None = None
    isHDR: bool | None = None
    isVR180: bool | None = None

    # Channel/date scoping
    oldestPostDate: str | None = None
    sortVideosBy: SortVideosBy | None = None


class DescriptionLink(BaseModel):
    model_config = ConfigDict(extra="allow")

    url: str | None = None
    text: str | None = None


class SubtitleTrack(BaseModel):
    model_config = ConfigDict(extra="allow")

    srtUrl: str | None = None
    type: str | None = None
    language: str | None = None
    srt: str | None = None


class Collaborator(BaseModel):
    model_config = ConfigDict(extra="allow")

    name: str | None = None
    username: str | None = None
    url: str | None = None


class VideoItem(BaseModel):
    """Apify output item. Unsourced fields default to ``None``/``[]``.

    ``extra="allow"`` keeps the contract open (Apify: "Other properties may be
    included") so gradually-added fields never break existing consumers.
    """

    model_config = ConfigDict(extra="allow")

    # Core video
    title: str | None = None
    id: str | None = None
    url: str | None = None
    viewCount: int | None = None
    date: str | None = None
    duration: str | None = None
    type: VideoContentType = "video"
    thumbnailUrl: str | None = None
    input: str | None = None
    fromYTUrl: str | None = None
    order: int | None = None

    # Engagement / content
    likes: int | None = None
    commentsCount: int | None = None
    commentsTurnedOff: bool | None = None
    text: str | None = None
    descriptionLinks: list[DescriptionLink] | None = None
    subtitles: list[SubtitleTrack] | None = None
    hashtags: list[str] = Field(default_factory=list)
    location: str | None = None

    # Translations
    translatedTitle: str | None = None
    translatedText: str | None = None

    # Video flags
    isMonetized: bool | None = None
    isMembersOnly: bool = False
    isPaidContent: bool = False
    isAgeRestricted: bool | None = None
    collaborators: list[Collaborator] | None = None

    # Channel
    channelName: str | None = None
    channelUrl: str | None = None
    channelUsername: str | None = None
    channelId: str | None = None
    numberOfSubscribers: int | None = None
    channelTotalVideos: int | None = None
    channelDescription: str | None = None
    channelLocation: str | None = None
    channelJoinedDate: str | None = None
    channelTotalViews: int | None = None
    isChannelVerified: bool | None = None
    channelBannerUrl: str | None = None
    channelAvatarUrl: str | None = None

    def to_output(self) -> dict[str, Any]:
        """Serialize to the flat dict shape Apify emits (keeps extras)."""
        return self.model_dump(exclude_none=False)


class YouTubeCommentsInput(BaseModel):
    """Apify "YouTube Comments Scraper" input surface.

    ``maxComments`` counts every emitted item (top-level comments and replies).
    Setting ``oldestCommentDate`` forces newest-first sort, matching Apify.
    """

    model_config = ConfigDict(extra="allow")

    startUrls: list[StartUrl] = Field(default_factory=list)
    maxComments: int = Field(default=1, ge=1, le=999999)
    sortCommentsBy: SortCommentsBy = "NEWEST_FIRST"
    oldestCommentDate: str | None = None


class CommentItem(BaseModel):
    """Apify comment output item (one per top-level comment or reply)."""

    model_config = ConfigDict(extra="allow")

    author: str | None = None
    comment: str | None = None
    cid: str | None = None
    replyCount: int | None = None
    replyToCid: str | None = None
    voteCount: int | None = None
    authorIsChannelOwner: bool | None = None
    publishedTimeText: str | None = None
    type: CommentType | None = None
    hasCreatorHeart: bool | None = None
    videoId: str | None = None
    pageUrl: str | None = None
    commentsCount: int | None = None
    title: str | None = None

    def to_output(self) -> dict[str, Any]:
        return self.model_dump(exclude_none=False)