import windows
import windows.generated_def as gdef

import struct

try:
    unichr # Py2/Py3 compat
except NameError:
    unichr = chr

# http://pubs.opengroup.org/onlinepubs/9629399/chap14.htm#tagcjh_19_03_07

## Array
# A conformant array is an array in which the maximum number of elements is not known beforehand and therefore is included in the representation of the array.
# A varying array is an array in which the actual number of elements passed in a given call varies and therefore is included in the representation of the array.

## Pointers

# NDR defines two classes of pointers that differ both in semantics and in representation
# - reference pointers, which cannot be null and cannot be aliases
# - full pointers, which can be null and can be an aliases
# - unique pointers, which can be null and cannot be aliases, and are transmitted as full pointers.


def pack_dword(x):
     return struct.pack("<I", x)


def dword_pad(s):
    if (len(s) % 4) == 0:
        return s
    return s + (b"P" * (4 - len(s) % 4))


class NdrUniquePTR(object):
    """Create a UNIQUE PTR around a given Ndr type"""
    def __init__(self, subcls):
        self.subcls = subcls

    def pack(self, data):
        subpack = self.subcls.pack(data)
        if subpack is None:
            return pack_dword(0)
        return pack_dword(0x02020202) + subpack

    def unpack(self, stream):
        ptr = NdrLong.unpack(stream)
        if not ptr:
            return None
        return self.subcls.unpack(stream)

    def pack_in_struct(self, data, id):
        if data is None:
            return pack_dword(0), None
        subpack = self.subcls.pack(data)
        if subpack is None:
            return pack_dword(0), None
        return pack_dword(0x01010101 * (id + 1)), subpack

    def unpack_in_struct(self, stream):
        ptr = NdrLong.unpack(stream)
        if not ptr:
            return 0, NdrUnpackNone
        return ptr, self.subcls

    def parse(self, stream):
        data = stream.partial_unpack("<I")
        if data[0] == 0:
            return None
        return self.subcls.parse(stream)

    def get_alignment(self):
        # 14.3.2 Alignment of Constructed Types
        # Pointer alignment is always modulo 4.
        return 4

class NdrUnpackNone(object):
    @classmethod
    def unpack(cls, stream):
        return None

class NdrRef(object):
    # TESTING
    def __init__(self, subcls):
        self.subcls = subcls

    def unpack(self, stream):
        ptr = NdrLong.unpack(stream)
        if not ptr:
            raise ValueError("Ndr REF cannot be NULL")
        return self.subcls.unpack(stream)

class NdrFixedArray(object):
    def __init__(self, subcls, size):
        self.subcls = subcls
        self.size = size

    def pack(self, data):
        data = list(data)
        assert len(data) == self.size
        return dword_pad(b"".join([self.subcls.pack(elt) for elt in data]))


    def unpack(self, stream):
        return [self.subcls.unpack(stream) for i in range(self.size)]

    def get_alignment(self):
        return self.subcls.get_alignment()


class NdrSID(object):
    @classmethod
    def pack(cls, psid):
        """Pack a PSID

        :param PSID psid:
        """
        subcount = windows.winproxy.GetSidSubAuthorityCount(psid)
        size = windows.winproxy.GetLengthSid(psid)
        sid_data = windows.current_process.read_memory(psid.value, size)
        return pack_dword(subcount[0]) + dword_pad(sid_data)

    @classmethod
    def unpack(cls, stream):
        """Unpack a PSID, partial implementation that returns a :class:`str` and not a PSID"""
        subcount = NdrLong.unpack(stream)
        return stream.read(8 + (subcount * 4))

    @classmethod
    def get_alignment(self):
        # Not sur, but it seems to contain an array of long
        return 4

class NdrVaryingCString(object):
    @classmethod
    def pack(cls, data):
        """Pack string ``data``. append ``\\x00`` if not present at the end of the string"""
        if data is None:
            return None
        if not data.endswith('\x00'):
            data += '\x00'
        l = len(data)
        result = struct.pack("<2I", 0, l)
        result += data
        return dword_pad(result)

    @classmethod
    def get_alignment(self):
        # Not sur, but size is on 4 bytes so...
        return 4

class NdrWString(object):
    @classmethod
    def pack(cls, data):
        """Pack string ``data``. append ``\\x00`` if not present at the end of the string"""
        if data is None:
            return None
        if not data.endswith('\x00'):
            data += '\x00'
        data = data.encode("utf-16-le")
        l = (len(data) // 2)
        result = struct.pack("<3I", l, 0, l)
        result += data
        return dword_pad(result)

    @classmethod
    def unpack(cls, stream):
        stream.align(4)
        size1, zero, size2 = stream.partial_unpack("<3I")
        assert size1 == size2
        assert zero == 0
        s = stream.read(size1 * 2)
        return s.decode("utf-16-le")

    @classmethod
    def get_alignment(self):
        # Not sur, but size is on 4 bytes so...
        return 4

class NdrCString(object):
    @classmethod
    def pack(cls, data):
        """Pack string ``data``. append ``\\x00`` if not present at the end of the string"""