from collections.abc import AsyncIterator, Callable

from pyrogram import Client
from pyrogram.errors import FloodPremiumWait, FloodWait, RPCError
from pyrogram.types import Story

from userbot.modules.capture.context import CaptureContext
from userbot.modules.jobs.context import JobContext
from userbot.modules.jobs.registry import register
from userbot.modules.stories.service import save_story

SAVE_EVERY = 10

StorySource = Callable[[], AsyncIterator[Story]]


def _sources(client: Client, peer_id: int, *, own: bool) -> dict[str, StorySource]:
    sources: dict[str, StorySource] = {
        "active": lambda: client.get_chat_stories(peer_id),
        "pinned": lambda: client.get_pinned_stories(peer_id),
    }
    if own:
        sources["archived"] = lambda: client.get_archived_stories(peer_id)
    return sources


async def _drain(
    ctx: JobContext, capture: CaptureContext, name: str, source: StorySource
) -> int:
    client = ctx.client
    if client is None:
        return 0
    saved = 0
    async for story in source():
        try:
            await save_story(client, capture, story)
        except (FloodWait, FloodPremiumWait):
            raise
        except RPCError:
            continue
        saved += 1
        if saved % SAVE_EVERY == 0:
            await ctx.report_progress({"saved": saved, "source": name})
            if await ctx.is_canceled():
                break
    return saved


@register("backfill_stories")
async def backfill_stories(ctx: JobContext) -> None:
    client = ctx.client
    if client is None:
        return
    capture = getattr(client, "capture", None)
    if capture is None:
        return
    peer_id = ctx.job.params["peer_id"]
    own = client.me is not None and client.me.id == peer_id
    saved = 0
    errors: dict[str, str] = {}
    for name, source in _sources(client, peer_id, own=own).items():
        try:
            saved += await _drain(ctx, capture, name, source)
        except (FloodWait, FloodPremiumWait):
            raise
        except RPCError as exc:
            errors[name] = type(exc).__name__
        if await ctx.is_canceled():
            break
    await ctx.report_progress({"saved": saved, "done": True, "errors": errors})