import asyncio
import contextlib
from io import BytesIO

from pyrogram import Client
from pyrogram.types import (
    Animation,
    Audio,
    Document,
    Message,
    Photo,
    Sticker,
    Story,
    Video,
    VideoNote,
    Voice,
)

from utils.logging import logger

type Downloadable = (
    str
    | Message
    | Story
    | Audio
    | Document
    | Photo
    | Sticker
    | Animation
    | Video
    | Voice
    | VideoNote
)

STOP_TIMEOUT = 5.0

_reset_lock = asyncio.Lock()


async def _reset_media_sessions(client: Client) -> None:
    async with _reset_lock:
        sessions = list(client.media_sessions.values())
        client.media_sessions.clear()
        for session in sessions:
            with contextlib.suppress(Exception):
                await asyncio.wait_for(session.stop(), STOP_TIMEOUT)
        logger.warning(f"[yellow]Dropped {len(sessions)} stale media session(s).[/]")


async def download_bytes(client: Client, target: Downloadable) -> bytes | None:
    try:
        buffer = await client.download_media(target, in_memory=True)
    except (TimeoutError, OSError):
        await _reset_media_sessions(client)
        buffer = await client.download_media(target, in_memory=True)
    return buffer.getvalue() if isinstance(buffer, BytesIO) else None