import struct
import zlib
from spock.mcp import mcdata
from spock.mcp import datautils

extensions = tuple(tuple({} for i in j) for j in mcdata.packet_structs)
def extension(ident):
    def inner(cl):
        extensions[ident[0]][ident[1]][ident[2]] = cl
        return cl
    return inner

#Login SERVER_TO_CLIENT 0x01 Encryption Request
@extension((mcdata.LOGIN_STATE, mcdata.SERVER_TO_CLIENT, 0x01))
class ExtensionLSTC01:
    def decode_extra(packet, bbuff):
        length = mcdata.unpack('short', bbuff)
        packet.data['public_key'] = bbuff.recv(length)
        length = mcdata.unpack('short', bbuff)
        packet.data['verify_token'] = bbuff.recv(length)
        return packet

    def encode_extra(packet):
        o  = pack('short', len(packet.data['public_key']))
        o += packet.data['public_key']
        o += pack('short', len(packet.data['verify_token']))
        o += packet.data['verify_token']
        return o

#Play  SERVER_TO_CLIENT 0x0E Spawn Object
@extension((mcdata.PLAY_STATE, mcdata.SERVER_TO_CLIENT, 0x0E))
class ExtensionPSTC0E:
    def decode_extra(packet, bbuff):
        if packet.data['obj_data']:
            packet['speed_x'] = mcdata.unpack('short', bbuff)
            packet['speed_y'] = mcdata.unpack('short', bbuff)
            packet['speed_z'] = mcdata.unpack('short', bbuff)
        return packet

    def encode_extra(packet):
        if packet.data['obj_data']:
            o  = mcdata.pack('short', packet['speed_x'])
            o += mcdata.pack('short', packet['speed_y'])
            o += mcdata.pack('short', packet['speed_z'])
        return o

#Play  SERVER_TO_CLIENT 0x13 Destroy Entities
@extension((mcdata.PLAY_STATE, mcdata.SERVER_TO_CLIENT, 0x13))
class ExtensionPSTC13:
    def decode_extra(packet, bbuff):
        count = mcdata.unpack('byte', bbuff)
        packet['eids'] = [
            mcdata.unpack('int', bbuff) for i in range(count)
        ]
        return packet

    def encode_extra(packet):
        o = mcdata.pack('int', len(packet.data['eids']))
        for eid in packet.data['eids']:
            o += mcdata.pack('int', eid)

#Play  SERVER_TO_CLIENT 0x20 Entity Properties
@extension((mcdata.PLAY_STATE, mcdata.SERVER_TO_CLIENT, 0x20))
class ExtensionPSTC20:
    def decode_extra(packet, bbuff):
        packet.data['properties'] = []
        for i in range(mcdata.unpack('int', bbuff)):
            prop = {
                'key': mcdata.unpack('string', bbuff),
                'value': mcdata.unpack('double', bbuff),
                'modifiers': []
            }
            for j in range(mcdata.unpack('short', bbuff)):
                a, b = struct.unpack('>QQ', bbuff.recv(16))
                prop['modifiers'].append({
                    'uuid': (a<<64)|b,
                    'amount': mcdata.unpack('double', bbuff),
                    'operation': mcdata.unpack('byte', bbuff),
                })
            packet.data['properties'].append(prop)
        return packet

    def encode_extra(packet):
        o = mcdata.pack('int', len(packet.data['properties']))
        for prop in packet.data['properties']:
            o += mcdata.pack('string', prop['key'])
            o += mcdata.pack('double', prop['value'])
            o += mcdata.pack('short', len(prop['modifiers']))
            for modifier in prop['modifiers']:
                o += struct.pack('>QQ',
                    (modifier['uuid']>>64)&((1<<64)-1),
                    modifier['uuid']&((1<<64)-1)
                )
                o += mcdata.pack('double', modifier['amount'])
                o += mcdata.pack('byte', modifier['operation'])
        return o

#Play  SERVER_TO_CLIENT 0x21 Chunk Data
@extension((mcdata.PLAY_STATE, mcdata.SERVER_TO_CLIENT, 0x21))
class ExtensionPSTC21:
    def decode_extra(packet, bbuff):
        packet.data['byte_array'] = zlib.decompress(
            bbuff.recv(mcdata.unpack('int', bbuff))
        )
        return packet

    def encode_extra(packet):
        data = zlib.compress(packet.data['byte_array'])
        o = mcdata.pack('int', len(data))
        o += data
        return o

#Play  SERVER_TO_CLIENT 0x22 Multi Block Change
@extension((mcdata.PLAY_STATE, mcdata.SERVER_TO_CLIENT, 0x21))
class ExtensionPSTC22:
    def decode_extra(packet, bbuff):
        count = mcdata.unpack('short', bbuff)
        packet.data['blocks'] = []
        assert(count == mcdata.unpack('int', bbuff))
        for i in range(count):
            data = unpack('uint', bbuff)
            packet.data['blocks'].append({
                'metadata': (data    )&0xF,
                'type':     (data>> 4)&0xFFF,
                'y':        (data>>16)&0xFF,
                'z':        (data>>24)&0xF,
                'x':        (data>>28)&0xF,
            })
        return packet

    def encode_extra(packet):
        o  = mcdata.pack('short', len(packet.data['blocks']))
        o += mcdata.pack('int', 4*len(packet.data['blocks']))
        for block in packet.data['blocks']:
            o += mcdata.pack('uint',
                block['metadata']  +
                (block['type']<<4) +
                (block['y'] << 16) +
                (block['z'] << 24) +
                (block['x'] << 28)
            )
        return o