"""
riffogg.py
Extract embedded OGG files from RIFF/WAVE files. Works on both Python 2 and Python 3.
Usage: riffogg.py input.wav output.ogg

Copyright (C) 2014 Shiz
Released under the terms of the WTFPL: see http://www.wtfpl.net/txt/copying/ for details.
"""
from __future__ import print_function
import sys
import os
import struct

RIFF_MAGIC = b'RIFF'
WAVE_MAGIC = b'WAVE'
OGG_MAGIC = b'OggS'
# Chunks that contain chunks.
META_CHUNKS = (b'RIFF', b'LIST')
# Chunks that don't contain anything.
GARBAGE_CHUNKS = (b'WAVE', b'INFO', b'adtl')


def validate_header(f):
    """ Make sure opened file is a RIFF WAVE file and return its full length, or -1 if it isn't. """
    f.seek(0, os.SEEK_SET)

    magic = f.read(4)
    if magic != RIFF_MAGIC:
        print('RIFF magic: expected {}, got {}'.format(RIFF_MAGIC, magic))
        return -1

    length = struct.unpack('<I', f.read(4))[0]
    if length <= 0:
        print('RIFF file: invalid length specified ({})'.format(length))
        return -1

    wave_magic = f.read(4)
    if wave_magic != WAVE_MAGIC:
        print('WAVE magic: expected {}, got {}'.format(WAVE_MAGIC, wave_magic))
        return -1

    return length + 8


def extract_chunks(f, max_length=None):
    """ Extract RIFF chunks from file. Yields (chunktype, data) tuples. """
    f.seek(0, os.SEEK_SET)
    read_length = 0

    while not max_length or read_length < max_length:
        chunk_type = f.read(4)
        if not chunk_type:
            # End of file.
            break
        read_length += 4

        # Neat, a garbage chunk. Can't do much with it, though.
        if chunk_type in GARBAGE_CHUNKS:
            print('got garbage chunk: {}'.format(chunk_type))
            yield (chunk_type, b'')
            continue
        else:
            chunk_length = struct.unpack('<I', f.read(4))[0]
            read_length += 4

        # Meta chunks are neat too, but we can't really process them well.
        if chunk_type in META_CHUNKS:
            print('got meta-chunk: {} (length = {})'.format(chunk_type, chunk_length))
            yield (chunk_type, b'')
        else:
            print('got chunk: {} (length = {})'.format(chunk_type, chunk_length))

            chunk_data = f.read(chunk_length)
            read_length += chunk_length
            yield (chunk_type, chunk_data)

            # Chunks are padded to even numbers.
            if chunk_length % 2:
                f.read(1)
                read_length += 1


if __name__ == '__main__':
    if len(sys.argv) < 3:
        exit('usage: {} FILE OUTFILE'.format(sys.argv[0]))

    infile = sys.argv[1]
    outfile = sys.argv[2]

    with open(infile, 'rb') as f:
        length = validate_header(f)
        if length <= 0:
            exit('not a RIFF file: {}'.format(infile))

        # Iterate over chunks to find a data chunk containing OGG data.
        for (type, data) in extract_chunks(f, length):
            if type == b'data' and data[:4] == OGG_MAGIC:
                # Gotcha!
                print('found embedded OGG data!', file=sys.stderr)
                print('writing data to {}...'.format(outfile), file=sys.stderr)

                with open(outfile, 'wb') as out:
                    out.write(data)

                break
        else:
                print('could not find embedded OGG data. :(', file=sys.stderr)