from filetypes.base import *
import malcat
from filetypes.BMP import RGB, RGBA


class FileHeader(Struct):

    def parse(self):
        yield String(6, name="Signature", comment="signature")
        w = yield UInt16(name="Width", comment="Logical screen width")
        if w == 0 or w > 6000:
            raise FatalError("Invalid width")
        h = yield UInt16(name="Height", comment="Logical screen height")
        if h == 0 or h > 6000:
            raise FatalError("Invalid height")
        flags = yield BitsField(
            Bit(name="GlobalMapSize0", comment="size of global map"),
            Bit(name="GlobalMapSize1", comment="size of global map"),
            Bit(name="GlobalMapSize2", comment="size of global map"),
            Bit(name="GlobalMapSorted", comment="is the global map sorted by decreasing importance?"),
            Bit(name="ColorDepth0", comment="bpp"),
            Bit(name="ColorDepth1", comment="bpp"),
            Bit(name="ColorDepth2", comment="bpp"),
            Bit(name="HasGlobalMap", comment="is global map present?"),
            name="Flags", comment="flags")
        yield UInt8(name="BackgroundColor", comment="background color")
        yield UInt8(name="PixelRatio", comment="pixel aspect ratio")

class UnknownExtension(Struct):

    def parse(self):
        yield UInt8(name="Separator")
        yield UInt8(name="ExtensionId")
        sz = yield UInt8(name="Size", comment="block size")
        if sz:
            yield Bytes(sz, name="Data")
        yield Unused(1, name="Terminator")

class GraphicControl(Struct):

    def parse(self):
        yield UInt8(name="Separator")
        yield UInt8(name="ExtensionId")
        sz = yield UInt8(name="Size", comment="block size")
        flags = yield BitsField(
            Bit(name="Transparent", comment="has transparency"),
            Bit(name="UserInput", comment="user input"),
            name="Flags", comment="flags")

        yield UInt16(name="Delay", comment="animation delay in milliseconds")
        yield UInt8(name="TransparentColor", comment="transparent color index")
        yield Unused(1, name="Terminator")


class Extension(Struct):

    def parse(self):
        sz = yield UInt8(name="Size")
        if sz:
            yield Bytes(sz, name="Data")


class ApplicationExtension(Struct):

    def parse(self):
        yield UInt8(name="Separator")
        yield UInt8(name="ExtensionId")
        yield PascalString8(name="Application")
        while True:
            ex = yield Extension()
            if not ex["Size"]:
                break

class Image(Struct):

    def parse(self): 
        yield UInt8(name="Separator")
        yield UInt16(name="Left")
        yield UInt16(name="Top")
        yield UInt16(name="Width")
        yield UInt16(name="Height")
        flags = yield BitsField(
            Bit(name="LocalMapSize0", comment="size of global map"),
            Bit(name="LocalMapSize1", comment="size of global map"),
            Bit(name="LocalMapSize2", comment="size of global map"),
            NullBits(2),
            Bit(name="LocalMapSorted", comment="is the global map sorted by decreasing importance?"),
            Bit(name="Interlaced", comment="is image interlaced?"),
            Bit(name="HasLocalMap", comment="is local map present?"),
            name="Flags", comment="flags")
        if flags["HasLocalMap"]:
            sz = 1 + flags["LocalMapSize0"] + flags["LocalMapSize1"] * 2 + flags["LocalMapSize2"] * 4
            yield Array(1 << sz, RGB(), name="Palette")
        yield UInt8(name="LzwMinimumCodeSize")
        while self.remaining():
            block = yield Extension(name="ImageBlock")
            if block["Size"] == 0:
                break
        
        

EXTENSIONS = {
        0xf9: (GraphicControl, Type.HEADER),
        0xff: (ApplicationExtension, Type.META),
}


class GIFAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.IMAGE
    name = "GIF"
    regexp = r"GIF8[79]a"

    def parse(self, hint):
        start = self.tell()
        fh = yield FileHeader(category=Type.HEADER)
        bpp = 1 + fh["Flags"]["ColorDepth0"] + fh["Flags"]["ColorDepth1"] * 2 + fh["Flags"]["ColorDepth2"] * 4
        if fh["Flags"]["HasGlobalMap"]:
            sz = 1 + fh["Flags"]["GlobalMapSize0"] + fh["Flags"]["GlobalMapSize1"] * 2 + fh["Flags"]["GlobalMapSize2"] * 4
            yield Array(1 << sz, RGB(), name="Palette")
        while self.remaining():
            prefix = self.read(self.tell(), 1)
            if prefix == b"!" and self.remaining() > 1:
                extid = self.read(self.tell() + 1, 1)[0]
                ext, cat = EXTENSIONS.get(extid, (UnknownExtension, Type.DATA))
                yield ext(category=cat)
            elif prefix == b",":
                yield Image(category=Type.DATA)
                self.confirm()
            elif prefix == b";":
                yield UInt8(name="EndOfFile", category=Type.HEADER)
            else:
                raise FatalError("Invalid block prefix: {}".format(prefix))

