"""
 * Copyright (C) 2013 Stellar Jockeys.
 * All Rights Reserved.
 *
 * NOTICE: All information contained herein is, and remains the property of
 * Stellar Jockeys and its suppliers, if any.  The intellectual and technical
 * concepts contained herein are proprietary to Stellar Jockeys and its suppliers.
 * Dissemination of this information or reproduction of this material is
 * strictly forbidden unless prior written permission is obtained from
 * Stellar Jockeys.
 *
 * @file genmanifest.py
 * @author Harry Hsiao
 * @date 01/28/2013
"""
import os
import re
import json
from functools import partial

"""
Shaders should have a .vert and .frag file with the same file name.
The extensions must be lowercase.

'//#include <[filename]>' lines are replaced with the contents of [filename].
No leading characters should be present. Any trailing characters are ignored.
[filename] is relative to the file which contains the statement.

In either the .vert or .frag file a '//#required' comment marks the
shader as required for the game. This is done after includes.

The includes are replaced before comments are stripped.
Inline comments are stripped before block comments.
Do not mix comments on the same line nor nest block comments.
Nested included files are not processed.

This script assumes valid glsl 1.20.8 and does not check for duplicate declarations
or unused variables.

For the shader global variables, the following are currently unsupported:
 * Multiple declarations per line (ie. vec3 foo, bar;)
 * Structures
 * Arrays
"""
# Output formatting
INDENT = 4
SEPARATORS = (',', ': ')
SORTED = True

# NOTE: Designed to be called from project root directory.
OUTPUT_FILE = "shaders/manifest.json"
INPUT_DIR = "shaders"

UNIFORMS = "uniforms"
TEXTURES = "textures"
ATTRIBS = "attributes"
VARYINGS = "varyings"
KEYS = [UNIFORMS, TEXTURES, ATTRIBS, VARYINGS]
VERT_PREFIX = "vertex_"
FRAG_PREFIX = "fragment_"
VERT_EXT = ".vert"
FRAG_EXT = ".frag"
REQUIRED = "required"
SHADER = "shader"
NAME = "name"
SIZE = "size"

QUALIFIER = r'(uniform|attribute|varying)'
TYPE = r'(float|int|bool|[ib]?vec[234]|mat[234](?:x[234])?|sampler(?:Cube|[123]D|[12]DShadow))'
IDENTIFIER = r'([a-zA-Z_][a-zA-Z0-9_]*)'
END = r'(?:;|[^a-zA-Z0-9_;][^;]*;)'
SPACE = r'(?:\s+)'
DECLARATION = QUALIFIER + SPACE + TYPE + SPACE + IDENTIFIER + END
BLOCK_COMMENT = r'/\*(?:[^\*]|\*(?!/))*\*/'
INLINE_COMMENT = r'//.*'
REQUIRED_PATTERN = r'//[ \t]*#required(?:\s|$)'
INCLUDE_PATTERN = r'(?:^|(?<=\n))//#include <([^>]*)>.*(?:\n|$)'
VECTOR = r'^[ib]?vec([234])$'
MATRIX = r'^mat([234])(?:x([234]))?$'

def GetFiles(root):
    """
    @brief Gets a list of (path, name) of shaders programs in the provided directory.

    The name excludes the extension.
    """
    shaders = []
    for path, dirs, files in os.walk(root):
        verts = [name[:-len(VERT_EXT)] for name in files if name.endswith(VERT_EXT)]
        frags = [name[:-len(FRAG_EXT)] for name in files if name.endswith(FRAG_EXT)]
        shaders.extend([(path, name) for name in verts if (name in frags)])
    return shaders


def TypeToComponents(typeString):
    """
    @brief Gets the number of components based on the type string.
    """
    result = re.match(VECTOR, typeString)
    if result is not None:
        return int(float(result.group(1)))
    result = re.match(MATRIX, typeString)
    if result is not None:
        columns = int(float(result.group(1)))
        rows = columns
        if result.group(2) is not None:
            rows = int(float(result.group(2)))
        return columns * rows
    return 1


def TypeToAttribs(typeString):
    """
    @brief Gets the number of attribs based on the type string.
    """
    result = re.match(MATRIX, typeString)
    if result is not None:
        return int(float(result.group(1)))
    return 1


def DeclToEntry(declaration):
    """
    @brief Creates a dictionary for the variable declaration.
    """
    size = 0
    if declaration[0] == "attribute":
        size = TypeToAttribs(declaration[1])
    else:
        size = TypeToComponents(declaration[1])
    entry = {}
    entry[NAME] = declaration[2]
    entry[SIZE] = size
    return entry


def DeclToKey(declaration):
    """
    @brief Gets the key string for a variable declaration.
    """
    if declaration[1].startswith("sampler"):
        return TEXTURES
    return declaration[0] + "s"


def ReplaceInclude(sourceFile, directory, match):
    """
    @brief Helper to replace include directives
    """
    filename = directory + '/' + match.group(1)
    source = ""
    try:
        with open(filename, 'r') as include:
            source = include.read()
    except IOError as e:
        print "ERROR: Failed to include '" + filename + "' in shader '" + sourceFile + "'."
        raise IOError(e)
    return source


def AnalyzeShader(filename, directory):
    """
    @brief Creates a dictionary for a shader file containing lists each key string.
    """
    source = ""
    with open(filename, 'r') as sourceFile:
        source = sourceFile.read()
    replacement = partial(ReplaceInclude, filename, directory)
    source = re.sub(INCLUDE_PATTERN, replacement, source)
    required = re.search(REQUIRED_PATTERN, source) is not None
    source = re.sub(INLINE_COMMENT, "", source)
    source = re.sub(BLOCK_COMMENT, "", source)
    declarations = re.findall(DECLARATION, source)
    results = {key:[] for key in KEYS}
    results[REQUIRED] = required
    for declaration in declarations:
        results[DeclToKey(declaration)].append(DeclToEntry(declaration))
    return results


def AnalyzeProgram(directory, name):
    """
    @brief Merges the results from analyzing the vertex and fragment shaders.
    """
    filename = directory + '/' + name
    vertFile = filename + VERT_EXT
    fragFile = filename + FRAG_EXT
    vert = AnalyzeShader(vertFile, directory)
    frag = AnalyzeShader(fragFile, directory)
    varyings = [v[SIZE] for v in frag[VARYINGS] if (v in vert[VARYINGS])]
    program = {}
    program[NAME] = name
    program[VARYINGS] = sum(varyings)
    program[REQUIRED] = vert[REQUIRED] or frag[REQUIRED]
    program[VERT_PREFIX + SHADER] = vertFile
    program[FRAG_PREFIX + SHADER] = fragFile
    for key in KEYS:
        if key == VARYINGS:
            continue
        program[VERT_PREFIX + key] = vert[key]
        if key != ATTRIBS:
            program[FRAG_PREFIX + key] = frag[key]
    return program


def GenerateManifest():
    manifest = [AnalyzeProgram(directory, name) for (directory, name) in GetFiles(INPUT_DIR)]
    with open(OUTPUT_FILE, 'w') as manifestFile:
        json.dump(manifest,
                  manifestFile,
                  indent=INDENT,
                  separators=SEPARATORS,
                  sort_keys=SORTED)


def Test(tests, function, name):
    success = True
    for (expected, test) in tests:
        actual = function(test)
        if (actual != expected):
            print "FAILED test '%s': Expected '%s', got '%s'. '%s'" \
                  % (name, expected, actual, test)
            success = False
    return success


def TestHelperCalculations():
    tests = [((1, 1), "float"), ((1, 1), "int"), ((1, 1), "bool"), ((1, 1), "sampler2D"),
             ((2, 1), "vec2"), ((3, 1), "vec3"), ((4, 1), "vec4"),
             ((4, 2), "mat2"), ((4, 2), "mat2x2"), ((6, 2), "mat2x3"), ((8, 2), "mat2x4"),
             ((9, 3), "mat3"), ((6, 3), "mat3x2"), ((9, 3), "mat3x3"), ((12, 3), "mat3x4"),
             ((16, 4), "mat4"), ((8, 4), "mat4x2"), ((12, 4), "mat4x3"), ((16, 4), "mat4x4")]
    def function(test):
        return (TypeToComponents(test), TypeToAttribs(test))
    return Test(tests, function, "Helpers")


def TestDeclarations():
    success = True
    tests = [([('attribute', 'float', 'foo'),
               ('uniform', 'vec3', 'bar'),
               ('varying', 'mat2x4', 'baz'),
               ('varying', 'int', 'corge5')],
              "attribute float foo; uniform vec3 bar = vec3(0.0, 0.0, 0.0);\n\
               centriod varying mat2x4 baz; const float qux; varying int corge5;")]
    def function(test):
        return re.findall(DECLARATION, test)
    return Test(tests, function, "Declarations")


def TestRequired():
    tests = [(False, "// nothing here. # req inc < >"),
             (True, "// blah\n //#required\r\n"),
             (True, "// #required")]
    def function(test):
        return re.search(REQUIRED_PATTERN, test) is not None
    return Test(tests, function, "Required")


def TestInclude():
    tests = [([("a")], "//#include <a>"),
             ([], "a //#include <asdf>"),
             ([], "// #include <asdf>"),
             ([], "//#include<asdf>"),
             ([("asdf/zxcv.txt")], "\n//#include <asdf/zxcv.txt>\n\t//#include <a>"),
             ([("asdf"), ("qwer")], "//#include <asdf> tgmjgks\ntyu\n//#include <qwer>\r")]
    def function(test):
        return re.findall(INCLUDE_PATTERN, test)
    return Test(tests, function, "Include")


def Tests():
    success = [TestHelperCalculations(),
               TestDeclarations(),
               TestRequired(),
               TestInclude()]
    return all(success)


if Tests():
    GenerateManifest()
    print "DONE."
