from email import policy
from email.parser import BytesParser
from email.headerregistry import Address, HeaderRegistry
import unittest

no_refold_policy = policy.default.clone(refold_source="none")


class TestFoldingBug(unittest.TestCase):
    message_before_q = b"From: =?utf-8?Q?a=2C=20123456789012345678901234567890123456?= <abc@example.com>\r\n"
    message_before_b = b"To: =?UTF-8?B?QWt0aXZhIChrw7ZyZW4gKyB0asOkbnN0bGVkaWdhICsgc3R5cmVsc2Up?= <abc@example.com>\r\n"

    message_parsed_q = BytesParser(policy=policy.default).parsebytes(message_before_q)
    message_parsed_b = BytesParser(policy=policy.default).parsebytes(message_before_b)

    expected_address_q = Address(
        display_name="a, 123456789012345678901234567890123456",
        username="abc",
        domain="example.com",
    )
    expected_address_b = Address(
        display_name="Aktiva (kören + tjänstlediga + styrelse)",
        username="abc",
        domain="example.com",
    )

    def test_q_refolded(self):
        message_refolded = self.message_parsed_q.as_bytes()
        parsed_message = BytesParser(policy=policy.default).parsebytes(message_refolded)

        print(message_refolded)

        self.assertEqual(parsed_message["From"].addresses, (self.expected_address_q,))

    def test_b_refolded(self):
        message_refolded = self.message_parsed_b.as_bytes()
        parsed_message = BytesParser(policy=policy.default).parsebytes(message_refolded)

        print(message_refolded)

        self.assertEqual(parsed_message.defects, [])
        self.assertEqual(parsed_message["To"].addresses, (self.expected_address_b,))

    def test_q_no_refold(self):
        message_not_refolded = self.message_parsed_q.as_bytes(policy=no_refold_policy)
        parsed_message = BytesParser(policy=policy.default).parsebytes(
            message_not_refolded
        )

        self.assertEqual(parsed_message["From"].addresses, (self.expected_address_q,))

    def test_b_no_refold(self):
        message_not_refolded = self.message_parsed_b.as_bytes(policy=no_refold_policy)
        parsed_message = BytesParser(policy=policy.default).parsebytes(
            message_not_refolded
        )

        self.assertEqual(parsed_message["To"].addresses, (self.expected_address_b,))


if __name__ == "__main__":
    unittest.main()