# Copyright (c) Twisted Matrix Laboratories.
# See LICENSE for details.

import email.message
import email.parser
from io import BytesIO, StringIO

from twisted.mail import bounce
from twisted.trial import unittest


class BounceTests(unittest.TestCase):
    """
    Bounce message generation
    """

    def test_bounceMessageUnicode(self):
        """
        L{twisted.mail.bounce.generateBounce} can accept L{unicode}.
        """
        fromAddress, to, s = bounce.generateBounce(
            StringIO(
                """\
From: Moshe Zadka <[email protected]>
To: [email protected]
Subject: test

"""
            ),
            "[email protected]",
            "[email protected]",
        )
        self.assertEqual(fromAddress, b"")
        self.assertEqual(to, b"[email protected]")
        emailParser = email.parser.Parser()
        mess = emailParser.parse(StringIO(s.decode("utf-8")))
        self.assertEqual(mess["To"], "[email protected]")
        self.assertEqual(mess["From"], "[email protected]")
        self.assertEqual(mess["subject"], "Returned Mail: see transcript for details")

    def test_bounceMessageBytes(self):
        """
        L{twisted.mail.bounce.generateBounce} can accept L{bytes}.
        """
        fromAddress, to, s = bounce.generateBounce(
            BytesIO(
                b"""\
From: Moshe Zadka <[email protected]>
To: [email protected]
Subject: test

"""
            ),
            b"[email protected]",
            b"[email protected]",
        )
        self.assertEqual(fromAddress, b"")
        self.assertEqual(to, b"[email protected]")
        emailParser = email.parser.Parser()
        mess = emailParser.parse(StringIO(s.decode("utf-8")))
        self.assertEqual(mess["To"], "[email protected]")
        self.assertEqual(mess["From"], "[email protected]")
        self.assertEqual(mess["subject"], "Returned Mail: see transcript for details")

    def test_bounceMessageCustomTranscript(self):
        """
        Pass a custom transcript message to L{twisted.mail.bounce.generateBounce}.
        """
        fromAddress, to, s = bounce.generateBounce(
            BytesIO(
                b"""\
From: Moshe Zadka <[email protected]>
To: [email protected]
Subject: test

"""
            ),
            b"[email protected]",
            b"[email protected]",
            "Custom transcript",
        )
        self.assertEqual(fromAddress, b"")
        self.assertEqual(to, b"[email protected]")
        emailParser = email.parser.Parser()
        mess = emailParser.parse(StringIO(s.decode("utf-8")))
        self.assertEqual(mess["To"], "[email protected]")
        self.assertEqual(mess["From"], "[email protected]")
        self.assertEqual(mess["subject"], "Returned Mail: see transcript for details")
        self.assertTrue(mess.is_multipart())
        parts = mess.get_payload()
        self.assertEqual(parts[0].get_payload(), "Custom transcript\n")

    def _bounceBigMessage(self, header, message, ioType):
        """
        Pass a really big message to L{twisted.mail.bounce.generateBounce}.
        """
        fromAddress, to, s = bounce.generateBounce(
            ioType(header + message), "[email protected]", "[email protected]"
        )
        emailParser = email.parser.Parser()
        mess = emailParser.parse(StringIO(s.decode("utf-8")))
        self.assertEqual(mess["To"], "[email protected]")
        self.assertEqual(mess["From"], "[email protected]")
        self.assertEqual(mess["subject"], "Returned Mail: see transcript for details")
        self.assertTrue(mess.is_multipart())
        parts = mess.get_payload()
        innerMessage = parts[1].get_payload()
        if isinstance(message, bytes):
            message = message.decode("utf-8")
        self.assertEqual(innerMessage[0].get_payload() + "\n", message)

    def test_bounceBigMessage(self):
        """
        L{twisted.mail.bounce.generateBounce} with big L{unicode} and
        L{bytes} messages.
        """
        header = b"""\
From: Moshe Zadka <[email protected]>
To: [email protected]
Subject: test

"""
        self._bounceBigMessage(header, b"Test test\n" * 10000, BytesIO)
        self._bounceBigMessage(header.decode("utf-8"), "More test\n" * 10000, StringIO)