#!/usr/bin/env python3

import mailbox
import os
import sys
from email import message, parser, policy
from io import BytesIO
from pathlib import Path

import pikepdf

MONITORED_SENDER = 'vyplatnice@disponero.cz'


def main(inbox: Path, payroll_save_dir: Path, password: str) -> None:
	print('Scanning for new payroll attachments')
	attachments = get_payroll_attachments(inbox)
	for attachment in attachments:
		saved_path = save_pdf(attachment, password, payroll_save_dir)
		if saved_path is not None:
			print(f'Saved {saved_path}')
	print('Done')


def get_payroll_attachments(inbox: Path) -> list[message.EmailMessage]:
	mbox = mailbox.mbox(inbox, create=False)
	mailbox_mails = (msg for msg in mbox.itervalues() if MONITORED_SENDER in msg['From'])
	mail_parser = parser.BytesParser(policy=policy.default)
	email_mails = (mail_parser.parsebytes(mail.as_bytes()) for mail in mailbox_mails)
	attachments = []
	for mail in email_mails:
		for attachment in mail.iter_attachments():
			content_type = attachment.get_content_type()
			if content_type not in ('application/pdf', 'application/octet-stream'):
				continue
			attachments.append(attachment)
	return attachments


def save_pdf(locked_pdf: message.EmailMessage, password: str, dest: Path) -> Path | None:
	dest_filename: Path = dest / locked_pdf.get_filename()
	if dest_filename.exists():
		return None
	pdf_bytes = BytesIO(locked_pdf.get_content())
	with pikepdf.open(pdf_bytes, password=password) as pdf:
		pdf.save(dest_filename)
	return dest_filename


if __name__ == '__main__':
	main(Path(sys.argv[1]), Path(sys.argv[2]), os.getenv('PAYROLL_PASSWORD'))