from pyrogram import Client
from pyrogram.errors import (
    ChannelInvalid,
    ChannelPrivate,
    FileIdInvalid,
    FileReferenceExpired,
    FileReferenceInvalid,
    PeerIdInvalid,
)
from pyrogram.types import ChatPhoto

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

STALE_FILE_ID = (FileIdInvalid, FileReferenceExpired, FileReferenceInvalid)
UNKNOWN_OWNER = (PeerIdInvalid, ChannelInvalid, ChannelPrivate)


async def _fresh_file_id(client: Client, owner_id: int, unique_id: str) -> str | None:
    try:
        chat = await client.get_chat(owner_id)
    except UNKNOWN_OWNER:
        return None
    photo = chat.photo if chat else None
    if photo and photo.big_photo_unique_id == unique_id:
        return photo.big_file_id
    async for item in client.get_chat_photos(owner_id):
        if isinstance(item, ChatPhoto) and item.big_photo_unique_id == unique_id:
            return item.big_file_id
    return None


async def _download(
    client: Client, owner_id: int, unique_id: str, file_id: str
) -> bytes | None:
    try:
        return await download_bytes(client, file_id)
    except STALE_FILE_ID:
        fresh = await _fresh_file_id(client, owner_id, unique_id)
        if fresh is None:
            return None
        return await download_bytes(client, fresh)


@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
    data = await _download(client, owner_id, unique_id, file_id)
    if data is None:
        return
    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),
    )