from io import BytesIO

from userbot.modules.avatars.repository import get_avatar_file, mark_avatar_downloaded
from userbot.modules.jobs.context import JobContext
from userbot.modules.jobs.registry import register


@register("fetch_avatar")
async def fetch_avatar(ctx: JobContext) -> None:
    client = ctx.client
    if client is None:
        return
    capture = getattr(client, "capture", None)
    if capture is None:
        return
    owner_kind = ctx.job.params["owner_kind"]
    owner_id = ctx.job.params["owner_id"]
    unique_id = ctx.job.params["unique_id"]
    found = await get_avatar_file(
        ctx.pool, ctx.account_id, owner_kind, owner_id, unique_id
    )
    if found is None:
        return
    file_id, downloaded = found
    if downloaded or file_id is None:
        return
    buffer = await client.download_media(file_id, in_memory=True)
    if not isinstance(buffer, BytesIO):
        return
    data = buffer.getvalue()
    storage_key = capture.storage.put(data)
    await mark_avatar_downloaded(
        ctx.pool,
        ctx.account_id,
        owner_kind,
        owner_id,
        unique_id,
        storage_key,
        len(data),
    )