import struct
import os


header = None
supported_versions = [0xA, 0xB]


class Header:
    def __init__(self, endianness='<'):
        self.format = '3I4xI'
        self.endianness = endianness

        self.magic = 0x53484141  # SHAA
        self.version = 0xB
        self.fileSize = 0
        self.name = ''

    def load(self, data, pos=0):
        (self.magic,
         self.version,
         self.fileSize,
         nameLen) = struct.unpack_from('%s%s' % (self.endianness, self.format), data, pos)

        assert self.magic == 0x53484141

        self.name = data[pos + 0x14:pos + 0x14 + nameLen].decode('utf-8').rstrip('\0')
        self.size = 0x14 + nameLen

        assert self.version in supported_versions


class ShaderProgram:
    class VariationMacro:
        def __init__(self, endianness='<'):
            self.format = '3I'
            self.endianness = endianness

            self.size = 0
            self.name = ''
            self.value =''

        def __str__(self):
            return 'Variation Macro'

        def load(self, data, pos):
            (self.size,
             nameLen,
             valueLen) = struct.unpack_from('%s%s' % (self.endianness, self.format), data, pos)

            pos += struct.calcsize(self.format)
            self.name = data[pos:pos + nameLen].decode('utf-8').rstrip('\0')

            pos += nameLen
            self.value = data[pos:pos + valueLen].decode('utf-8').rstrip('\0')

    class VariationSymbol:
        def __init__(self, endianness='<'):
            self.format = '4I'
            self.endianness = endianness

            self.size = 0

        def __str__(self):
            return 'Variation Symbol'

        def load(self, data, pos):
            (self.size,
             _4,
             _8,
             _C) = struct.unpack_from('%s%s' % (self.endianness, self.format), data, pos)

            assert self.size == 8  # Structure is incomplete

    class ShaderSymbol:
        class Variable:
            def __init__(self, name='', default=b'', maxLen=-1):
                self.name = name
                self.default = default
                self.maxLen = maxLen

        def __init__(self, endianness='<'):
            self.format = 'Ii4I'
            self.endianness = endianness

            self.size = 0

            self.name = ''
            self.variable = ShaderProgram.ShaderSymbol.Variable()
            self.variationFlags = []

        def __str__(self):
            return 'Variation Symbol'

        def load(self, data, pos):
            pos_ = pos
            (self.size,
             self.variable.maxLen,
             variableNameLen,
             nameLen,
             defaultValueLen,
             variationCount) = struct.unpack_from('%s%s' % (self.endianness, self.format), data, pos)

            pos += struct.calcsize(self.format)
            self.variable.name = data[pos:pos + variableNameLen].decode('utf-8').rstrip('\0')

            pos += variableNameLen
            self.name = data[pos:pos + nameLen].decode('utf-8').rstrip('\0')

            pos += nameLen
            self.variable.default = data[pos:pos + defaultValueLen]

            pos += defaultValueLen
            for _ in range(variationCount):
                self.variationFlags.append(bool(data[pos])); pos += 1

    def __init__(self, endianness='<'):
        self.format = '2I3i'
        self.endianness = endianness

        self.size = 0
        self.vtxShIdx = -1
        self.frgShIdx = -1
        self.geoShIdx = -1

        self.vertexVariations = List(self.endianness)
        self.fragmentVariations = List(self.endianness)
        self.geometryVariations = List(self.endianness)

        self.variationSymbols = List(self.endianness)
        self.variationSymbols2 = List(self.endianness)

        self.uniformVariables = List(self.endianness)
        self.uniformBlocks = List(self.endianness)
        self.samplerVariables = List(self.endianness)
        self.attribVariables = List(self.endianness)

        self.name = ''

    def __str__(self):
        return 'Shader Program'

    def load(self, data, pos):
        (self.size,
         nameLen,
         self.vtxShIdx,
         self.frgShIdx,
         self.geoShIdx) = struct.unpack_from('%s%s' % (self.endianness, self.format), data, pos)

        pos += struct.calcsize(self.format)
        self.name = data[pos:pos + nameLen].decode('utf-8').rstrip('\0')

        pos += nameLen
        self.vertexVariations.load(data, pos, ShaderProgram.VariationMacro)

        pos += self.vertexVariations.size
        self.fragmentVariations.load(data, pos, ShaderProgram.VariationMacro)

        pos += self.fragmentVariations.size
        self.geometryVariations.load(data, pos, ShaderProgram.VariationMacro)

        pos += self.geometryVariations.size
        self.variationSymbols.load(data, pos, ShaderProgram.VariationSymbol)
        pos += self.variationSymbols.size

        if header.version == 11:
            self.variationSymbols2.load(data, pos, ShaderProgram.VariationSymbol)
            pos += self.variationSymbols2.size

        self.uniformVariables.load(data, pos, ShaderProgram.ShaderSymbol)

        pos += self.uniformVariables.size
        self.uniformBlocks.load(data, pos, ShaderProgram.ShaderSymbol)

        pos += self.uniformBlocks.size
        self.samplerVariables.load(data, pos, ShaderProgram.ShaderSymbol)

        pos += self.samplerVariables.size
        self.attribVariables.load(data, pos, ShaderProgram.ShaderSymbol)


class ShaderCode:
    def __init__(self, endianness='<'):
        self.format = '4I'
        self.endianness = endianness

        self.size = 0
        self.name = ''
        self.code = ''

    def __str__(self):
        return 'Shader Code'

    def load(self, data, pos):
        (self.size,
         nameLen,
         codeLen,
         codeLen2) = struct.unpack_from('%s%s' % (self.endianness, self.format), data, pos)

        pos += struct.calcsize(self.format)
        self.name = data[pos:pos + nameLen].decode('utf-8').rstrip('\0')

        pos += nameLen; assert codeLen == codeLen2
        self.code = data[pos:pos + codeLen].decode('shift-jis')

        #with open(self.name, 'wb+') as out:
        #    out.write(self.code.encode('utf-8'))


class List:
    def __init__(self, endianness='<'):
        self.format = '2I'
        self.endianness = endianness

        self.size = 0
        self.count = 0
        self.items = []

    def __getitem__(self, i):
        if not isinstance(i, int):
            raise TypeError("index must be an integer")

        return self.items[i]

    def append(self, item):
        self.items.append(item)

    def extend(self, itemList):
        self.items.extend(itemList)

    def index(self, item):
        for i, oItem in enumerate(self.items):
            if item == oItem:
                return i

        return -1

    def load(self, data, pos, ItemClass=None):
        (self.size,
         self.count) = struct.unpack_from('%s%s' % (self.endianness, self.format), data, pos)

        pos += struct.calcsize(self.format)
        if ItemClass:
            for _ in range(self.count):
                item = ItemClass(self.endianness)
                item.load(data, pos)

                pos += item.size
                self.append(item)


def load(inb, pos=0):
    global header
    header = Header()
    header.load(inb, pos)

    pos += header.size

    progList = List()
    progList.load(inb, pos, ShaderProgram)

    pos += progList.size

    codeList = List()
    codeList.load(inb, pos, ShaderCode)

    pos += codeList.size

    return progList, codeList