from datetime import UTC, datetime

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.groups.repository import insert_chat_history
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.raw_json import raw_json

MEMBER_CAP = 200

_MISSING_SENDERS = """
SELECT DISTINCT sender_id FROM messages
WHERE account_id = $1 AND chat_id = $2 AND sender_id > 0
  AND NOT EXISTS (
    SELECT 1 FROM peers p WHERE p.account_id = $1 AND p.peer_id = messages.sender_id
  )
LIMIT 100
"""


async def _save_user(ctx: CaptureContext, user: User) -> None:
    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))
    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
        )


async def _enrich_chat_meta(client: Client, ctx: CaptureContext, chat_id: int) -> None:
    chat = await client.get_chat(chat_id)
    photo = chat.photo
    photo_unique_id = photo.big_photo_unique_id if photo else None
    photo_file_id = photo.big_file_id if photo else None
    await insert_chat_history(
        ctx.pool,
        ctx.account_id,
        chat_id,
        0,
        "meta",
        chat.title,
        photo_unique_id,
        None,
        datetime.now(UTC),
        raw_json(chat),
    )
    if photo_file_id and photo_unique_id:
        await note_avatar(
            ctx.pool, ctx.account_id, chat_id, "chat", photo_unique_id, photo_file_id
        )


async def _enrich_members(client: Client, ctx: CaptureContext, chat_id: int) -> None:
    try:
        async for member in client.get_chat_members(chat_id, limit=MEMBER_CAP):
            if isinstance(member.user, User):
                await _save_user(ctx, member.user)
    except (BadRequest, Forbidden):
        return


async def _fetch_users(client: Client, ids: list[int]) -> list[User]:
    try:
        users = await client.get_users(ids)
    except BadRequest:
        users = []
        for sender_id in ids:
            try:
                users.append(await client.get_users(sender_id))
            except BadRequest:
                continue
    if isinstance(users, User):
        return [users]
    return [user for user in users if isinstance(user, User)]


async def _enrich_senders(client: Client, ctx: CaptureContext, chat_id: int) -> None:
    rows = await ctx.pool.fetch(_MISSING_SENDERS, ctx.account_id, chat_id)
    ids = [row["sender_id"] for row in rows]
    if not ids:
        return
    for user in await _fetch_users(client, ids):
        await _save_user(ctx, user)


@register("enrich_chat")
async def enrich_chat(ctx: JobContext) -> None:
    client = ctx.client
    if client is None:
        return
    capture = getattr(client, "capture", None)
    if capture is None:
        return
    chat_id = ctx.job.params["chat_id"]
    await _enrich_chat_meta(client, capture, chat_id)
    await _enrich_members(client, capture, chat_id)
    await _enrich_senders(client, capture, chat_id)