chronicle.nixg.ru/telegram-archiver/app/models.py

"""
Pydantic models for Telegram Archiver.
Used for data validation and serialization.
"""

from pydantic import BaseModel, Field
from datetime import datetime
from typing import Optional, Literal
from enum import Enum


class MediaType(str, Enum):
    """Types of media files supported by the archiver."""
    PHOTO = "photo"
    VIDEO = "video"
    DOCUMENT = "document"
    AUDIO = "audio"
    VOICE = "voice"
    VIDEO_NOTE = "video_note"
    STICKER = "sticker"
    ANIMATION = "animation"  # GIF
    TIMEOUT = "timeout"  # File skipped due to timeout


class MediaFile(BaseModel):
    """Represents a media file attached to a Telegram message."""
    filename: str = Field(..., description="Original filename")
    type: MediaType = Field(..., description="Type of media")
    caption: Optional[str] = Field(None, description="Caption text for the media")
    size: Optional[int] = Field(None, description="File size in bytes")
    is_too_large: bool = Field(
        default=False,
        description="True if file exceeds MAX_FILE_SIZE limit"
    )
    download_url: Optional[str] = Field(
        None,
        description="Telegram URL for files that weren't downloaded"
    )

    model_config = {"use_enum_values": True}


class PostData(BaseModel):
    """
    Represents a single Telegram post with all metadata.
    This is the main data structure passed to the markdown generator.
    """
    message_id: int = Field(..., description="Telegram message ID")
    date: datetime = Field(..., description="Message timestamp")
    text: str = Field(default="", description="Message text in Markdown format")
    author: str = Field(..., description="Channel name or author")
    channel_username: str = Field(..., description="Channel username (for output path)")

    # Relationships
    reply_to: Optional[int] = Field(
        None,
        description="Message ID this post is replying to"
    )
    repost_from: Optional[int] = Field(
        None,
        description="Original message ID if this is a forwarded post"
    )
    repost_channel: Optional[str] = Field(
        None,
        description="Original channel name if this is a forwarded post"
    )

    # Media attachments
    media_files: list[MediaFile] = Field(
        default_factory=list,
        description="List of media files attached to the post"
    )

    # Metadata
    views: Optional[int] = Field(None, description="View count (if available)")
    has_large_files: bool = Field(
        default=False,
        description="True if some files were not downloaded due to size limit"
    )
    tags: list[str] = Field(
        default_factory=list,
        description="Hashtags extracted from the post"
    )

    @property
    def bundle_dir(self) -> str:
        """Return the directory name for this post's bundle."""
        return str(self.message_id)

    @property
    def index_path(self) -> str:
        """Return relative path to index.md within bundle."""
        return f"{self.message_id}/index.md"

    def has_media(self) -> bool:
        """Check if post has any media files."""
        return len(self.media_files) > 0

    def has_downloaded_media(self) -> bool:
        """Check if post has any downloaded (not too large) media files."""
        return any(not f.is_too_large for f in self.media_files)


class ArchiveRequest(BaseModel):
    """Request model for FastAPI endpoint."""
    channel: str = Field(..., description="Telegram channel username or ID")
    output_dir: Optional[str] = Field(
        None,
        description="Override default output directory"
    )
    limit: Optional[int] = Field(
        None,
        description="Limit number of posts to archive (for testing)"
    )
    from_message_id: Optional[int] = Field(
        None,
        description="Start archiving from this message ID"
    )
    force: bool = Field(
        default=False,
        description="Force re-download of already archived posts"
    )


class ArchiveResponse(BaseModel):
    """Response model for FastAPI endpoint."""
    channel: str
    channel_title: str
    posts_archived: int
    posts_skipped: int
    media_downloaded: int
    media_skipped_large: int
    output_path: str
    log_file: str
    large_files_report: Optional[str] = None
    duration_seconds: float


class HealthCheck(BaseModel):
    """Health check response."""
    status: str
    version: str
    session_active: bool