from pyrogram import Client
from pyrogram.errors import BadRequest, Forbidden
from pyrogram.types import User

from userbot.modules.capture.context import CaptureContext
from userbot.modules.jobs.context import JobContext
from userbot.modules.jobs.registry import register
from userbot.modules.profiles.snapshots import save_group, save_private, save_user
from utils.logging import logger

SAVE_EVERY = 100
USERS_BATCH = 200

_UPSERT_DIALOG = """
INSERT INTO dialogs (account_id, chat_id) VALUES ($1, $2)
ON CONFLICT (account_id, chat_id) DO UPDATE SET updated_at = now()
"""


async def _enrich_users(client: Client, ctx: CaptureContext, ids: list[int]) -> None:
    for start in range(0, len(ids), USERS_BATCH):
        batch = ids[start : start + USERS_BATCH]
        try:
            result = await client.get_users(batch)
        except (BadRequest, Forbidden):
            continue
        users = result if isinstance(result, list) else [result]
        for user in users:
            if isinstance(user, User):
                await save_user(ctx, user)


@register("sync_dialogs")
async def sync_dialogs(ctx: JobContext) -> None:
    client = ctx.client
    if client is None:
        return
    capture = getattr(client, "capture", None)
    if capture is None:
        return
    processed = ctx.job.progress.get("processed", 0)
    nameless: list[int] = []
    async for dialog in client.get_dialogs():
        chat = dialog.chat
        if chat is None or chat.id is None:
            continue
        chat_id = chat.id
        try:
            if chat_id > 0:
                if not await save_private(capture, chat):
                    nameless.append(chat_id)
            else:
                await save_group(capture, chat)
        except (BadRequest, Forbidden):
            pass
        except Exception:
            logger.exception(f"Failed to snapshot dialog {chat_id}")
        await ctx.pool.execute(_UPSERT_DIALOG, ctx.account_id, chat_id)
        processed += 1
        if processed % SAVE_EVERY == 0:
            await ctx.report_progress({"processed": processed})
    await _enrich_users(client, capture, nameless)
    await ctx.report_progress({"processed": processed, "done": True})