import email
import imaplib
import select
import sys

imaplib.Debug = 5


class MailStore:
    def __init__(self, hostname, username, password, callback):
        self.connection = imaplib.IMAP4_SSL(hostname)
        self.username = username
        self.password = password
        self.callback = callback
        self.running = False

    def _fetch(self):
        try:
            typ, data = self.connection.search(None, b'UNSEEN')
        except Exception as e:
            raise e

        if typ == 'OK':
            mail_ids = data[0]

            id_list = mail_ids.split()
            for i in id_list:
                typ, data = self.connection.fetch(i, b'(RFC822)')

                for response_part in data:
                    if isinstance(response_part, tuple):
                        yield email.message_from_bytes(response_part[1])
        else:
            raise ValueError('Server did accept command: typ=' + typ)

    def _idle(self):
        # noinspection PyProtectedMember
        tag = self.connection._new_tag()
        self.connection.select(b'inbox')
        self.connection.send(tag + b' IDLE\r\n')
        response = self.connection.readline()
        if response == b'+ idling\r\n':
            r, w, e = select.select([self.connection.socket().fileno()],
                                    [],
                                    [],
                                    60)
            if r:
                rc, data = self.connection.response(b'FETCH')
                if not rc:
                    raise ValueError(
                        'Unexpected reply when idling: rc={}, data={}'.format(
                            rc,
                            data),
                        *sys.exc_info())
                self._reply_done()
            else:
                self._reply_done()

        else:
            raise ValueError("IDLE not handled? : %s" % response)

    def _run(self):
        while self.running:
            try:
                self.connection.login(self.username, self.password)
                self.connection.select(b'inbox')

                while self.running:
                    for msg in self._fetch():
                        self.callback(msg)
                    self._idle()
            finally:
                self.connection.close()

    def _reply_done(self):
        self.connection.send(b'DONE\r\n')
        rc, data = self.connection.response(b'OK')
        if rc != b'OK':
            raise ValueError('Error while cancelling idle', *sys.exc_info())

    def start(self):
        self.running = True
        self._run()

    def stop(self):
        self.running = False

def main():
    m = MailStore('imap.gmail.com',
                  'test@test.com',
                  'hidemeimapassword',
                  lambda v: print(v))
    m.start()

if __name__ == '__main__':
    try:
        main()
    except (KeyboardInterrupt, SystemExit) as e:
        print('Shutting down application')