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

from userbot.modules.avatars import note_avatar
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.parse import snapshot_from_high_level
from userbot.modules.profiles.repository import write_profile
from userbot.modules.profiles.snapshots import save_group, save_private
from userbot.modules.raw_json import raw_json
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 not isinstance(user, User):
                continue
            fields, photo_file_id, photo_unique_id = snapshot_from_high_level(user)
            await write_profile(
                ctx.pool,
                ctx.account_id,
                user.id,
                fields,
                raw_json(user),
                history=ctx.resolve_peer(
                    user.id, is_bot=bool(user.is_bot)
                ).profile_history,
            )
            if photo_file_id and photo_unique_id:
                await note_avatar(
                    ctx.pool,
                    ctx.account_id,
                    user.id,
                    "peer",
                    photo_unique_id,
                    photo_file_id,
                )


@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})