from pyrogram import Client

from userbot.modules.avatars import repository
from userbot.modules.capture.context import CaptureContext
from userbot.modules.download import download_bytes
from utils.logging import logger


async def capture_avatar(  # noqa: PLR0913
    client: Client,
    ctx: CaptureContext,
    owner_id: int,
    owner_kind: str,
    file_id: str,
    unique_id: str,
    raw: str,
) -> None:
    if await repository.avatar_exists(ctx.pool, ctx.account_id, owner_id, unique_id):
        return
    storage_key: str | None = None
    file_size: int | None = None
    downloaded = False
    try:
        data = await download_bytes(client, file_id)
    except TimeoutError:
        logger.warning(f"[yellow]Avatar download timed out for {owner_id}.[/]")
        data = None
    if data is not None:
        storage_key = ctx.storage.put(data)
        file_size = len(data)
        downloaded = True
    await repository.insert_avatar(
        ctx.pool,
        ctx.account_id,
        owner_id,
        owner_kind,
        unique_id,
        storage_key,
        file_size,
        None,
        raw,
        downloaded=downloaded,
    )