from tqdm.auto import tqdm as tqdm_auto
from tqdm.utils import _range

class tqdm_telegram(tqdm_auto):
    def __new__(cls, *args, **kwargs):
        # workaround to avoid breaking nested progressbars
        try:
            cls._instances = tqdm_auto._instances
        except AttributeError:
            pass

        instance = super().__new__(cls, *args, **kwargs)

        tqdm_auto._instances = cls._instances
        return instance

    def __init__(self, *args, callback=None, **kwargs):
        """
        Parameters
        ----------
        token  : str, required. Telegram token.
        chat_id  : str, required. Telegram chat ID.
        callback  :  method that will receive the progressbar content on display
        See `tqdm.auto.tqdm.__init__` for other parameters.
        """
        self.callback = callback
        super().__init__(*args, **kwargs)

    def display(self, msg=None, **kwargs):
        super().display(msg=msg, **kwargs)
        fmt = self.format_dict
        if 'bar_format' in fmt and fmt['bar_format']:
            fmt['bar_format'] = fmt['bar_format'].replace('<bar/>', '{bar}')
        if self.callback is not None:
            self.callback(self.format_meter(**fmt))


def ttgrange(*args, **kwargs):
    """
    A shortcut for `tqdm.contrib.telegram.tqdm(xrange(*args), **kwargs)`.
    On Python3+, `range` is used instead of `xrange`.
    """
    return tqdm_telegram(_range(*args), **kwargs)


# Aliases
tqdm = tqdm_telegram
trange = ttgrange

class TelegramIO():
    def __init__(self, token, chat_id, show_last_update=True):
        self.token = token
        self.chat_id = chat_id
        self.session = requests.Session()
        self.text = '_Init_'
        self.message_id = self._sendMessage().json()['result']['message_id']
        self.show_last_update = show_last_update

    def _sendMessage(self):
        try:
            return self.session.post('https://api.telegram.org/bot%s/sendMessage' % self.token,
                                     data=dict(text=self.text, chat_id=self.chat_id, parse_mode='MarkdownV2'))
        except Exception as e:
            print(e)
            return None

    def _editMessageText(self, msg_id, text):
        try:
            return self.session.post('https://api.telegram.org/bot%s/editMessageText' % self.token,
                                     data=dict(text=text, chat_id=self.chat_id, message_id=msg_id, parse_mode='MarkdownV2'))
        except Exception as e:
            print(e)
            return None

    def write(self, s):
        if not s: return
        s = s.strip().replace('\r', '')
        if s == self.text: return  # avoid duplicate message Bot error
        self.text = s

        updated_text = '```\n{}```\n_Update: {}_'.format(self.text, datetime.now().strftime("%H:%M:%S")) if self.show_last_update else ''
        self._editMessageText(self.message_id, updated_text)

tg_bot = TelegramIO(TELEGRAM_BOT_TOKEN, TELEGRAM_CHAT_ID)
iterator = ttgrange(nb_epochs, desc='training', unit='epoch', callback=tg_bot.write)