#!/bin/env python
"""
Retrieve number and subjects of unread mail from Gmail.

External dependencies:
    Set a GMAIL_LOGIN environment variable as a json string with Gmail access credentials.
    An example in Zsh:
        export GMAIL_LOGIN='{"username":"[email protected]", "password":"app_password"}'

Usage example:
    In Zsh:
        $ chmod +x checkmail.py
        $ ./checkmail.py

        Or drop the .py and copy the script to $HOME/.local/bin, then:
        $ checkmail
"""

import os
import json
import imaplib
from typing import Iterable
from socket import gaierror
from typing import Generator
from email import policy, message_from_bytes

try:
    from prettytable import PrettyTable # pyright:ignore[reportMissingImports,reportUnknownVariableType]
except ImportError:
    print("Please install the prettytable package.")
    exit(-1)


def check_unread_email(*,
    username: str, password: str, gmail_handle: imaplib.IMAP4_SSL
) -> Generator[dict[str,str], None, None]:
    """Check if there's unread mail in Gmail

    Args:
        username: Gmail email address.
        password: Google app password https://support.google.com/accounts/answer/185833
        gmail_handle: An IMAP SSL object connected to Gmail.

    Returns:
        A generator of string(s).

    Side effect:
        Prints a message if there's no unread mail.

    Raises:
        None
    """

    try:
        gmail_handle.login(username, password)
        gmail_handle.select("INBOX")

        _, response = gmail_handle.search(None, "(UNSEEN)")

        if unread := response[0].split():
            for email_number in reversed(unread):
                status, email_data = gmail_handle.fetch(
                    email_number, "(BODY.PEEK[HEADER.FIELDS (FROM SUBJECT)])"
                )
                if status == "OK":
                    headers = message_from_bytes(
                        email_data[0][1], policy=policy.default  # pyright:ignore[reportOptionalSubscript,reportArgumentType]
                    )
                    yield dict(headers)
        else:
            print("Inbox zero!")
            exit(0)

    except imaplib.IMAP4.error as e:
        print(e)
        exit(-1)

    finally:
        gmail_handle.close()
        gmail_handle.logout()


def presentation(*, unread_mail: Iterable[dict[str,str]]):
    try:
        table = PrettyTable() # pyright:ignore[reportUnknownVariableType]
        table.field_names = ("#", "Sender", "Subject")

        for num, headers in enumerate(unread_mail,start=1):
            _from = headers["From"].split("<")[0].strip()
            _from = _from if len(_from) <=21 else _from[:21]+"..."

            _subject = headers["Subject"]
            _subject = _subject if len(_subject) <=55 else _subject[:55]+"..."
            
            table.add_row((num,_from,_subject)) # pyright:ignore[reportUnknownMemberType]
        
        table.align = "l"
        print(table) # pyright:ignore[reportUnknownArgumentType]

    except (gaierror, ConnectionResetError):
        print("Could not connect to Gmail")



if __name__ == "__main__":
    presentation(
        unread_mail=check_unread_email(
            gmail_handle=imaplib.IMAP4_SSL("imap.gmail.com"),
            **json.loads(os.getenv("GMAIL_LOGIN")),  # pyright:ignore[reportArgumentType]
        )
    )