from pyrogram import Client
from pyrogram.types import Story

from userbot.modules.capture.context import CaptureContext
from userbot.modules.download import download_bytes
from userbot.modules.raw_json import raw_json
from userbot.modules.stories import repository


def story_peer_id(story: Story) -> int:
    if story.chat is not None:
        return story.chat.id or 0
    if story.from_user is not None:
        return story.from_user.id or 0
    return 0


async def save_story(client: Client, capture: CaptureContext, story: Story) -> None:
    peer_id = story_peer_id(story)
    storage_key: str | None = None
    file_size: int | None = None
    downloaded = False
    stored = await repository.is_downloaded(
        capture.pool, capture.account_id, peer_id, story.id
    )
    if not (stored or story.deleted or story.media is None):
        data = await download_bytes(client, story)
        if data is not None:
            storage_key = capture.storage.put(data)
            file_size = len(data)
            downloaded = True
    await repository.upsert_story(
        capture.pool,
        capture.account_id,
        peer_id,
        story.id,
        story.date,
        story.expire_date,
        story.caption,
        story.media.name.lower() if story.media else None,
        storage_key,
        file_size,
        story.views,
        raw_json(story.raw),
        pinned=bool(story.pinned),
        deleted=bool(story.deleted),
        downloaded=downloaded,
    )