#!/usr/bin/env python3
# Copyright (c) 2018 softhyena.com
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import typing

#Can this be made into a stack based encoder? --Félix
def encode(input: typing.Any) -> bytes:
    """Encode a object to a bencode byte string"""
    t = type(input)
    if t == str:
        input = input.encode()
        return "{}:".format(len(input)).encode()+input

    elif t == bytes:
        return "{}:".format(len(input)).encode()+input

    elif t == int:
        return "i{}e".format(input).encode()

    elif t == list:
        output = b"l"
        for item in input:
            output = output + encode(item)

        return output+b"e"

    elif t == dict:
        output = b"d"
        for k,v in input.items():
            output = output + encode(k) + encode(v)

        return output+b"e"

    else:
        raise TypeError("Bencode cannot encode type {}".format(str(t)))

def decode(input: bytes, tryDecodeStrings: bool = False, disableStringKeys: bool = False) -> typing.Any:
    """Decode a bencode byte string"""
    stack = []

    offset = 0
    length = len(input)
    while True:
        try:
            data = None

            #Parse strings
            if "0" <= chr(input[offset]) <= "9":
                tmp = chr(input[offset])
                while True:
                    offset = offset + 1
                    if "0" <= chr(input[offset]) <= "9":
                        tmp = tmp + chr(input[offset])
                    else:
                        break

                if chr(input[offset]) != ":":
                    raise ValueError("Unexpected {} in bencode string decoding at {}".format(chr(input[offset]), offset))

                offset = offset + 1
                data = input[offset:offset+int(tmp)]
                offset = offset + int(tmp)

                if tryDecodeStrings:
                    #Try to decode strings
                    try:
                        data = data.decode()
                    except UnicodeDecodeError:
                        pass


            #Parse integers
            elif chr(input[offset]) == "i":
                offset = offset + 1

                if "0" <= chr(input[offset]) <= "9":
                    tmp = ""
                    while True:
                        if chr(input[offset]) == "e":
                            offset = offset + 1
                            break
                        elif "0" <= chr(input[offset]) <= "9":
                            tmp = tmp + chr(input[offset])
                            offset = offset + 1
                        else:
                            raise ValueError("Unexpected {} in bencode integer decoding at {}".format(chr(input[offset]), offset))

                else:
                    raise ValueError("Unexpected {} in bencode integer decoding at {}".format(chr(input[offset]), offset))

                data = int(tmp)

            #Setup lists
            elif chr(input[offset]) == "l":
                stack.append([])
                offset = offset + 1

            #Setup dictionary
            elif chr(input[offset]) == "d":
                stack.append({})
                stack.append(()) #Push a KV marker
                offset = offset + 1

            #Process stack lists or dictionaries
            elif chr(input[offset]) == "e":
                data = stack.pop()
                #KV should have been appended
                if type(data) == tuple:
                    if len(data):
                        raise ValueError("Stack unbalanced! Unexpected end of dict at {}!".format(offset))
                    else:
                        data = stack.pop()

                offset = offset + 1

            #Push data into lists or dictionaries
            if data:
                if len(stack) > 0:
                    destType = type(stack[-1])
                    if destType == list:
                        stack[-1].append(data)
                    elif destType == tuple:
                        stack[-1] = tuple(list(stack[-1]) + [data])
                        l = len(stack[-1])
                        if l == 2:
                            kv = stack.pop()

                            if type(kv[0]) == bytes and not disableStringKeys:
                                try:
                                    stack[-1][kv[0].decode()] = kv[1]
                                except UnicodeDecodeError:
                                    stack[-1][kv[0]] = kv[1]
                            else:
                                stack[-1][kv[0]] = kv[1]

                            stack.append(()) #Push the a new KV marker

                    else:
                        raise ValueError("Stack unbalanced! Unexpected {} at {}!".format(type(stack[-1]), offset))
                else:
                    return data

        except IndexError as e:
            raise ValueError("Unexpected end of string at {}!".format(offset))

#Aliases
loads = decode
dumps = encode

def load(file, *args, **kwargs):
    return decode(file.read(), *args, **kwargs)

def dump(input, file):
    file.write(encode(file.read(), *args, **kwargs))