from io import BytesIO

from pyrogram.types import Story

from userbot import PyroClient
from userbot.modules.stories import repository


def _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


@PyroClient.on_story()
async def on_story(client: PyroClient, story: Story) -> None:
    ctx = client.capture
    if ctx is None:
        return
    media_kind = story.media.name.lower() if story.media else None
    storage_key: str | None = None
    file_size: int | None = None
    downloaded = False
    if not story.deleted and story.media is not None:
        buffer = await client.download_media(story, in_memory=True)
        if isinstance(buffer, BytesIO):
            data = buffer.getvalue()
            storage_key = ctx.storage.put(data)
            file_size = len(data)
            downloaded = True
    await repository.upsert_story(
        ctx.pool,
        ctx.account_id,
        _peer_id(story),
        story.id,
        story.date,
        story.expire_date,
        story.caption,
        media_kind,
        storage_key,
        file_size,
        story.views,
        str(story.raw),
        pinned=bool(story.pinned),
        deleted=bool(story.deleted),
        downloaded=downloaded,
    )


handlers = on_story.handlers