# mautrix-telegram - A Matrix-Telegram puppeting bridge
# Copyright (C) 2021 Tulir Asokan
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
from __future__ import annotations

from typing import Type, Any, Union, TYPE_CHECKING
from abc import ABC, abstractmethod
import platform
import asyncio
import logging
import time

from telethon.sessions import Session
from telethon.network import (ConnectionTcpMTProxyRandomizedIntermediate, ConnectionTcpFull,
                              Connection)
from telethon.tl.patched import MessageService, Message
from telethon.tl.types import (
    Channel, Chat, MessageActionChannelMigrateFrom, PeerUser, TypeUpdate, UpdatePinnedMessages,
    UpdatePinnedChannelMessages, UpdateChatParticipantAdmin, UpdateChatParticipants, PeerChat,
    UpdateChatUserTyping, UpdateDeleteChannelMessages, UpdateNewMessage, UpdateDeleteMessages,
    UpdateEditChannelMessage, UpdateEditMessage, UpdateNewChannelMessage, UpdateReadHistoryOutbox,
    UpdateShortChatMessage, UpdateShortMessage, UpdateUserName, UpdateUserPhoto, UpdateUserStatus,
    UpdateUserTyping, User, UserStatusOffline, UserStatusOnline, UpdateReadHistoryInbox,
    UpdateReadChannelInbox, MessageEmpty, UpdateFolderPeers, UpdatePinnedDialogs,
    UpdateNotifySettings, UpdateChannelUserTyping)

from mautrix.types import UserID, PresenceState
from mautrix.errors import MatrixError
from mautrix.appservice import AppService
from mautrix.util.logging import TraceLogger
from mautrix.util.opt_prometheus import Histogram, Counter

from . import portal as po, puppet as pu, __version__
from .db import Message as DBMessage, PgSession
from .types import TelegramID
from .tgclient import MautrixTelegramClient
from .config import Config

if TYPE_CHECKING:
    from .bot import Bot
    from .__main__ import TelegramBridge

UpdateMessage = Union[UpdateShortChatMessage, UpdateShortMessage, UpdateNewChannelMessage,
                      UpdateNewMessage, UpdateEditMessage, UpdateEditChannelMessage]
UpdateMessageContent = Union[
    UpdateShortMessage, UpdateShortChatMessage, Message, MessageService, MessageEmpty
]

UPDATE_TIME = Histogram(
    name="bridge_telegram_update",
    documentation="Time spent processing Telegram updates",
    labelnames=("update_type",),
)
UPDATE_ERRORS = Counter(
    name="bridge_telegram_update_error",
    documentation="Number of fatal errors while handling Telegram updates",
    labelnames=("update_type",),
)


class AbstractUser(ABC):
    loop: asyncio.AbstractEventLoop = None
    log: TraceLogger
    az: AppService
    bridge: 'TelegramBridge'
    config: Config
    relaybot: 'Bot'
    ignore_incoming_bot_events: bool = True
    max_deletions: int = 10

    client: MautrixTelegramClient | None
    mxid: UserID | None

    tgid: TelegramID | None
    username: str | None
    is_bot: bool

    is_relaybot: bool

    puppet_whitelisted: bool
    whitelisted: bool
    relaybot_whitelisted: bool
    matrix_puppet_whitelisted: bool
    is_admin: bool

    def __init__(self) -> None:
        self.is_admin = False
        self.matrix_puppet_whitelisted = False
        self.puppet_whitelisted = False
        self.whitelisted = False
        self.relaybot_whitelisted = False
        self.client = None
        self.is_relaybot = False
        self.is_bot = False

    @property
    def connected(self) -> bool:
        return self.client and self.client.is_connected()

    @property
    def _proxy_settings(self) -> tuple[Type[Connection], tuple[Any, ...] | None]:
        proxy_type = self.config["telegram.proxy.type"].lower()
        connection = ConnectionTcpFull
        connection_data = (self.config["telegram.proxy.address"],
                           self.config["telegram.proxy.port"],
                           self.config["telegram.proxy.rdns"],
                           self.config["telegram.proxy.username"],
                           self.config["telegram.proxy.password"])
        if proxy_type == "disabled":
            connection_data = None
        elif proxy_type == "socks4":
            connection_data = (1,) + connection_data
        elif proxy_type == "socks5":
            connection_data = (2,) + connection_data
        elif proxy_type == "http":
            connection_data = (3,) + connection_data
        elif proxy_type == "mtproxy":
            connection = ConnectionTcpMTProxyRandomizedIntermediate
            connection_data = (connection_data[0], connection_data[1], connection_data[4])

        return connection, connection_data

    @classmethod
    def init_cls(cls, bridge: 'TelegramBridge') -> None:
        cls.bridge = bridge
        cls.config = bridge.config
        cls.loop = bridge.loop
        cls.az = bridge.az
        cls.ignore_incoming_bot_events = cls.config["bridge.relaybot.ignore_own_incoming_events"]
        cls.max_deletions = cls.config["bridge.max_telegram_delete"]

    async def _init_client(self) -> None:
        self.log.debug(f"Initializing client for {self.name}")

        session = await PgSession.get(self.name)
        if self.config["telegram.server.enabled"]:
            session.set_dc(self.config["telegram.server.dc"],
                           self.config["telegram.server.ip"],
                           self.config["telegram.server.port"])

        if self.is_relaybot:
            base_logger = logging.getLogger("telethon.relaybot")
        else:
            base_logger = logging.getLogger(f"telethon.{self.tgid or -hash(self.mxid)}")

        device = self.config["telegram.device_info.device_model"]
        sysversion = self.config["telegram.device_info.system_version"]
        appversion = self.config["telegram.device_info.app_version"]
        connection, proxy = self._proxy_settings

        assert isinstance(session, Session)

        self.client = MautrixTelegramClient(
            session=session,

            api_id=self.config["telegram.api_id"],
            api_hash=self.config["telegram.api_hash"],

            app_version=__version__ if appversion == "auto" else appversion,
            system_version=(MautrixTelegramClient.__version__
                            if sysversion == "auto" else sysversion),
            device_model=(f"{platform.system()} {platform.release()}"
                          if device == "auto" else device),

            timeout=self.config["telegram.connection.timeout"],
            connection_retries=self.config["telegram.connection.retries"],
            retry_delay=self.config["telegram.connection.retry_delay"],
            flood_sleep_threshold=self.config["telegram.connection.flood_sleep_threshold"],
            request_retries=self.config["telegram.connection.request_retries"],
            connection=connection,
            proxy=proxy,
            raise_last_call_error=True,

            loop=self.loop,
            base_logger=base_logger
        )
        self.client.add_event_handler(self._update_catch)

    @abstractmethod
    async def update(self, update: TypeUpdate) -> bool:
        return False

    @abstractmethod
    async def post_login(self) -> None:
        raise NotImplementedError()

    @abstractmethod
    async def register_portal(self, portal: po.Portal) -> None:
        raise NotImplementedError()

    @abstractmethod
    async def unregister_portal(self, tgid: int, tg_receiver: int) -> None:
        raise NotImplementedError()

    async def _update_catch(self, update: TypeUpdate) -> None:
        start_time = time.time()
        update_type = type(update).__name__
        try:
            if not await self.update(update):
                await self._update(update)
        except Exception:
            self.log.exception("Failed to handle Telegram update")
            UPDATE_ERRORS.labels(update_type=update_type).inc()
        UPDATE_TIME.labels(update_type=update_type).observe(time.time() - start_time)

    @property
    @abstractmethod
    def name(self) -> str:
        raise NotImplementedError()

    async def is_logged_in(self) -> bool:
        return (self.client and self.client.is_connected()
                and await self.client.is_user_authorized())

    async def has_full_access(self, allow_bot: bool = False) -> bool:
        return (self.puppet_whitelisted
                and (not self.is_bot or allow_bot)
                and await self.is_logged_in())

    async def start(self, delete_unless_authenticated: bool = False) -> AbstractUser:
        if not self.client:
            await self._init_client()
        await self.client.connect()
        self.log.debug(f"{'Bot' if self.is_relaybot else self.mxid} connected: {self.connected}")
        return self

    async def ensure_started(self, even_if_no_session=False) -> AbstractUser:
        if self.connected:
            return self
        if even_if_no_session or await PgSession.has(self.mxid):
            self.log.debug("Starting client due to ensure_started"
                           f"(even_if_no_session={even_if_no_session})")
            await self.start(delete_unless_authenticated=not even_if_no_session)
        return self

    async def stop(self) -> None:
        await self.client.disconnect()
        self.client = None

    # region Telegram update handling

    async def _update(self, update: TypeUpdate) -> None:
        asyncio.create_task(self._handle_entity_updates(getattr(update, "_entities", {})))
        if isinstance(update, (UpdateShortChatMessage, UpdateShortMessage, UpdateNewChannelMessage,
                               UpdateNewMessage, UpdateEditMessage, UpdateEditChannelMessage)):
            await self.update_message(update)
        elif isinstance(update, UpdateDeleteMessages):
            await self.delete_message(update)
        elif isinstance(update, UpdateDeleteChannelMessages):
            await self.delete_channel_message(update)
        elif isinstance(update, (UpdateChatUserTyping, UpdateChannelUserTyping, UpdateUserTyping)):
            await self.update_typing(update)
        elif isinstance(update, UpdateUserStatus):
            await self.update_status(update)
        elif isinstance(update, UpdateChatParticipantAdmin):
            await self.update_admin(update)
        elif isinstance(update, UpdateChatParticipants):
            await self.update_participants(update)
        elif isinstance(update, (UpdatePinnedMessages, UpdatePinnedChannelMessages)):
            await self.update_pinned_messages(update)
        elif isinstance(update, (UpdateUserName, UpdateUserPhoto)):
            await self.update_others_info(update)
        elif isinstance(update, UpdateReadHistoryOutbox):
            await self.update_read_receipt(update)
        elif isinstance(update, (UpdateReadHistoryInbox, UpdateReadChannelInbox)):
            await self.update_own_read_receipt(update)
        elif isinstance(update, UpdateFolderPeers):
            await self.update_folder_peers(update)
        elif isinstance(update, UpdatePinnedDialogs):
            await self.update_pinned_dialogs(update)
        elif isinstance(update, UpdateNotifySettings):
            await self.update_notify_settings(update)
        else:
            self.log.trace("Unhandled update: %s", update)

    async def update_folder_peers(self, update: UpdateFolderPeers) -> None:
        pass

    async def update_pinned_dialogs(self, update: UpdatePinnedDialogs) -> None:
        pass

    async def update_notify_settings(self, update: UpdateNotifySettings) -> None:
        pass

    async def update_pinned_messages(
        self, update: UpdatePinnedMessages | UpdatePinnedChannelMessages
    ) -> None:
        if isinstance(update, UpdatePinnedMessages):
            portal = await po.Portal.get_by_entity(update.peer, tg_receiver=self.tgid)
        else:
            portal = await po.Portal.get_by_tgid(TelegramID(update.channel_id))
        if portal and portal.mxid:
            await portal.receive_telegram_pin_ids(update.messages, self.tgid,
                                                  remove=not update.pinned)

    @staticmethod
    async def update_participants(update: UpdateChatParticipants) -> None:
        portal = await po.Portal.get_by_tgid(TelegramID(update.participants.chat_id))
        if portal and portal.mxid:
            await portal.update_power_levels(update.participants.participants)

    async def update_read_receipt(self, update: UpdateReadHistoryOutbox) -> None:
        if not isinstance(update.peer, PeerUser):
            self.log.debug("Unexpected read receipt peer: %s", update.peer)
            return

        portal = await po.Portal.get_by_tgid(
            TelegramID(update.peer.user_id), tg_receiver=self.tgid
        )
        if not portal or not portal.mxid:
            return

        # We check that these are user read receipts, so tg_space is always the user ID.
        message = await DBMessage.get_one_by_tgid(TelegramID(update.max_id), self.tgid,
                                                  edit_index=-1)
        if not message:
            return

        puppet = await pu.Puppet.get_by_tgid(TelegramID(update.peer.user_id))
        await puppet.intent.mark_read(portal.mxid, message.mxid)

    async def update_own_read_receipt(
        self, update: UpdateReadHistoryInbox | UpdateReadChannelInbox
    ) -> None:
        puppet = await pu.Puppet.get_by_tgid(self.tgid)
        if not puppet.is_real_user:
            return

        if isinstance(update, UpdateReadChannelInbox):
            portal = await po.Portal.get_by_tgid(TelegramID(update.channel_id))
        elif isinstance(update.peer, PeerChat):
            portal = await po.Portal.get_by_tgid(TelegramID(update.peer.chat_id))
        elif isinstance(update.peer, PeerUser):
            portal = await po.Portal.get_by_tgid(
                TelegramID(update.peer.user_id), tg_receiver=self.tgid
            )
        else:
            self.log.debug("Unexpected own read receipt peer: %s", update.peer)
            return

        if not portal or not portal.mxid:
            return

        tg_space = portal.tgid if portal.peer_type == "channel" else self.tgid
        message = await DBMessage.get_one_by_tgid(TelegramID(update.max_id), tg_space,
                                                  edit_index=-1)
        if not message:
            return

        await puppet.intent.mark_read(portal.mxid, message.mxid)

    async def update_admin(self, update: UpdateChatParticipantAdmin) -> None:
        # TODO duplication not checked
        portal = await po.Portal.get_by_tgid(TelegramID(update.chat_id))
        if not portal or not portal.mxid:
            return

        await portal.set_telegram_admin(TelegramID(update.user_id))

    async def update_typing(
        self, update: UpdateUserTyping | UpdateChatUserTyping | UpdateChannelUserTyping
    ) -> None:
        sender = None
        if isinstance(update, UpdateUserTyping):
            portal = await po.Portal.get_by_tgid(
                TelegramID(update.user_id), tg_receiver=self.tgid, peer_type="user"
            )
            sender = await pu.Puppet.get_by_tgid(TelegramID(update.user_id))
        elif isinstance(update, UpdateChannelUserTyping):
            portal = await po.Portal.get_by_tgid(TelegramID(update.channel_id))
        elif isinstance(update, UpdateChatUserTyping):
            portal = await po.Portal.get_by_tgid(TelegramID(update.chat_id))
        else:
            return

        if isinstance(update, (UpdateChannelUserTyping, UpdateChatUserTyping)):
            # Can typing notifications come from non-user peers?
            if not update.from_id.user_id:
                return
            sender = await pu.Puppet.get_by_tgid(TelegramID(update.from_id.user_id))

        if not sender or not portal or not portal.mxid:
            return

        await portal.handle_telegram_typing(sender, update)

    async def _handle_entity_updates(self, entities: dict[int, User | Chat | Channel]) -> None:
        try:
            users = (entity for entity in entities.values() if isinstance(entity, User))
            puppets = ((await pu.Puppet.get_by_tgid(TelegramID(user.id)), user) for user in users)
            await asyncio.gather(*[puppet.try_update_info(self, info)
                                   async for puppet, info in puppets if puppet])
        except Exception:
            self.log.exception("Failed to handle entity updates")

    async def update_others_info(self, update: UpdateUserName | UpdateUserPhoto) -> None:
        # TODO duplication not checked
        puppet = await pu.Puppet.get_by_tgid(TelegramID(update.user_id))
        if isinstance(update, UpdateUserName):
            puppet.username = update.username
            if await puppet.update_displayname(self, update):
                await puppet.save()
        elif isinstance(update, UpdateUserPhoto):
            if await puppet.update_avatar(self, update.photo):
                await puppet.save()
        else:
            self.log.warning(f"Unexpected other user info update: {type(update)}")

    async def update_status(self, update: UpdateUserStatus) -> None:
        puppet = await pu.Puppet.get_by_tgid(TelegramID(update.user_id))
        if isinstance(update.status, UserStatusOnline):
            await puppet.default_mxid_intent.set_presence(PresenceState.ONLINE)
        elif isinstance(update.status, UserStatusOffline):
            await puppet.default_mxid_intent.set_presence(PresenceState.OFFLINE)
        else:
            self.log.warning(f"Unexpected user status update: type({update})")
        return

    async def get_message_details(
        self, update: UpdateMessage
    ) -> tuple[UpdateMessageContent, pu.Puppet | None, po.Portal | None]:
        if isinstance(update, UpdateShortChatMessage):
            portal = po.Portal.get_by_tgid(TelegramID(update.chat_id))
            if not portal:
                self.log.warning(f"Received message in chat with unknown type {update.chat_id}")
            sender = await pu.Puppet.get_by_tgid(TelegramID(update.from_id))
        elif isinstance(update, UpdateShortMessage):
            portal = await po.Portal.get_by_tgid(
                TelegramID(update.user_id), tg_receiver=self.tgid, peer_type="user"
            )
            sender = await pu.Puppet.get_by_tgid(self.tgid if update.out else update.user_id)
        elif isinstance(update, (UpdateNewMessage, UpdateNewChannelMessage,
                                 UpdateEditMessage, UpdateEditChannelMessage)):
            update = update.message
            if isinstance(update, MessageEmpty):
                return update, None, None
            portal = await po.Portal.get_by_entity(update.peer_id, tg_receiver=self.tgid)
            if update.out:
                sender = await pu.Puppet.get_by_tgid(self.tgid)
            elif isinstance(update.from_id, PeerUser):
                sender = await pu.Puppet.get_by_tgid(TelegramID(update.from_id.user_id))
            else:
                sender = None
        else:
            self.log.warning("Unexpected message type in User#get_message_details: "
                             f"{type(update)}")
            return update, None, None
        return update, sender, portal

    @staticmethod
    async def _try_redact(message: DBMessage) -> None:
        portal = await po.Portal.get_by_mxid(message.mx_room)
        if not portal:
            return
        try:
            await portal.main_intent.redact(message.mx_room, message.mxid)
        except MatrixError:
            pass

    async def delete_message(self, update: UpdateDeleteMessages) -> None:
        if len(update.messages) > self.max_deletions:
            return

        for message_id in update.messages:
            for message in await DBMessage.get_all_by_tgid(TelegramID(message_id), self.tgid):
                if message.redacted:
                    continue
                await message.delete()
                number_left = await DBMessage.count_spaces_by_mxid(message.mxid, message.mx_room)
                if number_left == 0:
                    await self._try_redact(message)

    async def delete_channel_message(self, update: UpdateDeleteChannelMessages) -> None:
        if len(update.messages) > self.max_deletions:
            return

        channel_id = TelegramID(update.channel_id)

        for message_id in update.messages:
            for message in await DBMessage.get_all_by_tgid(TelegramID(message_id), channel_id):
                if message.redacted:
                    continue
                await message.delete()
                await self._try_redact(message)

    async def update_message(self, original_update: UpdateMessage) -> None:
        update, sender, portal = await self.get_message_details(original_update)
        if not portal:
            return
        elif portal and not portal.allow_bridging:
            self.log.debug(f"Ignoring message in portal {portal.tgid_log} (bridging disallowed)")
            return

        if self.is_relaybot:
            if update.is_private:
                if not self.config["bridge.relaybot.private_chat.invite"]:
                    self.log.debug(f"Ignoring private message to bot from {sender.id}")
                    return
            elif not portal.mxid and self.config["bridge.relaybot.ignore_unbridged_group_chat"]:
                self.log.debug("Ignoring message received by bot"
                               f" in unbridged chat {portal.tgid_log}")
                return

        if ((self.ignore_incoming_bot_events and self.relaybot
             and sender and sender.id == self.relaybot.tgid)):
            self.log.debug(f"Ignoring relaybot-sent message %s to %s", update.id, portal.tgid_log)
            return

        await portal.backfill_lock.wait(f"update {update.id}")

        if isinstance(update, MessageService):
            if isinstance(update.action, MessageActionChannelMigrateFrom):
                self.log.trace(f"Received %s in %s by %d, unregistering portal...",
                               update.action, portal.tgid_log, sender.id)
                await self.unregister_portal(update.action.chat_id, update.action.chat_id)
                await self.register_portal(portal)
                return
            self.log.trace("Handling action %s to %s by %d", update.action, portal.tgid_log,
                           (sender.id if sender else 0))
            return await portal.handle_telegram_action(self, sender, update)

        if isinstance(original_update, (UpdateEditMessage, UpdateEditChannelMessage)):
            return await portal.handle_telegram_edit(self, sender, update)
        return await portal.handle_telegram_message(self, sender, update)

    # endregion