from __future__ import annotations

import asyncio
import email
import os
import re
from datetime import datetime, timezone
from email import policy
from email.message import Message
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from threading import Thread

from aiosmtpd.controller import Controller


OUTPUT_DIR = Path(os.environ.get("MAILDUMP_OUTPUT_DIR", "/out"))
SMTP_PORT = int(os.environ.get("MAILDUMP_SMTP_PORT", "1025"))
HTTP_PORT = int(os.environ.get("MAILDUMP_HTTP_PORT", "8025"))


def safe_name(value: str) -> str:
    cleaned = re.sub(r"[^A-Za-z0-9._-]+", "-", value.strip())
    return cleaned.strip("-.") or "message"


def extract_part(message: Message, wanted: str) -> str:
    if message.is_multipart():
        for part in message.walk():
            if part.get_content_type() == wanted:
                payload = part.get_payload(decode=True) or b""
                charset = part.get_content_charset() or "utf-8"
                return payload.decode(charset, errors="replace")
        return ""

    if message.get_content_type() != wanted:
        return ""

    payload = message.get_payload(decode=True) or b""
    charset = message.get_content_charset() or "utf-8"
    return payload.decode(charset, errors="replace")


class MailDumpHandler:
    async def handle_DATA(self, server, session, envelope):
        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

        parsed = email.message_from_bytes(envelope.content, policy=policy.default)
        timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ")
        subject = safe_name(parsed.get("subject", "no-subject"))
        message_root = OUTPUT_DIR / f"{timestamp}-{subject}"
        message_root.mkdir(parents=True, exist_ok=True)

        (message_root / "message.eml").write_bytes(envelope.content)
        (message_root / "metadata.txt").write_text(
            "\n".join(
                [
                    f"mail_from: {envelope.mail_from}",
                    f"rcpt_tos: {', '.join(envelope.rcpt_tos)}",
                    f"subject: {parsed.get('subject', '')}",
                    f"date: {parsed.get('date', '')}",
                ]
            )
            + "\n",
            encoding="utf-8",
        )

        text_body = extract_part(parsed, "text/plain")
        html_body = extract_part(parsed, "text/html")

        if text_body:
            (message_root / "body.txt").write_text(text_body, encoding="utf-8")
        if html_body:
            (message_root / "body.html").write_text(html_body, encoding="utf-8")

        return "250 OK"


def serve_http() -> None:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    os.chdir(OUTPUT_DIR)
    server = ThreadingHTTPServer(("0.0.0.0", HTTP_PORT), SimpleHTTPRequestHandler)
    server.serve_forever()


async def main() -> None:
    http_thread = Thread(target=serve_http, daemon=True)
    http_thread.start()

    controller = Controller(MailDumpHandler(), hostname="0.0.0.0", port=SMTP_PORT)
    controller.start()

    try:
        while True:
            await asyncio.sleep(3600)
    finally:
        controller.stop()


if __name__ == "__main__":
    asyncio.run(main())