import smtplib
import imaplib
import imap_tools
import time
import ssl
import argparse
import uuid

parser = argparse.ArgumentParser()
parser.add_argument('host', type = str)
parser.add_argument('--mailfrom', default = 'foo@example.com')
parser.add_argument('--rcptto', default = 'awesome@vacu.store')
parser.add_argument('--submission', default = False, action='store_true')
parser.add_argument('--smtp-starttls', default = None, action='store_true')
parser.add_argument('--imap-insecure', default = False, action = 'store_true')
parser.add_argument('--imap-dir', default = None)
parser.add_argument('--username')
parser.add_argument('--password')
parser.add_argument('--expect-recipient-refused',
    dest = 'expect',
    action = 'store_const',
    const = 'recipient_refused',
    default = 'received'
)
parser.add_argument('--expect-sent', dest = 'expect', action = 'store_const', const = 'sent')
parser.add_argument('--expect-imap-error', dest = 'expect', action = 'store_const', const = 'imap_error')

args = parser.parse_args()

print(f"got args {args!r}")

# smtp_starttls = args.smtp_starttls
# if smtp_starttls is None:
#     smtp_starttls = args.submission

username = args.username
password = args.password
if password is None:
    password = username

if (username is None or password is None) and (args.submission or args.expect == 'received'):
    assert False, "Bad args"

msg_magic = str(uuid.uuid4())

def mk_ctx():
    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    return ctx

try:
    smtp = None
    if args.submission:
        smtp = smtplib.SMTP_SSL(args.host, port=465, context = mk_ctx())
    else:
        smtp = smtplib.SMTP(args.host, port=25)
    smtp.ehlo()
    if args.smtp_starttls:
        smtp.starttls(context = mk_ctx())
        smtp.ehlo()
    if args.submission:
        smtp.login(username, password)
    smtp.sendmail(args.mailfrom, args.rcptto, 'Subject: Test\n\n' + msg_magic)
    smtp.close()
except smtplib.SMTPRecipientsRefused:
    assert args.expect == 'recipient_refused', "Server rejected message as recipients refused"
else:
    assert (not args.expect == 'recipient_refused'), "Server was supposed to reject this message, but it didn't"

if args.expect == 'received' or args.expect == 'imap_error':
    time.sleep(3)
    try:
        any_okay = False
        with imap_tools.MailBox(args.host, ssl_context = mk_ctx()).login(username, password) as mailbox:
            if args.imap_dir is None:
                directories = [d.name for d in mailbox.folder.list()]
            else:
                directories = [args.imap_dir]
            for imap_dir in directories:
                print(f"checking in {imap_dir!r}")
                for d in mailbox.folder.list(imap_dir):
                    if d.name != imap_dir:
                        directories.append(d.name)
                try:
                    mailbox.folder.set(imap_dir)
                except imap_tools.errors.MailboxFolderSelectError as e:
                    continue
                for msg in mailbox.fetch():
                    print(f"found message with text {msg.text!r}")
                    any_okay = any_okay or msg.text.strip() == msg_magic
        assert any_okay, "Could not find the message in the mailbox"
    except imaplib.IMAP4.error as e:
        assert args.expect == 'imap_error', f"IMAP error: {e}"
    else:
        assert not args.expect == 'imap_error', "Expected an IMAP error, but didn't get one"