"""Write archived messages to disk in mbox or Maildir format."""

import mailbox
import time
from pathlib import Path


class MboxArchive:
    """
    Append messages to a standard mbox file.

    The mbox is opened once and kept open for the duration of the run for
    efficiency. Call .close() (or use as a context manager) when done.
    """

    def __init__(self, path: Path):
        self.path = path
        path.parent.mkdir(parents=True, exist_ok=True)
        self._mbox = mailbox.mbox(str(path))
        self._mbox.lock()

    def add(self, raw_bytes: bytes) -> None:
        msg = mailbox.mboxMessage(raw_bytes)
        self._mbox.add(msg)

    def flush(self) -> None:
        self._mbox.flush()

    def close(self) -> None:
        self._mbox.unlock()
        self._mbox.close()

    def __enter__(self):
        return self

    def __exit__(self, *_):
        self.close()


class MaildirArchive:
    """
    Write messages into a Maildir structure (cur/new/tmp).

    Each message is written as an individual file, which makes it trivially
    importable into Thunderbird, Dovecot, etc.
    """

    def __init__(self, path: Path):
        self.path = path
        path.parent.mkdir(parents=True, exist_ok=True)
        self._mdir = mailbox.Maildir(str(path), create=True)
        self._count = 0

    def add(self, raw_bytes: bytes) -> None:
        msg = mailbox.MaildirMessage(raw_bytes)
        msg.set_subdir("cur")
        self._mdir.add(msg)
        self._count += 1

    def flush(self) -> None:
        self._mdir.flush()

    def close(self) -> None:
        self._mdir.close()

    def __enter__(self):
        return self

    def __exit__(self, *_):
        self.close()


def open_archive(fmt: str, path: Path):
    """Factory: return the right archive object for *fmt* ('mbox' or 'maildir')."""
    if fmt == "mbox":
        return MboxArchive(path)
    elif fmt == "maildir":
        return MaildirArchive(path)
    else:
        raise ValueError(f"Unknown archive format: {fmt!r}. Use 'mbox' or 'maildir'.")
