from filetypes.base import *
import malcat

# http://www-mmsp.ece.mcgill.ca/Documents/AudioFormats/WAVE/WAVE.html

class RiffHeader(Struct):

    def parse(self):
        yield String(4, zero_terminated=False, name="ChunkId")
        yield UInt32(name="ChunkSize", comment="size of the rest of the chunk following this number. This is the size of the entire file in bytes minus 8 bytes")
        yield String(4, zero_terminated=False, name="Format")


class WaveHeader(Struct):
    def parse(self):
        yield String(4, zero_terminated=False, name="SubChunkId")
        sz = yield UInt32(name="SubChunkSize", comment="16 for PCM.  This is the size of the rest of the Subchunk which follows this number")
        fmt = yield UInt16(name="AudioFormat", values=[
            ("PCM", 1),
            ("MS-ADPCM", 2),
            ("IEEE_FLOAT", 3),
            ("ALAW", 6),
            ("MULAW", 7),
            ("EXTENSIBLE", 0xFFFE),
            ])
        yield UInt16(name="NumChannels", comment="Mono = 1, Stereo = 2, etc.")
        yield UInt32(name="SampleRate", comment="8000, 44100, etc.")
        yield UInt32(name="ByteRate", comment="= SampleRate * NumChannels * BitsPerSample/8")
        yield UInt16(name="BlockAlign", comment="= NumChannels * BitsPerSample/8")
        yield UInt16(name="BitsPerSample", comment="8 bits = 8, 16 bits = 16, etc.")
        if fmt != 1:
            extra = yield UInt16(name="ExtraParamSize", comment="size of extra parameters")
            if extra == 22:
                yield UInt16(name="ValidBitsPerSample", comment="number of valid bits")
                yield UInt32(name="ChannelMask", comment="speaker position mask")
                yield GUID(name="SubFormat", comment="GUID, including the data format code")
            elif extra:
                yield Bytes(extra, name="ExtraParam", comment="extra parameters")
        if len(self) - 8 < sz:
            yield Unused(sz - (len(self) - 8))

class DataHeader(Struct):
    def parse(self):
        yield String(4, zero_terminated=False, name="SubChunkId")
        datasz = yield UInt32(name="SubChunkSize", comment="size of data blob")
        yield Bytes(datasz, name="Data")


class ChunkHeader(Struct):
    def parse(self):
        yield String(4, zero_terminated=False, name="SubChunkId")
        datasz = yield UInt32(name="SubChunkSize", comment="size of data blob")
        yield Bytes(datasz, name="Data")

class WavAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.SOUND
    name = "WAV"
    regexp = r"RIFF....WAVEfmt "

    def __init__(self):
        FileTypeAnalyzer.__init__(self)

    def parse(self, hint):
        riff = yield RiffHeader(category=Type.HEADER)
        yield WaveHeader(category=Type.HEADER)
        while self.tell() < riff["ChunkSize"] + 8:
            start = self.tell()
            chunk_hdr = self.read(self.tell(), 4)
            if chunk_hdr == b"data":
                chk = yield DataHeader(category=Type.DATA)
            else:
                chk = yield ChunkHeader(category=Type.META)
            size = self.tell() - start
            self.add_section(chunk_hdr.decode("ascii", errors="ignore"), start, size)


