from pyrogram import Client
from pyrogram.errors import FloodPremiumWait, FloodWait, PeerIdInvalid
from pyrogram.types import Message

from userbot.modules.capture import capture_message
from userbot.modules.capture import repository as capture_repo
from userbot.modules.capture.chat_meta import meta_from_chat
from userbot.modules.capture.context import CaptureContext
from userbot.modules.jobs.context import JobContext
from userbot.modules.jobs.registry import register
from userbot.modules.stt import repository as stt_repo
from userbot.modules.stt import should_transcribe_on_backfill
from userbot.modules.stt.gate import safe_transcribe
from utils.logging import logger
from utils.policy.models import CaptureToggles

SAVE_EVERY = 100


async def resolve_min_id(ctx: JobContext, chat_id: int) -> int:
    cursor = ctx.job.cursor or {}
    if "min_id" in cursor:
        return int(cursor["min_id"])
    if ctx.job.params.get("full"):
        return 0
    newest = await capture_repo.max_message_id(ctx.pool, ctx.account_id, chat_id)
    return newest + 1 if newest else 0


async def maybe_transcribe(
    client: Client,
    capture: CaptureContext,
    chat_id: int,
    message: Message,
    self_id: int | None,
) -> None:
    if not (should_transcribe_on_backfill(message, self_id) and message.chat):
        return
    meta = meta_from_chat(message.chat, capture.contacts.ids)
    already = await stt_repo.is_transcribed(
        capture.pool, capture.account_id, chat_id, message.id
    )
    if capture.resolve(meta).stt and not already:
        await safe_transcribe(client, capture, chat_id, message.id)


@register("backfill")
async def backfill(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"]
    toggles = CaptureToggles(
        messages=True,
        media=bool(ctx.job.params.get("media")),
        self_destruct_media=False,
    )
    max_id = int((ctx.job.cursor or {}).get("max_id", 0))
    min_id = await resolve_min_id(ctx, chat_id)
    await ctx.save_cursor({"max_id": max_id, "min_id": min_id})
    processed = ctx.job.progress.get("processed", 0)
    errors = ctx.job.progress.get("errors", 0)
    self_id = client.me.id if client.me else None
    try:
        async for message in client.get_chat_history(
            chat_id, max_id=max_id, min_id=min_id
        ):
            try:
                await capture_message(client, message, capture, toggles)
                await maybe_transcribe(client, capture, chat_id, message, self_id)
            except (FloodWait, FloodPremiumWait, PeerIdInvalid):
                raise
            except Exception:
                errors += 1
                logger.exception(f"Backfill capture failed for {chat_id}/{message.id}")
            processed += 1
            if processed % SAVE_EVERY == 0:
                next_max = message.id - 1
                await ctx.save_cursor({"max_id": next_max, "min_id": min_id})
                await ctx.report_progress(
                    {"processed": processed, "errors": errors, "max_id": next_max}
                )
                if await ctx.is_canceled():
                    return
    except PeerIdInvalid:
        await ctx.report_progress({"processed": processed, "error": "peer_id_invalid"})
        return
    await ctx.report_progress({"processed": processed, "errors": errors, "done": True})