import asyncpg
from pyrogram import Client

from userbot.modules.capture.identity import ChatMetaCache, PeerIdentityCache
from userbot.modules.capture.pending import PendingDeletes
from userbot.modules.contacts import ContactCache
from userbot.modules.folders import FolderCache
from userbot.modules.watches import WatchCache
from utils.policy.models import CaptureToggles, ChatKind, ChatMeta, PolicySet
from utils.policy.repository import load_policy_set
from utils.policy.resolver import resolve
from utils.storage import ContentAddressedStorage


class CaptureContext:
    def __init__(
        self,
        account_id: int,
        pool: asyncpg.Pool,
        storage: ContentAddressedStorage,
        folders: FolderCache,
        contacts: ContactCache,
    ) -> None:
        self.account_id = account_id
        self.pool = pool
        self.storage = storage
        self.folders = folders
        self.contacts = contacts
        self.watches = WatchCache(pool, account_id)
        self.peer_identity = PeerIdentityCache()
        self.chat_meta = ChatMetaCache()
        self.pending_deletes = PendingDeletes()
        self.policies = PolicySet()

    async def reload_policies(self) -> None:
        self.policies = await load_policy_set(self.pool, self.account_id)

    def resolve(self, chat: ChatMeta) -> CaptureToggles:
        return resolve(chat, self.folders.folders, self.policies)

    def resolve_peer(self, peer_id: int, *, is_bot: bool = False) -> CaptureToggles:
        return self.resolve(
            ChatMeta(
                chat_id=peer_id,
                kind=ChatKind.DM,
                is_bot=is_bot,
                is_contact=peer_id in self.contacts.ids,
            )
        )


async def build_capture_context(
    client: Client,
    pool: asyncpg.Pool,
    storage: ContentAddressedStorage,
    account_id: int,
) -> CaptureContext:
    folders = FolderCache(client, pool, account_id)
    await folders.refresh()
    contacts = ContactCache(client)
    await contacts.refresh()
    ctx = CaptureContext(account_id, pool, storage, folders, contacts)
    await ctx.reload_policies()
    await ctx.watches.refresh()
    return ctx