#!/usr/bin/env python3
# coding=utf-8

# Copyright (c) 2019 Intel Corporation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

""" Generate in/out tests for 64 bit types. """

import os
import errno
import collections
import copy
import sys

from textwrap import dedent
from mako.template import Template
from mako import exceptions

import random_ubo



class VaryingType(object):
    __slots__ = ['type', 'members', 'name']

    def __init__(self, type, members=None):
        self.type = type
        self.members = members
        self.name = None


VaryingInstance = collections.namedtuple(
    'VaryingInstance', ['varying_type', 'setters', 'checkers', 'full_name'])


def scalar_assignment(type, name, data):
    """Return a GLSL code string to assign a scalar to its expected value."""
    if type == "bool":
        if int(data) == 0:
            return "{} = {};".format(name,
                                     'false' if int(data) == 0 else 'true')
    elif type == "uint":
        return "{} = {}u;".format(name, data)
    elif type == "int":
        return "{} = {};".format(name, data)
    elif type == "uint64_t":
        return "{} = {}ul;".format(name, data)
    elif type == "int64_t":
        return "{} = {}l;".format(name, data)
    elif type == "float":
        # Not all implementations support the bit-cast operators that are used
        # to do bit-exact comparisons.  For this reason float_match needs the
        # float data and the bit-exact data.

        bits = random_ubo.bit_exact_data(data, "float")
        return "{} = float_get({}, {}u);".format(name, data, bits)
    elif type == "double":
        bits = random_ubo.bit_exact_data(data, "double")

        # 0xHHHHHHHHLLLLLLLL
        # 012345678901234567

        hi = "0x" + bits[2:10]
        lo = "0x" + bits[10:18]

        return "{} = double_get(uvec2({}, {}));".format(name, lo, hi)
    else:
        raise Exception("Unknown scalar type {}".format(type))


def vector_assignment(type, name, data):
    """Return a list of GLSL code strings that assign each field of a vector
    to its expected value.
    """
    scalar = random_ubo.vector_base_type(type)
    components = ["x", "y", "z", "w"]

    return [scalar_assignment(scalar,
                              "{}.{}".format(name, "xyzw"[i]),
                              data[i])
            for i in range(random_ubo.vector_size(type))]


def matrix_assignment(type, name, data):
    """Return a list of GLSL code strings that assign each field of a matrix
    its expected value.
    """
    c, r = random_ubo.matrix_dimensions(type)

    if type[0] == 'd':
        column_type = "dvec{}".format(r)
    else:
        column_type = "vec{}".format(r)

    data_pairs = []

    for i in range(c):
        data_pairs.extend(vector_assignment(
            column_type,
            "{}[{}]".format(name, i),
            data[(i * r):(i * r) + r]))

    return data_pairs


def create_array(base_type, size):
    return "{}[{}]".format(base_type, size)


def array_base_type(varying_type):
    t = copy.copy(varying_type)
    t.type = random_ubo.array_base_type(t.type)
    t.name = None
    return t


def create_getters_and_setters(varying_type, full_name, instances):
    if not full_name:
        full_name = varying_type.name
    elif varying_type.name is not None:
        full_name += '.' + varying_type.name

    if random_ubo.isarray(varying_type.type):
        base_type = array_base_type(varying_type)
        for i in range(random_ubo.array_elements(varying_type.type)):
            indexed_name = full_name + '[' + str(i) + ']'
            create_getters_and_setters(base_type, indexed_name, instances)

    elif random_ubo.isstructure(varying_type.type):
        for m in varying_type.members:
            create_getters_and_setters(m, full_name, instances)

    else:
        v = VaryingInstance(varying_type, [], [], full_name)

        raw_data = random_ubo.random_data(varying_type.type, full_name, 0)
        data = raw_data.split(" ")

        if random_ubo.isscalar(varying_type.type):
            v.checkers.append(random_ubo.scalar_derp(varying_type.type,
                                                     full_name,
                                                     data[0]))
            v.setters.append(
                scalar_assignment(varying_type.type, full_name, data[0]))
        elif random_ubo.isvector(varying_type.type):
            v.checkers.extend(random_ubo.vector_derp(varying_type.type,
                                                     full_name,
                                                     data))
            v.setters.extend(
                vector_assignment(varying_type.type, full_name, data))
        elif random_ubo.ismatrix(varying_type.type):
            v.checkers.extend(random_ubo.matrix_derp(varying_type.type,
                                                     full_name,
                                                     data))
            v.setters.extend(
                matrix_assignment(varying_type.type, full_name, data))

        instances.append(v)


def assign_names(varying_type, names):
    if random_ubo.isarray(varying_type.type):
        varying_type.name = names.get_name(
            random_ubo.without_array(varying_type.type))
    else:
        varying_type.name = names.get_name(varying_type.type)

    if varying_type.members:
        for m in varying_type.members:
            assign_names(m, names)


def gather_structs(varying_type, structs):
    if random_ubo.isarray(varying_type.type):
        base_type = array_base_type(varying_type)
        gather_structs(base_type, structs)
    elif random_ubo.isstructure(varying_type.type):
        for m in varying_type.members:
            gather_structs(m, structs)

        structs.append(varying_type)


def stringify_array_dimensions(type):
    s = ""
    base_type = type
    while random_ubo.isarray(base_type):
        s += "@{}".format(random_ubo.array_elements(base_type))
        base_type = random_ubo.array_base_type(base_type)
    return s


def stringify_varying_type(varying_type):
    if random_ubo.isarray(varying_type.type):
        name = "{}{}".format(random_ubo.without_array(varying_type.type),
                             stringify_array_dimensions(varying_type.type))
    else:
        name = varying_type.type

    if varying_type.members:
        for m in varying_type.members:
            name += "-" + stringify_varying_type(m)

    return name


def generate_file_name(root_types, explicit_locations):
    name = "vs-out-fs-in-"
    name += "-and-".join(
        [stringify_varying_type(t) for t in root_types])
    if explicit_locations:
        name += "-location-0"
    return name + ".shader_test"


def do_test(root_types, explicit_locations, glsl_version, extensions,
            tests_path):
    if explicit_locations and len(root_types) > 1:
        return

    basename = generate_file_name(root_types, explicit_locations)
    fullname = os.path.join(tests_path, basename)

    print(fullname)

    shader_file = open(fullname, "w", buffering=1)

    names = random_ubo.unique_name_dict()

    instances = []
    structs = []

    for type in root_types:
        assign_names(type, names)
        create_getters_and_setters(type, None, instances)
        gather_structs(type, structs)

    t = Template(dedent("""\
    [require]
    GLSL >= ${glsl_version // 100}.${glsl_version % 100}
    % for ext in extensions:
    ${ext}
    % endfor

    % if explicit_locations:
    GL_ARB_explicit_attrib_location
    % endif

    [vertex shader]
    % for ext in extensions:
    #extension ${ext}: require
    % endfor
    #extension GL_ARB_shader_bit_encoding: enable
    #extension GL_ARB_gpu_shader5: enable

    % if explicit_locations:
    #extension GL_ARB_explicit_attrib_location : require
    % endif

    precision highp float;
    % for s in structures:

    struct ${s.type} {
        % for m in s.members:
        ${"{:<15}".format(m.type)} ${m.name};
        % endfor
    };
    % endfor

    % for r in root_types:
    % if explicit_locations:
    layout(location = 0)
    % endif
    flat out ${"{:<10}".format(r.type)} ${r.name};
    % endfor

    in vec4 piglit_vertex;

    #if defined(GL_ARB_shader_bit_encoding) || defined(GL_ARB_gpu_shader5) || __VERSION__ >= 430
    float float_get(float f, uint bits) { return uintBitsToFloat(bits); }
    #else
    float float_get(float f, uint bits) { return f; }
    #endif
    % if glsl_version >= 400 or "GL_ARB_gpu_shader_fp64" in extensions:
    double double_get(uvec2 bits) { return packDouble2x32(bits); }
    %endif

    void main()
    {
    % for inst in instances:
        % for s in inst.setters:
        ${s}
        % endfor
    % endfor

        gl_Position = piglit_vertex;
    }

    [fragment shader]
    % for ext in extensions:
    #extension ${ext}: require
    % endfor
    #extension GL_ARB_shader_bit_encoding: enable
    #extension GL_ARB_gpu_shader5: enable

    % if explicit_locations:
    #extension GL_ARB_explicit_attrib_location : require
    % endif

    precision highp float;

    % for s in structures:
    struct ${s.type} {
        % for m in s.members:
        ${"{:<15}".format(m.type)} ${m.name};
        % endfor
    };
    % endfor

    % for r in root_types:
    % if explicit_locations:
    layout(location = 0)
    % endif
    flat in ${"{:<11}".format(r.type)} ${r.name};
    % endfor

    out vec4 piglit_fragcolor;

    #if defined(GL_ARB_shader_bit_encoding) || defined(GL_ARB_gpu_shader5) || __VERSION__ >= 430
    bool float_match(float u, float f, uint bits) { return floatBitsToUint(u) == bits; }
    #else
    bool float_match(float u, float f, uint bits) { return u == f; }
    #endif
    % if glsl_version >= 400 or "GL_ARB_gpu_shader_fp64" in extensions:
    bool double_match(double u, uvec2 bits) { return unpackDouble2x32(u) == bits; }
    %endif

    void main()
    {
        bool pass = true;

    % for inst in instances:
        % for s in inst.checkers:
        if (${s})
            pass = false;
        % endfor
    % endfor

        piglit_fragcolor = pass ? vec4(0, 1, 0, 1) : vec4(1, 0, 0, 1);
    }

    [test]
    link success

    draw rect -1 -1 2 2
    probe all rgba 0.0 1.0 0.0 1.0"""))

    try:
        shader = t.render(glsl_version=glsl_version,
                        root_types=root_types,
                        explicit_locations=explicit_locations,
                        structures=structs,
                        extensions=extensions,
                        instances=instances)
    except:
        print(exceptions.text_error_template().render(), file=sys.stderr)
        raise

    shader_file.write(shader)
    shader_file.close()


def do_test_permutations(root_types, glsl_version, extensions, tests_path):
    without_outer_struct = []
    for type in root_types:
        if random_ubo.isstructure(type.type):
            without_outer_struct.extend(m for m in type.members)
        else:
            without_outer_struct.append(type)

    do_test(root_types, False, glsl_version, extensions, tests_path)
    do_test(root_types, True, glsl_version,
            extensions + ["GL_ARB_separate_shader_objects"], tests_path)
    do_test(without_outer_struct, False, glsl_version, extensions, tests_path)


fp64_ext_path = os.path.join("spec", "arb_gpu_shader_fp64", "execution",
                             "inout")
fp64_core_path = os.path.join("spec", "glsl-4.00", "execution",
                              "inout")
int64_path = os.path.join("spec", "arb_gpu_shader_int64", "execution",
                          "inout")

for path in [fp64_ext_path, fp64_core_path, int64_path]:
    try:
        os.makedirs(path)
    except OSError as exc:
        if exc.errno != errno.EEXIST or not os.path.isdir(path):
            raise

INT64_VEC_TYPES = [None, "int64_t", "i64vec2", "i64vec3", "i64vec4"]
UINT64_VEC_TYPES = [None, "uint64_t", "u64vec2", "u64vec3", "u64vec4"]
FLOAT_VEC_TYPES = [None, "float", "vec2", "vec3", "vec4"]
DOUBLE_VEC_TYPES = [None, "double", "dvec2", "dvec3", "dvec4"]

DOUBLE_MAT_TYPES = ["dmat2x2", "dmat2x3", "dmat2x4", "dmat3x2", "dmat3x3",
                    "dmat3x4", "dmat4x2", "dmat4x3", "dmat4x4"]

VT = VaryingType


def do_common_64bit_tests(glsl_version, extensions, name_arr, tests_path):
    tests = []
    for i in range(1, 4):
        tests.append([VT("S1", [VT(name_arr[i])])])

        tests.append([VT("S1", [VT(name_arr[i] + "[4]")])])

        for j in range(1, 4):
            tests.append([VT("S1", [VT(FLOAT_VEC_TYPES[j]), VT(name_arr[i])])])

        tests.append([VaryingType("S1",
                                  [VT("float"), VT("float"), VT("float"),
                                   VT(name_arr[i])])])

        for j in range(2, 5):
            tests.append([VT("S1",
                             [VT("float[{}]".format(j)),
                              VT(name_arr[i])])])

        tests.append([VT("S1", [VT(name_arr[i] + "[3]")])])

        tests.append([VT("S1",
                         [VT("S2[3]", [VT(name_arr[i]), VT("float")])])])

        tests.append([VT("S1",
                         [VT("S2", [VT(name_arr[i])])])])

    for i in range(1, 2):
        tests.append([VT("S1",
                         [VT("S2[2]",
                             [VT("S3[2]", [
                                 VT("float"),
                                 VT(name_arr[i])])])])])

        tests.append([VT("S1",
                         [VT("S2[2]",
                             [VT("S3[2]", [
                                 VT("vec3"),
                                 VT(name_arr[i])])])])])

        tests.append([VT("S1[2]",
                         [VT("S2[2]",
                             [VT("S3[2]", [
                                 VT(name_arr[i])])])])])

    for t in tests:
        do_test_permutations(t, glsl_version, extensions, tests_path)

    tests = []

    for i in range(1, 2):
        tests.append([VT("S1",
                         [VT(name_arr[i] + "[3][2]")])])
    for i in range(1, 2):
        for j in range(1, 4):
            tests.append([VT("S1",
                             [VT(FLOAT_VEC_TYPES[j]),
                              VT(name_arr[i] + "[3][2]")])])

    for i in range(1, 2):
        for j in range(1, 4):
            tests.append([VT("S1", [VT("S2[2][2]",
                                       [VT(FLOAT_VEC_TYPES[j]),
                                        VT(name_arr[i])])])])

    for i in range(3, 4):
        tests.append([VT("S1", [VT(name_arr[i] + "[2][2]")])])

    for t in tests:
        do_test_permutations(t, glsl_version,
                             extensions + ["GL_ARB_arrays_of_arrays"],
                             tests_path)


def do_fp64_specific_tests():
    tests = []

    for t in DOUBLE_MAT_TYPES:
        tests.append([VT("S1", [VT(t)])])

    for t in ["dmat2x2", "dmat2x3", "dmat2x4", "dmat3x2", "dmat3x3"]:
        for j in range(1, 4):
            tests.append([VT("S1", [VT(FLOAT_VEC_TYPES[j]), VT(t)])])

        for j in range(1, 7):
            tests.append([VT("S1",
                             [VT("float[{}]".format(j)), VT(t)])])

    for j in range(1, 4):
        tests.append([VT("S1", [VT("S2[2]",
                                   [VT(FLOAT_VEC_TYPES[j]),
                                    VT("dmat2x2")])])])

    tests.append([VT("S1",
                     [VT("double"), VT("float"), VT("double[2]"),
                      VT("float[3]"), VT("dmat2x2")])])

    tests.append([VT("S1",
                     [VT("S2", [VT("double")]),
                      VT("S3", [VT("float")]),
                      VT("S4", [VT("dmat3x3")])])])

    for t in tests:
        do_test_permutations(t, 150, ["GL_ARB_gpu_shader_fp64"], fp64_ext_path)
        do_test_permutations(t, 400, [], fp64_core_path)


do_common_64bit_tests(400, [], DOUBLE_VEC_TYPES, fp64_core_path)
do_common_64bit_tests(150, ["GL_ARB_gpu_shader_fp64"], DOUBLE_VEC_TYPES,
                      fp64_ext_path)
do_common_64bit_tests(150, ["GL_ARB_gpu_shader_int64"], INT64_VEC_TYPES,
                      int64_path)
do_common_64bit_tests(150, ["GL_ARB_gpu_shader_int64"], UINT64_VEC_TYPES,
                      int64_path)

do_fp64_specific_tests()
