import imaplib, ssl, re, time, sys, select as sel, email, html

ctx = ssl.create_default_context()
ctx.check_hostname = True
ctx.verify_mode = ssl.CERT_REQUIRED

HOST = "imap.zick-mail.casa"
PORT = 993
TIMEOUT = 120

email_addr = sys.argv[1] if len(sys.argv) > 1 else ""
password = sys.argv[2] if len(sys.argv) > 2 else ""

if not email_addr or not password:
    print("Usage: python script.py email@zick-store.pro password")
    sys.exit(1)

m = imaplib.IMAP4_SSL(HOST, PORT, ssl_context=ctx)
m.login(email_addr, password)
m.select("INBOX")

_, nums = m.search(None, "ALL")
start_count = len(nums[0].split())

def idle_wait(timeout):
    tag = m._new_tag()
    m.send(tag + b" IDLE\r\n")
    while True:
        line = m.readline()
        if line.startswith(b"+"):
            break
    start = time.time()
    found = False
    while time.time() - start < timeout:
        r, _, _ = sel.select([m.sock], [], [], 0.5)
        if r:
            line = m.readline()
            if b"EXISTS" in line:
                found = True
                break
    m.send(b"DONE\r\n")
    while True:
        line = m.readline()
        if line.startswith(tag):
            break
    return found

def strip_html(h):
    h = re.sub(r'<style[^>]*>[\s\S]*?</style>', '', h, flags=re.I)
    h = re.sub(r'<script[^>]*>[\s\S]*?</script>', '', h, flags=re.I)
    h = h.replace('&nbsp;', ' ').replace('&amp;', '&').replace('&lt;', '<').replace('&gt;', '>')
    h = re.sub(r'&#\d+;', ' ', h)
    h = re.sub(r'<br\s*/?>', '\n', h, flags=re.I)
    h = re.sub(r'</p>', '\n', h, flags=re.I)
    h = re.sub(r'</div>', '\n', h, flags=re.I)
    h = re.sub(r'</tr>', '\n', h, flags=re.I)
    h = re.sub(r'<[^>]+>', ' ', h)
    h = html.unescape(h)
    h = re.sub(r'[ \t]+', ' ', h)
    h = re.sub(r'\n{3,}', '\n\n', h)
    return h.strip()

def find_code(text):
    codes = [c for c in re.findall(r'\b(\d{4,8})\b', text) if not re.match(r'^20\d\d$', c)]
    return max(codes, key=len) if codes else None

def extract_code(raw_bytes):
    msg = email.message_from_bytes(raw_bytes)
    plain = ""
    html_part = ""
    if msg.is_multipart():
        for part in msg.walk():
            ct = part.get_content_type()
            if ct == "text/plain" and not plain:
                payload = part.get_payload(decode=True)
                if payload:
                    plain = payload.decode("utf-8", "ignore")
            elif ct == "text/html" and not html_part:
                payload = part.get_payload(decode=True)
                if payload:
                    html_part = payload.decode("utf-8", "ignore")
    else:
        payload = msg.get_payload(decode=True)
        if payload:
            p = payload.decode("utf-8", "ignore")
            if msg.get_content_type() == "text/html":
                html_part = p
            else:
                plain = p
    code = find_code(plain) if plain else None
    if code:
        return code
    if html_part:
        return find_code(strip_html(html_part))
    return None

print("Waiting for email...")
if idle_wait(TIMEOUT):
    _, nums = m.search(None, "ALL")
    ids = nums[0].split()
    if len(ids) > start_count:
        _, d = m.fetch(ids[-1], "(RFC822)")
        raw_bytes = d[0][1]
        code = extract_code(raw_bytes)
        if code:
            print("Code:", code)
        else:
            print("Email received, code not found")
    else:
        print("No new emails")
else:
    print("Timeout - no email received in", TIMEOUT, "sec")

m.logout()