Messenger_1Million_Stats/million/model/message.py

from datetime import datetime
from typing import Any, List
from uuid import uuid4
from pydantic import BaseModel, Field, PrivateAttr, computed_field, field_validator


class Reaction(BaseModel):
    reaction: str
    actor: str


class AudioFile(BaseModel):
    uri: str
    creation_timestamp: int


class Video(BaseModel):
    uri: str
    creation_timestamp: int


class Photo(BaseModel):
    uri: str
    creation_timestamp: int


class Gif(BaseModel):
    uri: str


class Share(BaseModel):
    link: str
    share_text: str


class Sticker(BaseModel):
    uri: str
    ai_stickers: List[Any]


class Message(BaseModel):
    sender_name: str
    date_time: datetime = Field(alias="timestamp_ms")
    content: str | None = None
    sticker: Sticker | None = None
    share: Share | None = None
    photos: List[Photo] | None = None
    videos: List[Video] | None = None
    gifs: List[Gif] | None = None
    audio_files: List[AudioFile] | None = None
    call_duration: int | None = None
    reactions: List[Reaction] | None = None
    is_unsent: bool | None = None
    is_geoblocked_for_viewer: bool

    _id: str = PrivateAttr(default_factory=lambda: str(uuid4()))

    def __str__(self) -> str:
        dt_str = self.date_time.strftime("%d/%m/%Y, %H:%M:%S")

        msg_str = f"{self.sender_name}({dt_str})"

        if self.content:
            msg_str += " : " + self.content
        if self.photos:
            msg_str += f" [PHOTOS {len(self.photos)}]"
        if self.videos:
            msg_str += f" [VIDEOS {len(self.videos)}]"
        if self.gifs:
            msg_str += f" [GIFS {len(self.gifs)}]"

        return msg_str

    def __hash__(self) -> int:
        return hash(self.item_id)

    @computed_field
    @property
    def item_id(self) -> str:
        return self._id

    @field_validator("date_time")
    def parse_timestamp(cls, v):
        if isinstance(v, int):
            return datetime.fromtimestamp(v / 1000)
        return v