import base64
import html2text
from enum import Enum

class Message:
    class Format(Enum):
        PLAIN = 0
        HTML = 1

    def __init__(self, fragment):
        self._next_fragment = 0
        self._raw_content = ""

        parts = fragment.split()

        if len(parts) != 2:
            raise Exception("Invalid fragment!")

        if parts[0] == "PLAIN":
            self._msg_fmt = Message.Format.PLAIN
        elif parts[0] == "HTML":
            self._msg_fmt = Message.Format.HTML
        else:
            raise Exception("Unknown message type %s!" % msg_fmt)

        try:
            self._msg_len = int(parts[1])
        except:
            raise Exception("Invalid message length %s!" % parts[1])

        if self._msg_len < 1 or self._msg_len > 1000:
            raise Exception("Invalid message length %d!" % self._msg_len)

    def feed_fragment(self, fragment):
        parts = fragment.split()

        if len(parts) != 2:
            raise Exception("Invalid fragment!")

        try:
            fragment_number = int(parts[0])
        except:
            raise Exception("Invalid fragment number %s!" % parts[0])

        if fragment_number == self._next_fragment and fragment_number < self._msg_len:
            self._next_fragment += 1
        else:
            raise Exception("Unexpected fragment!")

        self._raw_content += parts[1]

    def is_complete(self):
        return self._next_fragment == self._msg_len

    def get_content(self):
        try:
            msgB64Bytes = self._raw_content.encode("UTF-8")
            msgBytes = base64.b64decode(msgB64Bytes)
            msgStr = msgBytes.decode("UTF-8")
        except:
            raise Exception("Message decoding error!")

        if self._msg_fmt == Message.Format.PLAIN:
            return {
                "msgtype": "m.text",
                "body": msgStr,
            }
        elif self._msg_fmt == Message.Format.HTML:
            return {
                "format": "org.matrix.custom.html",
                "msgtype": "m.text",
                "body": html2text.html2text(msgStr),
                "formatted_body": msgStr
            }