# Copyright 2023 <omicron.me@protonmail.com>
# Copyright 2023 <andreasruden91@gmail.com>
#
# This file is part of d2warehouse.
#
# d2warehouse is free software: you can redistribute it and/or modify it under the
# terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# d2warehouse is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
# A PARTICULAR PURPOSE. See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with
# Mercator. If not, see <https://www.gnu.org/licenses/>.

from bitarray import bitarray, decodetree
import itertools

code = {
    "a": bitarray("11110"),
    "b": bitarray("0101"),
    "c": bitarray("01000"),
    "d": bitarray("110001"),
    "e": bitarray("110000"),
    "f": bitarray("010011"),
    "g": bitarray("11010"),
    "h": bitarray("00011"),
    "i": bitarray("1111110"),
    "j": bitarray("000101110"),
    "k": bitarray("010010"),
    "l": bitarray("11101"),
    "m": bitarray("01101"),
    "n": bitarray("001101"),
    "o": bitarray("1111111"),
    "p": bitarray("11001"),
    "q": bitarray("11011001"),
    "r": bitarray("11100"),
    "s": bitarray("0010"),
    "t": bitarray("01100"),
    "u": bitarray("00001"),
    "v": bitarray("1101110"),
    "w": bitarray("00000"),
    "x": bitarray("00111"),
    "y": bitarray("0001010"),
    "z": bitarray("11011000"),
    "0": bitarray("11111011"),
    "1": bitarray("1111100"),
    "2": bitarray("001100"),
    "3": bitarray("1101101"),
    "4": bitarray("11111010"),
    "5": bitarray("00010110"),
    "6": bitarray("1101111"),
    "7": bitarray("01111"),
    "8": bitarray("000100"),
    "9": bitarray("01110"),
    " ": bitarray("10"),
}

decode_tree = decodetree(code)


def decode(bits: bitarray, n) -> tuple[str, int]:
    text = "".join(itertools.islice(bits.iterdecode(decode_tree), n))
    length = len(encode(text))
    return text, length


def encode(s: str) -> bitarray:
    bits = bitarray(endian="little")
    bits.encode(code, s)
    bits.encode(code, " ")
    return bits