#!/usr/bin/python2.4

""" Add a label to every message over a certain size """

import sys
import imaplib
import socket
import readline
import getpass

#imaplib.Debug = 1


def usage(argv0):
  print "Usage: %s [-i]" % argv0
  print "   -i  Work on the INBOX instead of the All Mail label"
  print "   -h  Help -> this message"


def main(argv):
  doAllMail = 1
  min_size = 5
  max_size = 25
  iter = 2

  for arg in argv[1:]:
    if arg == "-i":
      doAllMail = 0
    elif arg == "-h":
      usage(argv[0])
      return

  print "IMAP Size labeler.\n"

  email_address = raw_input('Email address: ')
  print "Password will not be echoed to the terminal"
  password = getpass.getpass('Password: ')

  remote_ipaddr = socket.gethostbyname('imap.gmail.com')
  print "IP Address for imap.gmail.com: %s" % remote_ipaddr

  print "Connecting to imap.gmail.com"
  imapconn = imaplib.IMAP4_SSL('imap.gmail.com', 993)

  print "Banner: " + imapconn.welcome

  print "Connected, attempting login"
  print imapconn.login(email_address, password)

  if doAllMail:
    print "Attempting to select All Mail"
    print imapconn.select("[Gmail]/All Mail")
  else:
    print "Attempting to select INBOX"
    print imapconn.select("INBOX")

  seen_uids = set()
  current_size = max_size
  mb = 1024*1024
  while current_size >= min_size:
    print "Checking for messages >%dMB" % (current_size)
    r = imapconn.uid("SEARCH", "larger", str(current_size * mb))
    if r[0] != 'OK':
      print r
      break
    matching_uids = set(r[1][0].split())
    working_uids = matching_uids.difference(seen_uids)
    seen_uids.update(matching_uids)

    if len(working_uids) > 0:
      label = 'Size_%dMB' % (current_size)
      print "Add label %s to %d messages" % (label, len(working_uids))
      imapconn.create(label)
      imapconn.uid('COPY', ','.join(working_uids), label)

    current_size -= iter

  print "Success, logging out"
  print imapconn.logout()


if __name__ == '__main__':
  main(sys.argv)